fix: drain runtime streams before releasing turn ownership

Cancel abandoned runs and drain their event channels before returning.
Keep API streaming locks and ACP turn tokens held until teardown
finishes; apply the same lifecycle contract to A2A, CLI, and
synchronous runtime wrappers. Regression tests cover cancellation,
handler errors, and follow-up turns.

Signed-off-by: David Gageot <david.gageot@docker.com>
Assisted-By: Claude (Anthropic)

pkg/a2a/adapter.go

@@ -2,6 +2,7 @@ package a2a
 
 import (
 	"cmp"
+	"context"
 	"errors"
 	"fmt"
 	"iter"
@@ -127,8 +128,14 @@ func runDockerAgent(ctx agent.InvocationContext, t *team.Team, agentName string,
 			return
 		}
 
-		// Run the agent and collect events
-		eventsChan := rt.RunStream(ctx, sess)
+		// Early iterator exits must cancel and join the runtime.
+		streamCtx, cancel := context.WithCancel(ctx)
+		eventsChan := rt.RunStream(streamCtx, sess)
+		defer func() {
+			cancel()
+			for range eventsChan {
+			}
+		}()
 
 		// Track accumulated content for chunked responses
 		var contentBuilder strings.Builder
@@ -152,7 +159,7 @@ func runDockerAgent(ctx agent.InvocationContext, t *team.Team, agentName string,
 
 		for event := range eventsChan {
 			if ctx.Ended() {
-				slog.Debug("Invocation ended, stopping agent", "agent", agentName)
+				slog.DebugContext(ctx, "Invocation ended, stopping agent", "agent", agentName)
 				return
 			}
 

pkg/a2a/adapter_drain_test.go

@@ -0,0 +1,114 @@
+package a2a
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/chat"
+	"github.com/docker/docker-agent/pkg/model/provider/base"
+	"github.com/docker/docker-agent/pkg/modelsdev"
+	"github.com/docker/docker-agent/pkg/servesafety"
+	"github.com/docker/docker-agent/pkg/session"
+	"github.com/docker/docker-agent/pkg/tools"
+)
+
+// stallingStream yields its chunks and then parks until the model call is
+// cancelled, standing in for a provider that only stops on cancellation.
+type stallingStream struct {
+	ctx    context.Context //nolint:containedctx // the model-call context, inspected after the adapter returns
+	chunks []string
+	closed atomic.Bool
+}
+
+func (s *stallingStream) Recv() (chat.MessageStreamResponse, error) {
+	if len(s.chunks) > 0 {
+		chunk := s.chunks[0]
+		s.chunks = s.chunks[1:]
+		return chat.MessageStreamResponse{
+			Choices: []chat.MessageStreamChoice{{Delta: chat.MessageDelta{Content: chunk}}},
+		}, nil
+	}
+	<-s.ctx.Done()
+	return chat.MessageStreamResponse{}, s.ctx.Err()
+}
+
+func (s *stallingStream) Close() { s.closed.Store(true) }
+
+type stallingProvider struct {
+	chunks []string
+
+	mu     sync.Mutex
+	stream *stallingStream
+}
+
+func (p *stallingProvider) ID() modelsdev.ID { return modelsdev.NewID("test", "mock-model") }
+
+func (p *stallingProvider) CreateChatCompletionStream(ctx context.Context, _ []chat.Message, _ []tools.Tool) (chat.MessageStream, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.stream = &stallingStream{ctx: ctx, chunks: p.chunks}
+	return p.stream, nil
+}
+
+func (p *stallingProvider) BaseConfig() base.Config { return base.Config{} }
+
+func (p *stallingProvider) MaxTokens() int { return 0 }
+
+func (p *stallingProvider) current() *stallingStream {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.stream
+}
+
+// When the ADK consumer stops early, the adapter must cancel the runtime it
+// started and wait for the stream to close before returning, rather than
+// leaving the model call running in the background.
+func TestRunDockerAgent_EarlyStopCancelsAndDrainsRuntime(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name   string
+		chunks []string
+		stop   func(ctx *fakeInvocationContext) bool
+	}{
+		{
+			name:   "consumer breaks",
+			chunks: []string{"first"},
+			stop:   func(*fakeInvocationContext) bool { return true },
+		},
+		{
+			name:   "invocation ended",
+			chunks: []string{"first", "second"},
+			stop:   func(ctx *fakeInvocationContext) bool { ctx.EndInvocation(); return false },
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			prov := &stallingProvider{chunks: tc.chunks}
+			tm, root := newTeamWithProvider(prov)
+			store := session.NewInMemorySessionStore()
+			ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-drain-"+tc.name, "Hi")
+
+			var yielded int
+			for _, err := range runDockerAgent(ctx, tm, root.Name(), root, store, servesafety.Resolved{Policy: session.SafetyPolicyRestricted}, testWorkspaceRoot) {
+				require.NoError(t, err)
+				yielded++
+				if tc.stop(ctx) {
+					break
+				}
+			}
+			assert.Equal(t, 1, yielded)
+
+			stream := prov.current()
+			require.NotNil(t, stream, "the model must have been called")
+			require.ErrorIs(t, stream.ctx.Err(), context.Canceled, "the adapter must cancel the runtime it abandons")
+			assert.True(t, stream.closed.Load(), "the adapter must drain the runtime stream before returning")
+		})
+	}
+}

pkg/acp/agent.go

@@ -702,7 +702,14 @@ func (a *Agent) runAgent(ctx context.Context, acpSess *Session) error {
 		slog.DebugContext(ctx, "Failed to emit available commands", "error", err)
 	}
 
-	eventsChan := acpSess.rt.RunStream(ctx, acpSess.sess)
+	runCtx, cancel := context.WithCancel(ctx)
+	eventsChan := acpSess.rt.RunStream(runCtx, acpSess.sess)
+	// Cancel on handler errors too, before waiting for runtime teardown.
+	defer func() {
+		cancel()
+		for range eventsChan {
+		}
+	}()
 	toolCallArgs := map[string]string{}
 
 	for event := range eventsChan {

pkg/acp/runagent_drain_test.go

@@ -0,0 +1,83 @@
+package acp
+
+import (
+	"context"
+	"testing"
+	"testing/synctest"
+
+	acpsdk "github.com/coder/acp-go-sdk"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/runtime"
+	"github.com/docker/docker-agent/pkg/session"
+)
+
+type drainingPromptRuntime struct {
+	fakeRuntime
+
+	started chan struct{}
+	release chan struct{}
+	stopped chan struct{}
+	fail    bool
+}
+
+func (r *drainingPromptRuntime) RunStream(ctx context.Context, _ *session.Session) <-chan runtime.Event {
+	events := make(chan runtime.Event)
+	close(r.started)
+	go func() {
+		defer close(events)
+		defer close(r.stopped)
+		if r.fail {
+			events <- &runtime.ToolCallResponseEvent{ToolCallID: "missing-start"}
+		}
+		<-ctx.Done()
+		events <- runtime.Warning("teardown started", "root")
+		<-r.release
+		for range 256 {
+			events <- runtime.Warning("teardown still running", "root")
+		}
+	}()
+	return events
+}
+
+func TestPrompt_DrainsBeforeReleasingTurn(t *testing.T) {
+	t.Parallel()
+	for _, fail := range []bool{false, true} {
+		name := "cancel"
+		if fail {
+			name = "handler error"
+		}
+		t.Run(name, func(t *testing.T) {
+			synctest.Test(t, func(t *testing.T) {
+				rt := &drainingPromptRuntime{started: make(chan struct{}), release: make(chan struct{}), stopped: make(chan struct{}), fail: fail}
+				agent, sess, _ := newPromptTestAgent(t, rt)
+				done := promptAsync(agent, t.Context(), promptRequest("first"))
+				<-rt.started
+				if !fail {
+					require.NoError(t, agent.Cancel(t.Context(), acpsdk.CancelNotification{SessionId: testSessionID}))
+				}
+				synctest.Wait()
+				select {
+				case <-sess.turns:
+					t.Error("turn token released before runtime teardown")
+				default:
+				}
+				close(rt.release)
+				synctest.Wait()
+				result := <-done
+				if fail {
+					require.ErrorContains(t, result.err, "missing tool call arguments")
+				} else {
+					require.NoError(t, result.err)
+					assert.Equal(t, acpsdk.StopReasonCancelled, result.response.StopReason)
+				}
+				select {
+				case <-rt.stopped:
+				default:
+					t.Error("teardown events were not drained")
+				}
+			})
+		})
+	}
+}

pkg/cli/runner.go

@@ -118,23 +118,32 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess
 		sess.AddMessage(userMsg)
 		sess.AddAttachedFile(attachedPath)
 
+		// Stop only this turn; subsequent user messages still need the parent context.
+		runCtx, stopRun := context.WithCancel(ctx)
+		events := rt.RunStream(runCtx, sess)
+		defer func() {
+			stopRun()
+			for range events {
+			}
+		}()
+
 		if cfg.OutputJSON {
-			for event := range rt.RunStream(ctx, sess) {
+			for event := range events {
 				switch e := event.(type) {
 				case *runtime.ToolCallConfirmationEvent:
 					// JSON mode has no user at stdin — reject unconditionally.
 					// A confirmation event under AutoApprove means a
 					// preempt-yolo hook overrode --yolo; the safe answer is
 					// still Reject (the hook said Ask, not Approve).
-					rt.Resume(ctx, runtime.ResumeReject(""))
+					rt.Resume(runCtx, runtime.ResumeReject(""))
 				case *runtime.ElicitationRequestEvent:
-					_ = rt.ResumeElicitation(ctx, "decline", nil, e.ElicitationID)
+					_ = rt.ResumeElicitation(runCtx, "decline", nil, e.ElicitationID)
 				case *runtime.MaxIterationsReachedEvent:
 					switch handleMaxIterationsAutoApprove(cfg.AutoApprove, &autoExtensions, e.MaxIterations) {
 					case maxIterContinue:
-						rt.Resume(ctx, runtime.ResumeApprove())
+						rt.Resume(runCtx, runtime.ResumeApprove())
 					default: // maxIterStop or maxIterPrompt (no interactive prompt in JSON mode)
-						rt.Resume(ctx, runtime.ResumeReject(""))
+						rt.Resume(runCtx, runtime.ResumeReject(""))
 						return nil
 					}
 				case *runtime.ErrorEvent:
@@ -154,7 +163,7 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess
 		firstLoop := true
 		lastAgent := rt.CurrentAgentName(ctx)
 		var lastConfirmedToolCallID string
-		for event := range rt.RunStream(ctx, sess) {
+		for event := range events {
 			agentName := event.GetAgentName()
 			if agentName != "" && (firstLoop || lastAgent != agentName) {
 				if !firstLoop {
@@ -178,15 +187,15 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess
 				lastConfirmedToolCallID = e.ToolCall.ID // Store the ID to avoid duplicate printing
 				switch result {
 				case ConfirmationApprove:
-					rt.Resume(ctx, runtime.ResumeApprove())
+					rt.Resume(runCtx, runtime.ResumeApprove())
 				case ConfirmationApproveBalanced:
 					sess.SetSafetyPolicy(session.SafetyPolicyBalanced)
-					rt.Resume(ctx, runtime.ResumeApproveBalanced())
+					rt.Resume(runCtx, runtime.ResumeApproveBalanced())
 				case ConfirmationApproveSession:
 					sess.SetSafetyPolicy(session.SafetyPolicyAutonomous)
-					rt.Resume(ctx, runtime.ResumeApproveAutonomous())
+					rt.Resume(runCtx, runtime.ResumeApproveAutonomous())
 				case ConfirmationReject:
-					rt.Resume(ctx, runtime.ResumeReject(""))
+					rt.Resume(runCtx, runtime.ResumeReject(""))
 					lastConfirmedToolCallID = "" // Clear on reject since tool won't execute
 				case ConfirmationAbort:
 					// Stop the agent loop immediately
@@ -228,20 +237,20 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess
 			case *runtime.MaxIterationsReachedEvent:
 				switch handleMaxIterationsAutoApprove(cfg.AutoApprove, &autoExtensions, e.MaxIterations) {
 				case maxIterContinue:
-					rt.Resume(ctx, runtime.ResumeApprove())
+					rt.Resume(runCtx, runtime.ResumeApprove())
 				case maxIterStop:
-					rt.Resume(ctx, runtime.ResumeReject(""))
+					rt.Resume(runCtx, runtime.ResumeReject(""))
 					return nil
 				case maxIterPrompt:
 					result := out.PromptMaxIterationsContinue(ctx, e.MaxIterations)
 					switch result {
 					case ConfirmationApprove:
-						rt.Resume(ctx, runtime.ResumeApprove())
+						rt.Resume(runCtx, runtime.ResumeApprove())
 					case ConfirmationReject:
-						rt.Resume(ctx, runtime.ResumeReject(""))
+						rt.Resume(runCtx, runtime.ResumeReject(""))
 						return nil
 					case ConfirmationAbort:
-						rt.Resume(ctx, runtime.ResumeReject(""))
+						rt.Resume(runCtx, runtime.ResumeReject(""))
 						return nil
 					}
 				}
@@ -250,7 +259,7 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess
 				if !ok || serverURL == "" {
 					// Keep draining after declining forms so follow-up events cannot stall the turn.
 					slog.WarnContext(ctx, "Declining elicitation without form support in CLI mode", "message", e.Message)
-					_ = rt.ResumeElicitation(ctx, "decline", nil, e.ElicitationID)
+					_ = rt.ResumeElicitation(runCtx, "decline", nil, e.ElicitationID)
 					continue
 				}
 
@@ -262,9 +271,9 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess
 
 				switch result {
 				case ConfirmationApprove:
-					_ = rt.ResumeElicitation(ctx, "accept", nil, e.ElicitationID)
+					_ = rt.ResumeElicitation(runCtx, "accept", nil, e.ElicitationID)
 				case ConfirmationReject:
-					_ = rt.ResumeElicitation(ctx, "decline", nil, e.ElicitationID)
+					_ = rt.ResumeElicitation(runCtx, "decline", nil, e.ElicitationID)
 					return errors.New("OAuth authorization rejected by user")
 				}
 			}

pkg/cli/runner_drain_test.go

@@ -0,0 +1,128 @@
+package cli
+
+import (
+	"bytes"
+	"context"
+	"strings"
+	"testing"
+
+	"gotest.tools/v3/assert"
+
+	"github.com/docker/docker-agent/pkg/runtime"
+	"github.com/docker/docker-agent/pkg/session"
+)
+
+// stallingRunStream emits the trigger events, waits for the run context to be
+// cancelled, then floods more events than any buffer holds before closing.
+// done is closed only once every event has been consumed, so Run can only
+// observe it closed if it cancelled the producer and drained the stream.
+func stallingRunStream(triggers []runtime.Event, done chan struct{}) func(context.Context, *session.Session) <-chan runtime.Event {
+	return func(ctx context.Context, _ *session.Session) <-chan runtime.Event {
+		ch := make(chan runtime.Event)
+		go func() {
+			defer close(ch)
+			defer close(done)
+			for _, e := range triggers {
+				ch <- e
+			}
+			<-ctx.Done()
+			for range 256 {
+				ch <- runtime.Warning("trailing", "test")
+			}
+		}()
+		return ch
+	}
+}
+
+func repeatMaxIterEvents(n int) []runtime.Event {
+	events := make([]runtime.Event, n)
+	for i := range events {
+		events[i] = maxIterEvent(10)
+	}
+	return events
+}
+
+// Every early return out of a turn (error, safety cap) must cancel the
+// runtime stream and drain it before the turn ends.
+func TestRunEarlyReturnCancelsAndDrainsStream(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name     string
+		cfg      Config
+		triggers []runtime.Event
+		wantErr  string
+	}{
+		{
+			name:     "json error",
+			cfg:      Config{OutputJSON: true},
+			triggers: []runtime.Event{runtime.Error("model failed")},
+			wantErr:  "model failed",
+		},
+		{
+			name:     "json max iterations cap",
+			cfg:      Config{OutputJSON: true, AutoApprove: true},
+			triggers: repeatMaxIterEvents(maxAutoExtensions + 1),
+		},
+		{
+			name:     "text max iterations cap",
+			cfg:      Config{AutoApprove: true},
+			triggers: repeatMaxIterEvents(maxAutoExtensions + 1),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			done := make(chan struct{})
+			rt := &mockRuntime{runStreamFn: stallingRunStream(tc.triggers, done)}
+
+			var buf bytes.Buffer
+			err := Run(t.Context(), NewPrinter(&buf), tc.cfg, rt, session.New(), []string{"hello"})
+			if tc.wantErr == "" {
+				assert.NilError(t, err)
+			} else {
+				assert.ErrorContains(t, err, tc.wantErr)
+			}
+
+			select {
+			case <-done:
+			default:
+				t.Fatal("the turn ended before cancelling and draining the runtime stream")
+			}
+			assert.Check(t, !strings.Contains(buf.String(), "trailing"), "drained events must not be printed: %q", buf.String())
+		})
+	}
+}
+
+// A turn that ends early must only cancel its own stream: later user messages
+// still run with a live context.
+func TestRunEarlyReturnKeepsFollowUpTurnsAlive(t *testing.T) {
+	t.Parallel()
+
+	var turns int
+	var secondTurnCtxErr error
+	rt := &mockRuntime{
+		runStreamFn: func(ctx context.Context, _ *session.Session) <-chan runtime.Event {
+			turns++
+			ch := make(chan runtime.Event, maxAutoExtensions+1)
+			if turns == 1 {
+				for _, e := range repeatMaxIterEvents(maxAutoExtensions + 1) {
+					ch <- e
+				}
+			} else {
+				secondTurnCtxErr = ctx.Err()
+				ch <- runtime.AgentChoice("test", "sess", "second turn answer")
+			}
+			close(ch)
+			return ch
+		},
+	}
+
+	var buf bytes.Buffer
+	err := Run(t.Context(), NewPrinter(&buf), Config{AutoApprove: true}, rt, session.New(), []string{"one", "two"})
+	assert.NilError(t, err)
+
+	assert.Equal(t, turns, 2)
+	assert.NilError(t, secondTurnCtxErr)
+	assert.Check(t, strings.Contains(buf.String(), "second turn answer"), "the follow-up turn must run: %q", buf.String())
+}

pkg/runtime/loop.go

@@ -1152,7 +1152,14 @@ func (r *LocalRuntime) runTurn(
 // messages. This is a convenience wrapper around RunStream for non-streaming
 // callers.
 func (r *LocalRuntime) Run(ctx context.Context, sess *session.Session) ([]session.Message, error) {
+	ctx, cancel := context.WithCancel(ctx)
 	events := r.RunStream(ctx, sess)
+	// Cancel before draining so error returns cannot strand the producer.
+	defer func() {
+		cancel()
+		for range events {
+		}
+	}()
 	for event := range events {
 		if errEvent, ok := event.(*ErrorEvent); ok {
 			return nil, fmt.Errorf("%s", errEvent.Error)

pkg/runtime/remote_runtime.go

@@ -313,7 +313,14 @@ func (r *RemoteRuntime) RunStream(ctx context.Context, sess *session.Session) <-
 
 // Run starts the agent's interaction loop and returns the final messages
 func (r *RemoteRuntime) Run(ctx context.Context, sess *session.Session) ([]session.Message, error) {
+	ctx, cancel := context.WithCancel(ctx)
 	eventsChan := r.RunStream(ctx, sess)
+	// Cancel before draining so error returns cannot strand the forwarder.
+	defer func() {
+		cancel()
+		for range eventsChan {
+		}
+	}()
 
 	for event := range eventsChan {
 		if errEvent, ok := event.(*ErrorEvent); ok {

pkg/runtime/run_drain_test.go

@@ -0,0 +1,107 @@
+package runtime
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"testing/synctest"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/agent"
+	"github.com/docker/docker-agent/pkg/api"
+	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/docker/docker-agent/pkg/session"
+	"github.com/docker/docker-agent/pkg/team"
+)
+
+// stallingObserver stands in for teardown work (e.g. persistence) that only
+// completes once the run context is cancelled: OnEvent parks on StreamStopped
+// until ctx is done, so the channel cannot close before the consumer cancels.
+type stallingObserver struct{ done chan struct{} }
+
+func (stallingObserver) OnRunStart(context.Context, *session.Session) {}
+
+func (o stallingObserver) OnEvent(ctx context.Context, _ *session.Session, event Event) {
+	if _, ok := event.(*StreamStoppedEvent); ok {
+		<-ctx.Done()
+		close(o.done)
+	}
+}
+
+// Run returning on the first ErrorEvent must not abandon the stream: it has
+// to cancel the runtime it started and wait for the channel to close, so no
+// teardown work outlives the call.
+func TestLocalRuntime_RunCancelsAndDrainsStreamOnError(t *testing.T) {
+	t.Parallel()
+
+	synctest.Test(t, func(t *testing.T) {
+		prov := &failingProvider{id: "test/failing", err: errors.New("401 unauthorized")}
+		root := agent.New("root", "test", agent.WithModel(prov))
+		obs := stallingObserver{done: make(chan struct{})}
+		rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)),
+			WithSessionCompaction(false),
+			WithModelStore(mockModelStore{}),
+			WithEventObserver(obs),
+		)
+		require.NoError(t, err)
+
+		_, err = rt.Run(t.Context(), session.New(session.WithUserMessage("hi")))
+		require.ErrorContains(t, err, "401 unauthorized")
+
+		select {
+		case <-obs.done:
+		default:
+			t.Fatal("Run returned before cancelling and draining its stream")
+		}
+	})
+}
+
+// stallingRemoteClient serves a stream that fails, waits for the caller to
+// cancel, then floods more events than the runtime buffer holds. Only a
+// consumer that cancels and drains lets the producer finish.
+type stallingRemoteClient struct {
+	stubRemoteClient
+
+	done chan struct{}
+}
+
+func (c *stallingRemoteClient) RunAgent(ctx context.Context, _, _ string, _ []api.Message, _ string) (<-chan Event, error) {
+	ch := make(chan Event)
+	go func() {
+		defer close(ch)
+		defer close(c.done)
+		ch <- Error("remote failure")
+		<-ctx.Done()
+		for range 2 * defaultEventChannelCapacity {
+			ch <- Warning("trailing", "test")
+		}
+	}()
+	return ch, nil
+}
+
+func (c *stallingRemoteClient) RunAgentWithAgentName(ctx context.Context, sessionID, agentFile, _ string, msgs []api.Message, model string) (<-chan Event, error) {
+	return c.RunAgent(ctx, sessionID, agentFile, msgs, model)
+}
+
+func TestRemoteRuntime_RunCancelsAndDrainsStreamOnError(t *testing.T) {
+	t.Parallel()
+
+	synctest.Test(t, func(t *testing.T) {
+		client := &stallingRemoteClient{
+			stubRemoteClient: stubRemoteClient{cfg: &latest.Config{Agents: latest.Agents{{Name: "test"}}}},
+			done:             make(chan struct{}),
+		}
+		rt, err := NewRemoteRuntime(client)
+		require.NoError(t, err)
+
+		_, err = rt.Run(t.Context(), session.New(session.WithUserMessage("hi")))
+		require.EqualError(t, err, "remote failure")
+
+		select {
+		case <-client.done:
+		default:
+			t.Fatal("Run returned before cancelling and draining its stream")
+		}
+	})
+}

pkg/runtime/runtime.go

@@ -77,7 +77,9 @@ type Runtime interface {
 	EmitAgentInfo(ctx context.Context, events EventSink)
 	// ResetStartupInfo resets the startup info emission flag, allowing re-emission
 	ResetStartupInfo()
-	// RunStream starts the agent's interaction loop and returns a channel of events
+	// RunStream starts the agent's interaction loop and returns a channel of events.
+	// Consumers must drain it to closure before releasing turn ownership, even
+	// after cancellation. Cancel the supplied context before draining on errors.
 	RunStream(ctx context.Context, sess *session.Session) <-chan Event
 	// Run starts the agent's interaction loop and returns the final messages
 	Run(ctx context.Context, sess *session.Session) ([]session.Message, error)

pkg/server/session_manager.go

@@ -1167,6 +1167,12 @@ func (sm *SessionManager) RunSession(ctx context.Context, sessionID, agentFilena
 		}
 
 		stream := runtimeSession.runtime.RunStream(streamCtx, sess)
+		// Teardown must finish before the deferred streaming.Unlock.
+		defer func(events <-chan runtime.Event) {
+			cancel()
+			for range events {
+			}
+		}(stream)
 		for stream != nil {
 			select {
 			case event, ok := <-titleEvents:

pkg/server/session_manager_drain_test.go

@@ -0,0 +1,88 @@
+package server
+
+import (
+	"context"
+	"testing"
+	"testing/synctest"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/api"
+	"github.com/docker/docker-agent/pkg/runtime"
+	"github.com/docker/docker-agent/pkg/session"
+)
+
+type drainingSessionRuntime struct {
+	fakeRuntime
+
+	started chan struct{}
+	release chan struct{}
+	stopped chan struct{}
+}
+
+func (r *drainingSessionRuntime) RunStream(ctx context.Context, _ *session.Session) <-chan runtime.Event {
+	events := make(chan runtime.Event)
+	close(r.started)
+	go func() {
+		defer close(events)
+		defer close(r.stopped)
+		<-ctx.Done()
+		events <- runtime.Warning("teardown started", "root")
+		<-r.release
+		for range 256 {
+			events <- runtime.Warning("teardown still running", "root")
+		}
+	}()
+	return events
+}
+
+func TestRunSession_CancellationDrainsBeforeUnlock(t *testing.T) {
+	t.Parallel()
+	synctest.Test(t, func(t *testing.T) {
+		ctx, cancel := context.WithCancel(t.Context())
+		defer cancel()
+		rt := &drainingSessionRuntime{started: make(chan struct{}), release: make(chan struct{}), stopped: make(chan struct{})}
+		sess := session.New()
+		sm := newTestSessionManager(t, sess, rt)
+		output, err := sm.RunSession(ctx, sess.ID, "agent", "root", []api.Message{{Content: "first"}}, "")
+		require.NoError(t, err)
+		outputClosed := make(chan struct{})
+		go func() {
+			for range output {
+			}
+			close(outputClosed)
+		}()
+		<-rt.started
+		cancel()
+		synctest.Wait()
+		rs, ok := sm.runtimeSessions.Load(sess.ID)
+		require.True(t, ok)
+		unlocked := rs.streaming.TryLock()
+		if unlocked {
+			rs.streaming.Unlock()
+		}
+		assert.False(t, unlocked, "turn ownership released before runtime teardown")
+		if !unlocked {
+			_, err = sm.RunSession(t.Context(), sess.ID, "agent", "root", []api.Message{{Content: "second"}}, "")
+			require.ErrorIs(t, err, ErrSessionBusy)
+		}
+		select {
+		case <-outputClosed:
+			t.Error("forwarder returned before runtime teardown")
+		default:
+		}
+		close(rt.release)
+		synctest.Wait()
+		select {
+		case <-rt.stopped:
+		default:
+			t.Error("teardown events were not drained")
+		}
+		select {
+		case <-outputClosed:
+		default:
+			t.Error("forwarder did not finish")
+		}
+	})
+}