diff --git a/pkg/acp/agent.go b/pkg/acp/agent.go index 67f7c3f1d0..4523abb8b2 100644 --- a/pkg/acp/agent.go +++ b/pkg/acp/agent.go @@ -121,8 +121,11 @@ func (s *Session) close(ctx context.Context) <-chan struct{} { } s.initTurns() s.cleanupDone = make(chan struct{}) + var terminalsClosed chan error if s.terminals != nil { - s.terminals.stop(ctx) + s.terminals.retire() + terminalsClosed = make(chan error, 1) + go func() { terminalsClosed <- s.terminals.cleanup(ctx) }() } var mcpClosed chan error if s.clientMCP != nil { @@ -143,8 +146,8 @@ func (s *Session) close(ctx context.Context) <-chan struct{} { s.cleanupErr = errors.Join(s.cleanupErr, fmt.Errorf("closing ACP runtime: %w", err)) } } - if s.terminals != nil { - s.cleanupErr = errors.Join(s.cleanupErr, s.terminals.wait()) + if terminalsClosed != nil { + s.cleanupErr = errors.Join(s.cleanupErr, <-terminalsClosed) } if mcpClosed != nil { s.cleanupErr = errors.Join(s.cleanupErr, <-mcpClosed) diff --git a/pkg/acp/terminal.go b/pkg/acp/terminal.go index d42f812e36..6c8c2a2eb0 100644 --- a/pkg/acp/terminal.go +++ b/pkg/acp/terminal.go @@ -26,8 +26,6 @@ type terminalManager struct { unreleased map[string]error uncertain error closing bool - done chan struct{} - closeErr error pending sync.WaitGroup } @@ -146,40 +144,32 @@ func (op *terminalOperation) kill(ctx context.Context) error { return err } -func (m *terminalManager) stop(ctx context.Context) { +// retire closes admission before the session starts asynchronous cleanup. +func (m *terminalManager) retire() { m.mu.Lock() defer m.mu.Unlock() if m.closing { return } m.closing = true - m.done = make(chan struct{}) for op := range m.active { op.cancel() } - go func() { - m.pending.Wait() - // Operations have relinquished ownership; only this final sweep can release now. - m.mu.Lock() - ids := make([]string, 0, len(m.unreleased)) - for id := range m.unreleased { - ids = append(ids, id) - } - m.mu.Unlock() - for _, id := range ids { - _ = m.release(ctx, id) - } - m.mu.Lock() - defer m.mu.Unlock() - m.closeErr = m.failureLocked() - close(m.done) - }() } -func (m *terminalManager) wait() error { +// cleanup is run once by the owning session, after retirement. +func (m *terminalManager) cleanup(ctx context.Context) error { + m.retire() + m.pending.Wait() + // Operations have relinquished ownership; only this final sweep can release now. m.mu.Lock() - done := m.done + ids := make([]string, 0, len(m.unreleased)) + for id := range m.unreleased { + ids = append(ids, id) + } m.mu.Unlock() - <-done - return m.closeErr + for _, id := range ids { + _ = m.release(ctx, id) + } + return m.failure() } diff --git a/pkg/acp/terminal_cleanup_test.go b/pkg/acp/terminal_cleanup_test.go new file mode 100644 index 0000000000..6594bf4b92 --- /dev/null +++ b/pkg/acp/terminal_cleanup_test.go @@ -0,0 +1,195 @@ +package acp + +import ( + "context" + "encoding/json" + "errors" + "strings" + "sync" + "testing" + "testing/synctest" + + acpsdk "github.com/coder/acp-go-sdk" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/propagation" + + "github.com/docker/docker-agent/pkg/config/latest" +) + +func TestTerminalRetirementIsSynchronous(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + manager := newTerminalManager(&acpsdk.AgentSideConnection{}, "owner") + s := &Session{terminals: manager} + ctx, op, err := manager.acquire(t.Context()) + require.NoError(t, err) + closed := s.close(t.Context()) + require.ErrorIs(t, ctx.Err(), context.Canceled) + _, _, err = manager.acquire(t.Context()) + require.ErrorIs(t, err, errSessionClosed) + synctest.Wait() + select { + case <-closed: + t.Fatal("cleanup returned while an operation still owns a terminal") + default: + } + op.finish() + <-closed + require.NoError(t, s.cleanupErr) + }) +} + +func TestTerminalCleanupRunsBeforeTurnDrainAndAlongsideRuntime(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + s := &Session{} + _, _, _, peer := terminalFixtureForSession(t, latest.Toolset{}, s) + manager := s.terminals + manager.unreleased["owned"] = nil + _, finishTurn, err := s.startTurn(t.Context()) + require.NoError(t, err) + releaseRequest := make(chan terminalRequest, 1) + peer.respond = func(r terminalRequest) (any, *acpsdk.RequestError) { + if r.method == "terminal/release" { + releaseRequest <- r + return nil, nil + } + return map[string]any{}, nil + } + rt := &barrierRuntime{started: make(chan struct{}), release: make(chan struct{})} + s.rt = rt + closed := s.close(t.Context()) + req := <-releaseRequest + select { + case <-rt.started: + t.Fatal("runtime closed before the active turn drained") + default: + } + finishTurn() + <-rt.started + synctest.Wait() + require.NoError(t, peer.reply(req.id, map[string]any{}, nil)) + synctest.Wait() + select { + case <-closed: + t.Fatal("session cleanup returned before runtime cleanup") + default: + } + close(rt.release) + <-closed + require.NoError(t, s.cleanupErr) + assert.Equal(t, 1, strings.Count(strings.Join(peer.methods(), ","), "terminal/release")) + }) +} + +func TestTerminalCleanupFailureIsCachedBySession(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + s := &Session{} + _, _, _, peer := terminalFixtureForSession(t, latest.Toolset{}, s) + s.terminals.unreleased["owned"] = errors.New("initial release failed") + releaseRequest := make(chan terminalRequest, 1) + peer.respond = func(r terminalRequest) (any, *acpsdk.RequestError) { + releaseRequest <- r + return nil, nil + } + ctx := propagation.TraceContext{}.Extract(t.Context(), propagation.MapCarrier{"traceparent": testTraceParent}) + ctx, cancel := context.WithCancel(ctx) + defer cancel() + closed := s.close(ctx) + req := <-releaseRequest + cancel() + var waiters sync.WaitGroup + for range 10 { + waiters.Go(func() { + assert.Equal(t, closed, s.close(t.Context())) + }) + } + waiters.Wait() + synctest.Wait() + select { + case <-closed: + t.Fatal("canceling the caller abandoned the final release") + default: + } + require.NoError(t, peer.reply(req.id, nil, acpsdk.NewInternalError("persistent failure"))) + <-closed + require.ErrorContains(t, s.cleanupErr, "cleanup is unresolved") + firstErr := s.cleanupErr + for range 3 { + <-s.close(t.Context()) + assert.Same(t, firstErr, s.cleanupErr) + } + assert.Equal(t, []string{"terminal/release"}, peer.methods()) + var params struct { + Meta map[string]any `json:"_meta"` + } + require.NoError(t, json.Unmarshal(req.params, ¶ms)) + assert.Equal(t, testTraceParent, params.Meta["traceparent"]) + }) +} + +func TestTerminalAdmissionRacingSessionClose(t *testing.T) { + t.Parallel() + for range 20 { + manager := newTerminalManager(&acpsdk.AgentSideConnection{}, "owner") + s := &Session{terminals: manager} + start := make(chan struct{}) + var callers sync.WaitGroup + for range 10 { + callers.Go(func() { + <-start + ctx, op, err := manager.acquire(t.Context()) + if err != nil { + assert.ErrorIs(t, err, errSessionClosed) + return + } + defer op.finish() + <-ctx.Done() + }) + } + close(start) + closed := s.close(t.Context()) + callers.Wait() + <-closed + require.NoError(t, s.cleanupErr) + manager.mu.Lock() + assert.Empty(t, manager.active) + manager.mu.Unlock() + } +} + +func TestTerminalCleanupConcurrentFirstClosers(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + s := &Session{} + _, _, _, peer := terminalFixtureForSession(t, latest.Toolset{}, s) + s.terminals.unreleased["owned"] = errors.New("initial failure") + requests := make(chan terminalRequest, 10) + peer.respond = func(req terminalRequest) (any, *acpsdk.RequestError) { + requests <- req + return nil, nil + } + start := make(chan struct{}) + closures := make(chan (<-chan struct{}), 10) + for range cap(closures) { + go func() { + <-start + closures <- s.close(t.Context()) + }() + } + close(start) + closed := <-closures + for range cap(closures) - 1 { + assert.Equal(t, closed, <-closures) + } + req := <-requests + synctest.Wait() + assert.Empty(t, requests, "concurrent closes must start one final sweep") + require.NoError(t, peer.reply(req.id, nil, acpsdk.NewInternalError("persistent failure"))) + <-closed + require.Error(t, s.cleanupErr) + assert.Equal(t, []string{"terminal/release"}, peer.methods()) + }) +} diff --git a/pkg/acp/terminal_test.go b/pkg/acp/terminal_test.go index d63ad34794..3489f4c608 100644 --- a/pkg/acp/terminal_test.go +++ b/pkg/acp/terminal_test.go @@ -123,7 +123,12 @@ func (p *terminalPeer) methods() []string { return names } -func newTerminalFixture(t *testing.T, cfg latest.Toolset) (*terminalToolset, *terminalManager, context.Context, *terminalPeer) { +func newTerminalFixture(t *testing.T, cfg latest.Toolset) (*terminalToolset, *Session, context.Context, *terminalPeer) { + t.Helper() + return terminalFixtureForSession(t, cfg, &Session{}) +} + +func terminalFixtureForSession(t *testing.T, cfg latest.Toolset, owner *Session) (*terminalToolset, *Session, context.Context, *terminalPeer) { t.Helper() wd := t.TempDir() a := &Agent{clientTerminal: true} @@ -131,12 +136,12 @@ func newTerminalFixture(t *testing.T, cfg latest.Toolset) (*terminalToolset, *te peer := &terminalPeer{peer: writer} conn := a.NewConnection(peer, reader) conn.SetLogger(slog.New(slog.DiscardHandler)) - owner := newTerminalManager(conn, "owner") - t.Cleanup(func() { owner.stop(t.Context()); _ = owner.wait(); _ = writer.Close(); <-conn.Done() }) + owner.terminals = newTerminalManager(conn, "owner") + t.Cleanup(func() { <-owner.close(t.Context()); _ = writer.Close(); <-conn.Done() }) rc := &config.RuntimeConfig{Config: config.Config{WorkingDir: wd}, EnvProviderOverride: environment.NewMapEnvProvider(map[string]string{"CONFIGURED": "expanded=value"})} toolset, err := newTerminalToolset(t.Context(), cfg, rc) require.NoError(t, err) - return toolset.(*terminalToolset), owner, context.WithValue(t.Context(), terminalOwnerKey{}, owner), peer + return toolset.(*terminalToolset), owner, context.WithValue(t.Context(), terminalOwnerKey{}, owner.terminals), peer } func callTerminalTool(t *testing.T, ctx context.Context, ts *terminalToolset, args string) *tools.ToolCallResult { @@ -315,9 +320,9 @@ func TestClientTerminalTimeoutAndUnknownCreate(t *testing.T) { result := callTerminalTool(t, ctx, ts, `{"cmd":"command","timeout":1}`) assert.True(t, result.IsError) assert.Contains(t, result.Output, "outcome is unknown") - require.Error(t, owner.failure()) - owner.stop(t.Context()) - require.Error(t, owner.wait()) + require.Error(t, owner.terminals.failure()) + <-owner.close(t.Context()) + require.Error(t, owner.cleanupErr) assert.Equal(t, []string{"terminal/create"}, peer.methods(), "never retry an ambiguous execution") }) synctest.Test(t, func(t *testing.T) { @@ -364,20 +369,19 @@ func TestClientTerminalReleaseRetriesOnlyDuringFinalCleanup(t *testing.T) { } result := callTerminalTool(t, ctx, ts, `{"cmd":"command"}`) assert.True(t, result.IsError) - require.Error(t, owner.failure()) + require.Error(t, owner.terminals.failure()) again := callTerminalTool(t, ctx, ts, `{"cmd":"must not run"}`) assert.True(t, again.IsError) assert.Equal(t, 1, releases) - owner.stop(t.Context()) - err := owner.wait() + <-owner.close(t.Context()) + err := owner.cleanupErr if persistent { require.Error(t, err) } else { require.NoError(t, err) } assert.Equal(t, 2, releases) - owner.stop(t.Context()) - _ = owner.wait() + <-owner.close(t.Context()) assert.Equal(t, 2, releases) } } @@ -411,9 +415,10 @@ func TestClientTerminalConcurrentOwnershipAndClose(t *testing.T) { } <-entered <-entered - owner.stop(t.Context()) + owner.close(t.Context()) wg.Wait() - require.NoError(t, owner.wait()) + <-owner.close(t.Context()) + require.NoError(t, owner.cleanupErr) assert.Equal(t, 2, strings.Count(strings.Join(peer.methods(), ","), "terminal/release")) result := callTerminalTool(t, ctx, ts, `{"cmd":"stale handler"}`) assert.True(t, result.IsError) @@ -630,12 +635,12 @@ func TestClientTerminalDefinitiveRejectionDoesNotPoisonOwner(t *testing.T) { result := callTerminalTool(t, ctx, ts, `{"cmd":"command"}`) assert.True(t, result.IsError) assert.Contains(t, result.Output, "creation rejected") - require.NoError(t, owner.failure()) + require.NoError(t, owner.terminals.failure()) peer.respond = nil result = callTerminalTool(t, ctx, ts, `{"cmd":"retry"}`) assert.False(t, result.IsError, result.Output) - owner.stop(t.Context()) - require.NoError(t, owner.wait()) + <-owner.close(t.Context()) + require.NoError(t, owner.cleanupErr) } } @@ -652,9 +657,8 @@ func TestClientTerminalCloseJoinsLateCreateAndBackgroundWait(t *testing.T) { for _, lateCreate := range []bool{false, true} { t.Run(map[bool]string{false: "background wait", true: "late create"}[lateCreate], func(t *testing.T) { synctest.Test(t, func(t *testing.T) { - ts, owner, ctx, peer := newTerminalFixture(t, latest.Toolset{}) - a, s := newResumeFixture(t, ts.workingDir) - s.terminals = owner + a, s := newResumeFixture(t, t.TempDir()) + ts, _, ctx, peer := terminalFixtureForSession(t, latest.Toolset{}, s) entered := make(chan terminalRequest, 1) peer.respond = func(r terminalRequest) (any, *acpsdk.RequestError) { switch r.method { @@ -689,6 +693,7 @@ func TestClientTerminalCloseJoinsLateCreateAndBackgroundWait(t *testing.T) { assert.Contains(t, (<-done).Output, "cancelled") require.NoError(t, <-closed) assert.Equal(t, "terminal/release", peer.methods()[len(peer.methods())-1]) + assert.Equal(t, 1, strings.Count(strings.Join(peer.methods(), ","), "terminal/release")) before := len(peer.snapshot()) result := callTerminalTool(t, ctx, ts, `{"cmd":"stale owner"}`) assert.True(t, result.IsError) @@ -720,12 +725,12 @@ func TestClientTerminalWaitErrorAndMissingCreateID(t *testing.T) { result := callTerminalTool(t, ctx, ts, `{"cmd":"command"}`) assert.True(t, result.IsError) if missingID { - require.Error(t, owner.failure()) + require.Error(t, owner.terminals.failure()) assert.Equal(t, []string{"terminal/create"}, peer.methods()) } else { assert.Contains(t, result.Output, "Error waiting") assert.Equal(t, []string{"terminal/create", "terminal/wait_for_exit", "terminal/kill", "terminal/output", "terminal/release"}, peer.methods()) - require.NoError(t, owner.failure()) + require.NoError(t, owner.terminals.failure()) } } } @@ -758,7 +763,7 @@ func TestClientTerminalSessionIsolationAndDisconnectNoFallback(t *testing.T) { assert.Equal(t, []acpsdk.SessionId{acpsdk.SessionId(first.id), second.SessionId}, ids) // A missing connection is a client execution error, not a reason to run locally. ts, owner, ctx, _ := newTerminalFixture(t, latest.Toolset{}) - owner.conn = nil + owner.terminals.conn = nil result := callTerminalTool(t, ctx, ts, `{"cmd":"echo must-not-run-locally"}`) assert.True(t, result.IsError) assert.Contains(t, result.Output, "connection unavailable") diff --git a/pkg/acp/tracing_test.go b/pkg/acp/tracing_test.go index bf8b140452..bbddc1f141 100644 --- a/pkg/acp/tracing_test.go +++ b/pkg/acp/tracing_test.go @@ -246,12 +246,12 @@ func TestACPTraceOutboundBoundaries(t *testing.T) { } peer.mu.Unlock() ts, owner, _, terminal := newTerminalFixture(t, latest.Toolset{}) - terminalCtx := context.WithValue(ctx, terminalOwnerKey{}, owner) + terminalCtx := context.WithValue(ctx, terminalOwnerKey{}, owner.terminals) result := callTerminalTool(t, terminalCtx, ts, `{"cmd":"echo test"}`) require.False(t, result.IsError) - opCtx, op, err := owner.acquire(terminalCtx) + opCtx, op, err := owner.terminals.acquire(terminalCtx) require.NoError(t, err) - require.NoError(t, op.create(opCtx, acpsdk.CreateTerminalRequest{SessionId: owner.sid, Command: "unused", Meta: map[string]any{"traceparent": "stale", "baggage": "private"}})) + require.NoError(t, op.create(opCtx, acpsdk.CreateTerminalRequest{SessionId: owner.terminals.sid, Command: "unused", Meta: map[string]any{"traceparent": "stale", "baggage": "private"}})) require.NoError(t, op.kill(opCtx)) require.NoError(t, op.release(opCtx)) op.finish()