From bf23048b36075fb54933db206e50c4840e161cd7 Mon Sep 17 00:00:00 2001 From: Djordje Lukic Date: Sun, 25 Oct 2026 21:33:30 +0100 Subject: [PATCH 1/2] fix: allow MCP supervisors to start after an explicit stop Signed-off-by: Djordje Lukic --- pkg/tools/lifecycle/supervisor.go | 51 ++++++++-- pkg/tools/lifecycle/supervisor_resume_test.go | 96 +++++++++++++++++++ pkg/tools/lifecycle/supervisor_test.go | 5 +- pkg/tools/mcp/session_stdio_test.go | 23 +++++ 4 files changed, 163 insertions(+), 12 deletions(-) create mode 100644 pkg/tools/lifecycle/supervisor_resume_test.go diff --git a/pkg/tools/lifecycle/supervisor.go b/pkg/tools/lifecycle/supervisor.go index 8b1a874112..ca27c404cf 100644 --- a/pkg/tools/lifecycle/supervisor.go +++ b/pkg/tools/lifecycle/supervisor.go @@ -200,12 +200,14 @@ type Supervisor struct { startMu sync.Mutex // mu guards the rest of the fields. - mu sync.Mutex - session Session - stopping bool - watcherAlive bool - forceRestart bool // set by RestartAndWait so the watcher reconnects - restarted chan struct{} // closed and replaced on each successful restart + mu sync.Mutex + session Session + stopping bool + stopDone chan struct{} // closed after session close and watcher exit + connectReaped chan struct{} // closed after a stopped startup handshake is reaped + watcherAlive bool + forceRestart bool // set by RestartAndWait so the watcher reconnects + restarted chan struct{} // closed and replaced on each successful restart // done is closed when the supervisor enters a terminal state (Stopped // or Failed) so RestartAndWait can wake up promptly. Replaced with a // fresh channel by Start when transitioning out of a terminal state. @@ -292,6 +294,7 @@ func (s *Supervisor) Restarted() <-chan struct{} { // // When policy.StartupTimeout is non-zero the connect is bounded by it; on // expiry Start returns ErrInitTimeout and stays in StateStopped. +// After Stop, Start waits for teardown before reconnecting. func (s *Supervisor) Start(ctx context.Context) error { s.startMu.Lock() defer s.startMu.Unlock() @@ -302,8 +305,18 @@ func (s *Supervisor) Start(ctx context.Context) error { return nil } if s.stopping { + stopDone, connectReaped := s.stopDone, s.connectReaped s.mu.Unlock() - return ErrNotStarted + // Do not reuse the connector until the old watcher and any late + // handshake have finished closing their shared transport. + if err := waitForWatcher(ctx, stopDone); err != nil { + return err + } + if err := waitForWatcher(ctx, connectReaped); err != nil { + return err + } + s.mu.Lock() + s.stopping = false } // A crash loop detected by watch reports itself here exactly once // (see crashLoopErr's doc): the very next Start after this one attempts @@ -474,16 +487,23 @@ func (s *Supervisor) reapPendingConnect(ctx context.Context, p *pendingConnect) func (s *Supervisor) Stop(ctx context.Context) error { s.mu.Lock() if s.stopping { - watchDone := s.watchDone + stopDone := s.stopDone s.mu.Unlock() - return waitForWatcher(ctx, watchDone) + return waitForWatcher(ctx, stopDone) } s.stopping = true + stopDone := make(chan struct{}) + s.stopDone = stopDone sess := s.session s.session = nil watchDone := s.watchDone pending := s.inflightConnect s.inflightConnect = nil + var connectReaped chan struct{} + if pending != nil { + connectReaped = make(chan struct{}) + } + s.connectReaped = connectReaped // A deliberate Stop discards any pending crash-loop verdict and its // history: a future Start begins a fresh window rather than resuming // one that spans an intentional stop/start cycle. @@ -498,7 +518,10 @@ func (s *Supervisor) Stop(ctx context.Context) error { // session is closed rather than leaked. Runs in the background: a wedged // handshake must not block Stop. if pending != nil { - go s.reapPendingConnect(context.WithoutCancel(ctx), pending) + go func() { + defer close(connectReaped) + s.reapPendingConnect(context.WithoutCancel(ctx), pending) + }() } var closeErr error @@ -506,6 +529,14 @@ func (s *Supervisor) Stop(ctx context.Context) error { closeErr = sess.Close(context.WithoutCancel(ctx)) } waitErr := waitForWatcher(ctx, watchDone) + if waitErr == nil { + close(stopDone) + } else { + go func() { + _ = waitForWatcher(context.WithoutCancel(ctx), watchDone) + close(stopDone) + }() + } if closeErr != nil && ctx.Err() == nil { return closeErr } diff --git a/pkg/tools/lifecycle/supervisor_resume_test.go b/pkg/tools/lifecycle/supervisor_resume_test.go new file mode 100644 index 0000000000..8c03c95328 --- /dev/null +++ b/pkg/tools/lifecycle/supervisor_resume_test.go @@ -0,0 +1,96 @@ +package lifecycle_test + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "gotest.tools/v3/assert" + + "github.com/docker/docker-agent/pkg/tools/lifecycle" +) + +func TestSupervisor_StartAfterStop(t *testing.T) { + t.Parallel() + for _, stopFirst := range []bool{false, true} { + t.Run(map[bool]string{false: "started", true: "never-started"}[stopFirst], func(t *testing.T) { + sessions := []*fakeSession{newFakeSession(), newFakeSession(), newFakeSession()} + c := newScriptedConnector(scriptStep{session: sessions[0]}, scriptStep{session: sessions[1]}, scriptStep{session: sessions[2]}) + s := lifecycle.New("test", c, lifecycle.Policy{}) + if stopFirst { + assert.NilError(t, s.Stop(t.Context())) + } + for _, sess := range sessions { + assert.NilError(t, s.Start(t.Context())) + assert.Equal(t, s.State().State, lifecycle.StateReady) + sess.waitParked(t) + assert.NilError(t, s.Stop(t.Context())) + assert.Equal(t, s.State().State, lifecycle.StateStopped) + sess.waitClosed(t) + } + assert.Equal(t, c.Calls(), len(sessions)) + }) + } +} + +func TestSupervisor_StartWaitsForStoppedWatcher(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + disconnected, release := make(chan struct{}), make(chan struct{}) + first, second := newFakeSession(), newFakeSession() + c := newScriptedConnector(scriptStep{session: first}, scriptStep{session: second}) + s := lifecycle.New("test", c, lifecycle.Policy{OnDisconnect: func(error) { close(disconnected); <-release }}) + assert.NilError(t, s.Start(t.Context())) + first.fail(errors.New("disconnected")) + <-disconnected + ctx, cancel := context.WithCancel(t.Context()) + cancel() + assert.ErrorIs(t, s.Stop(ctx), context.Canceled) + assert.ErrorIs(t, s.Start(ctx), context.Canceled) + assert.Equal(t, c.Calls(), 1) + resumed := make(chan error, 1) + go func() { resumed <- s.Start(t.Context()) }() + synctest.Wait() + assert.Equal(t, c.Calls(), 1) + close(release) + assert.NilError(t, <-resumed) + assert.Equal(t, c.Calls(), 2) + assert.NilError(t, s.Stop(t.Context())) + }) +} + +type resumeConnector struct { + calls atomic.Int32 + release chan struct{} + first, second *fakeSession +} + +func (c *resumeConnector) Connect(context.Context) (lifecycle.Session, error) { + if c.calls.Add(1) == 1 { + <-c.release + return c.first, nil + } + return c.second, nil +} + +func TestSupervisor_StartAfterStopWaitsForLateConnect(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + c := &resumeConnector{release: make(chan struct{}), first: newFakeSession(), second: newFakeSession()} + s := lifecycle.New("test", c, lifecycle.Policy{StartupTimeout: time.Millisecond}) + assert.ErrorIs(t, s.Start(t.Context()), lifecycle.ErrInitTimeout) + assert.NilError(t, s.Stop(t.Context())) + resumed := make(chan error, 1) + go func() { resumed <- s.Start(t.Context()) }() + synctest.Wait() + assert.Equal(t, c.calls.Load(), int32(1)) + close(c.release) + assert.NilError(t, <-resumed) + c.first.waitClosed(t) + assert.Equal(t, c.calls.Load(), int32(2)) + assert.NilError(t, s.Stop(t.Context())) + }) +} diff --git a/pkg/tools/lifecycle/supervisor_test.go b/pkg/tools/lifecycle/supervisor_test.go index 8dd1a16ae9..58221ea544 100644 --- a/pkg/tools/lifecycle/supervisor_test.go +++ b/pkg/tools/lifecycle/supervisor_test.go @@ -1080,8 +1080,9 @@ func TestSupervisor_CrashLoopStopIsClean(t *testing.T) { assert.NilError(t, s.Stop(t.Context())) assert.Check(t, is.Equal(s.State().State, lifecycle.StateStopped)) - // A permanently-stopped supervisor never restarts, crash loop or not. - assert.Check(t, errors.Is(s.Start(t.Context()), lifecycle.ErrNotStarted)) + assert.Check(t, s.PendingCrashLoopError() == nil) + // The next start reaches the connector instead of the old crash-loop gate. + assert.ErrorContains(t, s.Start(t.Context()), "scripted connector exhausted") }) } diff --git a/pkg/tools/mcp/session_stdio_test.go b/pkg/tools/mcp/session_stdio_test.go index f3c5918f49..cdea174946 100644 --- a/pkg/tools/mcp/session_stdio_test.go +++ b/pkg/tools/mcp/session_stdio_test.go @@ -88,6 +88,29 @@ func sessionTestToolset(t *testing.T, mode string) (*Toolset, string) { return ts, marker } +func TestSessionStdioStartAfterStop(t *testing.T) { + t.Parallel() + ts, _ := sessionTestToolset(t, "normal") + for range 3 { + require.NoError(t, ts.Start(t.Context())) + available, err := ts.Tools(t.Context()) + require.NoError(t, err) + found := false + for _, tool := range available { + if tool.Name != "session_ping" { + continue + } + result, err := tool.Handler(t.Context(), tools.ToolCall{Function: tools.FunctionCall{Name: tool.Name, Arguments: "{}"}}, tools.NopRuntime{}) + require.NoError(t, err) + assert.Equal(t, "pong", result.Output) + found = true + } + require.True(t, found) + require.NoError(t, ts.Stop(t.Context())) + require.False(t, ts.IsStarted()) + } +} + func TestSessionStdioSetupCancellation(t *testing.T) { t.Parallel() ts, marker := sessionTestToolset(t, "hang-init") From a09be2001611b584d9ce830b1ed90ee4ff64df06 Mon Sep 17 00:00:00 2001 From: David Gageot Date: Tue, 27 Oct 2026 15:50:45 +0100 Subject: [PATCH 2/2] fix(mcp): make stop and resume transport-safe --- pkg/tools/lifecycle/supervisor.go | 91 +++++++++++---- pkg/tools/lifecycle/supervisor_resume_test.go | 107 ++++++++++++++++++ pkg/tools/mcp/mcp.go | 56 +++++---- pkg/tools/mcp/session_remote_test.go | 27 +++++ pkg/tools/mcp/session_stdio_test.go | 86 +++++++++++++- 5 files changed, 317 insertions(+), 50 deletions(-) diff --git a/pkg/tools/lifecycle/supervisor.go b/pkg/tools/lifecycle/supervisor.go index ca27c404cf..9a4406d648 100644 --- a/pkg/tools/lifecycle/supervisor.go +++ b/pkg/tools/lifecycle/supervisor.go @@ -137,7 +137,7 @@ type Policy struct { MaxAttempts int // 0 = default (5); negative = unlimited Backoff Backoff // zero fields use Backoff defaults - // StartupTimeout bounds the initial Connect performed by Start. Zero + // StartupTimeout bounds teardown waiting and the initial Connect in Start. Zero // means no timeout. It is enforced by racing Connect against a timer // rather than via a context deadline, so it still interrupts connectors // (notably MCP) that detach the context they are handed with @@ -162,6 +162,11 @@ type Policy struct { // OnFailed is called once when the supervisor enters StateFailed. OnFailed func(err error) + // Shutdown wraps session close and watcher teardown during Stop. It must + // finish all transport cleanup before returning. Concurrent Stops share + // one invocation, and Start waits for it before reusing the connector. + Shutdown func(ctx context.Context, teardown func() error) error + // Logger is used for lifecycle logs. Defaults to slog.Default(). Logger *slog.Logger } @@ -203,8 +208,9 @@ type Supervisor struct { mu sync.Mutex session Session stopping bool - stopDone chan struct{} // closed after session close and watcher exit - connectReaped chan struct{} // closed after a stopped startup handshake is reaped + stopDone chan struct{} // closed after session close and watcher exit + stopCancel context.CancelFunc // lets concurrent Stops interrupt the same teardown + connectReaped chan struct{} // closed after a stopped startup handshake is reaped watcherAlive bool forceRestart bool // set by RestartAndWait so the watcher reconnects restarted chan struct{} // closed and replaced on each successful restart @@ -292,7 +298,7 @@ func (s *Supervisor) Restarted() <-chan struct{} { // the watcher goroutine is launched (if not already alive) and state moves // to Ready. Concurrent Start calls serialize. // -// When policy.StartupTimeout is non-zero the connect is bounded by it; on +// When policy.StartupTimeout is non-zero teardown and connect share its budget; on // expiry Start returns ErrInitTimeout and stays in StateStopped. // After Stop, Start waits for teardown before reconnecting. func (s *Supervisor) Start(ctx context.Context) error { @@ -304,15 +310,21 @@ func (s *Supervisor) Start(ctx context.Context) error { s.mu.Unlock() return nil } + var timeout <-chan time.Time + if s.policy.StartupTimeout > 0 { + timer := time.NewTimer(s.policy.StartupTimeout) + defer timer.Stop() + timeout = timer.C + } if s.stopping { stopDone, connectReaped := s.stopDone, s.connectReaped s.mu.Unlock() // Do not reuse the connector until the old watcher and any late // handshake have finished closing their shared transport. - if err := waitForWatcher(ctx, stopDone); err != nil { + if err := s.waitForTeardown(ctx, stopDone, timeout); err != nil { return err } - if err := waitForWatcher(ctx, connectReaped); err != nil { + if err := s.waitForTeardown(ctx, connectReaped, timeout); err != nil { return err } s.mu.Lock() @@ -331,7 +343,7 @@ func (s *Supervisor) Start(ctx context.Context) error { s.tracker.Set(StateStarting) - sess, err := s.connect(ctx) + sess, err := s.connect(ctx, timeout) if err != nil { s.tracker.Fail(StateStopped, err) return err @@ -409,7 +421,7 @@ type pendingConnect struct { // spawning a new attempt per turn would accumulate orphaned handshakes and, // for stdio, leak subprocesses. A handshake that eventually completes is // adopted by the next Start. -func (s *Supervisor) connect(ctx context.Context) (Session, error) { +func (s *Supervisor) connect(ctx context.Context, timeout <-chan time.Time) (Session, error) { if s.policy.StartupTimeout <= 0 { return s.connector.Connect(ctx) } @@ -434,9 +446,6 @@ func (s *Supervisor) connect(ctx context.Context) (Session, error) { } s.mu.Unlock() - timer := time.NewTimer(s.policy.StartupTimeout) - defer timer.Stop() - select { case <-p.done: s.mu.Lock() @@ -459,8 +468,8 @@ func (s *Supervisor) connect(ctx context.Context) (Session, error) { return nil, res.err } return res.sess, res.err - case <-timer.C: - return nil, wrap(ErrInitTimeout, fmt.Errorf("%q did not connect within %s", s.name, s.policy.StartupTimeout)) + case <-timeout: + return nil, s.initTimeoutError() } } @@ -483,14 +492,21 @@ func (s *Supervisor) reapPendingConnect(ctx context.Context, p *pendingConnect) } // Stop tears the supervisor down. Idempotent. Blocks until the underlying -// session is closed. +// session is closed. Cancellation by any concurrent caller interrupts teardown. func (s *Supervisor) Stop(ctx context.Context) error { s.mu.Lock() if s.stopping { - stopDone := s.stopDone + stopDone, cancel := s.stopDone, s.stopCancel s.mu.Unlock() - return waitForWatcher(ctx, stopDone) + err := waitForWatcher(ctx, stopDone) + if ctx.Err() != nil { + cancel() + } + return err } + ctx, cancel := context.WithCancel(ctx) + defer cancel() + s.stopCancel = cancel s.stopping = true stopDone := make(chan struct{}) s.stopDone = stopDone @@ -524,12 +540,24 @@ func (s *Supervisor) Stop(ctx context.Context) error { }() } - var closeErr error - if sess != nil { - closeErr = sess.Close(context.WithoutCancel(ctx)) + teardown := func() error { + var closeErr error + if sess != nil { + closeErr = sess.Close(context.WithoutCancel(ctx)) + } + waitErr := waitForWatcher(ctx, watchDone) + if closeErr != nil && ctx.Err() == nil { + return closeErr + } + return waitErr } - waitErr := waitForWatcher(ctx, watchDone) - if waitErr == nil { + var err error + if shutdown := s.policy.Shutdown; shutdown != nil { + err = shutdown(ctx, teardown) + } else { + err = teardown() + } + if waitForWatcher(ctx, watchDone) == nil { close(stopDone) } else { go func() { @@ -537,10 +565,25 @@ func (s *Supervisor) Stop(ctx context.Context) error { close(stopDone) }() } - if closeErr != nil && ctx.Err() == nil { - return closeErr + return err +} + +func (s *Supervisor) initTimeoutError() error { + return wrap(ErrInitTimeout, fmt.Errorf("%q did not connect within %s", s.name, s.policy.StartupTimeout)) +} + +func (s *Supervisor) waitForTeardown(ctx context.Context, done <-chan struct{}, timeout <-chan time.Time) error { + if done == nil { + return nil + } + select { + case <-done: + return nil + case <-ctx.Done(): + return ctx.Err() + case <-timeout: + return s.initTimeoutError() } - return waitErr } func waitForWatcher(ctx context.Context, done <-chan struct{}) error { diff --git a/pkg/tools/lifecycle/supervisor_resume_test.go b/pkg/tools/lifecycle/supervisor_resume_test.go index 8c03c95328..afed14c1d4 100644 --- a/pkg/tools/lifecycle/supervisor_resume_test.go +++ b/pkg/tools/lifecycle/supervisor_resume_test.go @@ -66,6 +66,7 @@ type resumeConnector struct { calls atomic.Int32 release chan struct{} first, second *fakeSession + delay time.Duration } func (c *resumeConnector) Connect(context.Context) (lifecycle.Session, error) { @@ -73,6 +74,7 @@ func (c *resumeConnector) Connect(context.Context) (lifecycle.Session, error) { <-c.release return c.first, nil } + time.Sleep(c.delay) //nolint:forbidigo // Simulate handshake latency inside a synctest bubble. return c.second, nil } @@ -94,3 +96,108 @@ func TestSupervisor_StartAfterStopWaitsForLateConnect(t *testing.T) { assert.NilError(t, s.Stop(t.Context())) }) } + +func TestSupervisor_ResumeTimeoutWhileReaping(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + c := &resumeConnector{release: make(chan struct{}), first: newFakeSession(), second: newFakeSession()} + release := c.release + defer func() { + if release != nil { + close(release) + } + }() + s := lifecycle.New("test", c, lifecycle.Policy{StartupTimeout: time.Millisecond}) + assert.ErrorIs(t, s.Start(t.Context()), lifecycle.ErrInitTimeout) + assert.NilError(t, s.Stop(t.Context())) + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + for range 2 { + started := time.Now() + assert.ErrorIs(t, s.Start(ctx), lifecycle.ErrInitTimeout) + assert.Equal(t, time.Since(started), time.Millisecond) + assert.Equal(t, c.calls.Load(), int32(1)) + } + close(c.release) + release = nil + synctest.Wait() + assert.NilError(t, s.Start(t.Context())) + assert.Equal(t, c.calls.Load(), int32(2)) + assert.NilError(t, s.Stop(t.Context())) + }) +} + +func TestSupervisor_ResumeWaitsForShutdown(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + first, second := newFakeSession(), newFakeSession() + c := newScriptedConnector(scriptStep{session: first}, scriptStep{session: second}) + cleaning, release := make(chan struct{}), make(chan struct{}) + var shutdowns atomic.Int32 + s := lifecycle.New("test", c, lifecycle.Policy{ + StartupTimeout: time.Millisecond, + Shutdown: func(_ context.Context, teardown func() error) error { + err := teardown() + if shutdowns.Add(1) == 1 { + close(cleaning) + <-release + } + return err + }, + }) + assert.NilError(t, s.Start(t.Context())) + first.waitParked(t) + stopped := make(chan error, 1) + go func() { stopped <- s.Stop(t.Context()) }() + <-cleaning + joined := make(chan error, 1) + go func() { joined <- s.Stop(t.Context()) }() + ctx, cancel := context.WithCancel(t.Context()) + cancel() + assert.ErrorIs(t, s.Start(ctx), context.Canceled) + assert.ErrorIs(t, s.Stop(ctx), context.Canceled) + assert.ErrorIs(t, s.Start(t.Context()), lifecycle.ErrInitTimeout) + assert.Equal(t, c.Calls(), 1) + assert.Equal(t, shutdowns.Load(), int32(1)) + close(release) + assert.NilError(t, <-stopped) + assert.NilError(t, <-joined) + assert.NilError(t, s.Start(t.Context())) + assert.Equal(t, c.Calls(), 2) + assert.NilError(t, s.Stop(t.Context())) + }) +} + +func TestSupervisor_ResumeSharesStartupBudget(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + first, second := newFakeSession(), newFakeSession() + release := make(chan struct{}) + close(release) + c := &resumeConnector{release: release, first: first, second: second, delay: 6 * time.Millisecond} + var shutdowns atomic.Int32 + s := lifecycle.New("test", c, lifecycle.Policy{ + StartupTimeout: 10 * time.Millisecond, + Shutdown: func(_ context.Context, teardown func() error) error { + err := teardown() + if shutdowns.Add(1) == 1 { + time.Sleep(6 * time.Millisecond) //nolint:forbidigo // Consume part of the fake startup budget. + } + return err + }, + }) + assert.NilError(t, s.Start(t.Context())) + first.waitParked(t) + stopped := make(chan error, 1) + go func() { stopped <- s.Stop(t.Context()) }() + synctest.Wait() + started := time.Now() + assert.ErrorIs(t, s.Start(t.Context()), lifecycle.ErrInitTimeout) + assert.Equal(t, time.Since(started), 10*time.Millisecond) + assert.NilError(t, <-stopped) + synctest.Sleep(2 * time.Millisecond) + assert.NilError(t, s.Start(t.Context())) + assert.Equal(t, c.calls.Load(), int32(2)) + assert.NilError(t, s.Stop(t.Context())) + }) +} diff --git a/pkg/tools/mcp/mcp.go b/pkg/tools/mcp/mcp.go index 3b485839fd..8a6092f914 100644 --- a/pkg/tools/mcp/mcp.go +++ b/pkg/tools/mcp/mcp.go @@ -357,7 +357,7 @@ func firstOrZero[T any](s []T) T { // newSupervisor constructs a Supervisor wired to the toolset's mcpClient // using the provided policy as a base. Internal callbacks (OnDisconnect, -// OnRestart, Logger) are always overridden so the supervisor can +// OnRestart, Shutdown, Logger) are always overridden so the supervisor can // invalidate caches and refresh tools/prompts on reconnect; any values // passed in for those fields are ignored. // @@ -367,6 +367,10 @@ func firstOrZero[T any](s []T) T { func newSupervisor(ts *Toolset, base lifecycle.Policy) *lifecycle.Supervisor { connector := &clientConnector{ts: ts} policy := base + policy.Shutdown = nil + if ts.sessionOwned { + policy.Shutdown = ts.shutdown + } policy.Logger = slog.With("component", "mcp.supervisor", "server", ts.logID) policy.OnDisconnect = func(error) { ts.mu.Lock() @@ -531,28 +535,7 @@ func (ts *Toolset) Stop(ctx context.Context) error { return nil } if ts.sessionOwned { - var forceClose func() - var finalize func(context.Context) error - switch client := ts.mcpClient.(type) { - case *stdioMCPClient: - forceClose = client.cancelProcess - case *sessionRemoteClient: - forceClose = client.forceClose - finalize = client.Close - } - forced := make(chan struct{}) - stopForce := context.AfterFunc(ctx, func() { forceClose(); close(forced) }) - defer func() { - if !stopForce() { - <-forced - } - forceClose() - }() - err := ts.supervisor.Stop(ctx) - if finalize != nil { - err = errors.Join(err, finalize(ctx)) - } - return errors.Join(err, ctx.Err()) + return ts.supervisor.Stop(ctx) } if err := ts.supervisor.Stop(ctx); err != nil && ctx.Err() == nil { slog.ErrorContext(ctx, "Failed to stop MCP toolset", "server", ts.logID, "error", err) @@ -562,6 +545,32 @@ func (ts *Toolset) Stop(ctx context.Context) error { return nil } +// shutdown keeps session-owned transport cleanup inside the supervisor's stop barrier. +func (ts *Toolset) shutdown(ctx context.Context, teardown func() error) error { + var forceClose func() + var finalize func(context.Context) error + switch client := ts.mcpClient.(type) { + case *stdioMCPClient: + forceClose = client.cancelProcess + case *sessionRemoteClient: + forceClose = client.forceClose + finalize = client.Close + } + forced := make(chan struct{}) + stopForce := context.AfterFunc(ctx, func() { forceClose(); close(forced) }) + defer func() { + if !stopForce() { + <-forced + } + forceClose() + }() + err := teardown() + if finalize != nil { + err = errors.Join(err, finalize(ctx)) + } + return errors.Join(err, ctx.Err()) +} + // clientConnector adapts an mcpClient to the lifecycle.Connector interface. // It owns the initialize handshake and exposes the shared mcpClient as a // Session. @@ -680,6 +689,7 @@ func (c *clientConnector) Connect(ctx context.Context) (lifecycle.Session, error slog.DebugContext(ctx, "Started MCP toolset successfully", "server", ts.logID) ts.mu.Lock() defer ts.mu.Unlock() + ts.invalidateCache() ts.instructions = result.Instructions return &clientSession{client: ts.mcpClient}, nil diff --git a/pkg/tools/mcp/session_remote_test.go b/pkg/tools/mcp/session_remote_test.go index fbc3861f90..5e592544e4 100644 --- a/pkg/tools/mcp/session_remote_test.go +++ b/pkg/tools/mcp/session_remote_test.go @@ -518,3 +518,30 @@ func TestSessionRemoteDoesNotResumeInterruptedStream(t *testing.T) { require.Error(t, err) assert.Zero(t, gets.Load(), "SDK must not resume a client-owned SSE stream") } + +func TestSessionRemoteStartAfterStop(t *testing.T) { + t.Parallel() + for _, transport := range []string{"streamable", "sse"} { + t.Run(transport, func(t *testing.T) { + t.Parallel() + server := gomcp.NewServer(&gomcp.Implementation{Name: "resume", Version: "1"}, nil) + gomcp.AddTool(server, &gomcp.Tool{Name: "ping"}, func(context.Context, *gomcp.CallToolRequest, struct{}) (*gomcp.CallToolResult, any, error) { + return &gomcp.CallToolResult{Content: []gomcp.Content{&gomcp.TextContent{Text: "pong"}}}, nil, nil + }) + hs := httptest.NewServer(sessionRemoteHandler(transport, server)) + defer hs.Close() + ts := localSessionRemote(t, hs.URL, transport, nil) + defer func() { require.NoError(t, ts.Stop(t.Context())) }() + for range 3 { + require.NoError(t, ts.Start(t.Context())) + available, err := ts.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, available, 1) + result, err := available[0].Handler(t.Context(), tools.ToolCall{Function: tools.FunctionCall{Name: available[0].Name, Arguments: "{}"}}, tools.NopRuntime{}) + require.NoError(t, err) + assert.Equal(t, "pong", result.Output) + require.NoError(t, ts.Stop(t.Context())) + } + }) + } +} diff --git a/pkg/tools/mcp/session_stdio_test.go b/pkg/tools/mcp/session_stdio_test.go index cdea174946..76c77502ea 100644 --- a/pkg/tools/mcp/session_stdio_test.go +++ b/pkg/tools/mcp/session_stdio_test.go @@ -7,6 +7,7 @@ import ( "os" "path/filepath" "strings" + "sync/atomic" "testing" "time" @@ -15,6 +16,7 @@ import ( "github.com/stretchr/testify/require" "github.com/docker/docker-agent/pkg/tools" + "github.com/docker/docker-agent/pkg/tools/lifecycle" ) func TestSessionStdioHelper(t *testing.T) { @@ -67,6 +69,19 @@ func TestSessionStdioHelper(t *testing.T) { } } } + if os.Getenv("MCP_MODE") == "discovery" { + name, err := os.ReadFile(marker + ".capabilities") + if err != nil { + os.Exit(2) + } + server.RemoveTools("ping", "exit") + gomcp.AddTool(server, &gomcp.Tool{Name: string(name)}, func(context.Context, *gomcp.CallToolRequest, struct{}) (*gomcp.CallToolResult, any, error) { + return &gomcp.CallToolResult{}, nil, nil + }) + server.AddPrompt(&gomcp.Prompt{Name: string(name)}, func(context.Context, *gomcp.GetPromptRequest) (*gomcp.GetPromptResult, error) { + return &gomcp.GetPromptResult{}, nil + }) + } _ = server.Run(t.Context(), &gomcp.StdioTransport{}) _ = os.WriteFile(marker+".stopped", []byte("stopped"), 0o600) os.Exit(0) @@ -174,13 +189,15 @@ func TestSessionStdioStopUnblocksPipeWrite(t *testing.T) { _, err := tool.Handler(t.Context(), tools.ToolCall{Function: tools.FunctionCall{Name: tool.Name, Arguments: string(args)}}, tools.NopRuntime{}) called <- err }() + stopped := make(chan error, 1) + go func() { stopped <- ts.Stop(t.Context()) }() + require.Eventually(t, func() bool { return ts.State().State == lifecycle.StateStopped }, 5*time.Second, time.Millisecond) ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) defer cancel() - stopped := make(chan error, 1) - go func() { stopped <- ts.Stop(ctx) }() + require.ErrorIs(t, ts.Stop(ctx), context.DeadlineExceeded) select { case err := <-stopped: - require.ErrorIs(t, err, context.DeadlineExceeded) + require.ErrorIs(t, err, context.Canceled) case <-time.After(5 * time.Second): t.Fatal("stop did not kill the blocked subprocess") } @@ -191,3 +208,66 @@ func TestSessionStdioStopUnblocksPipeWrite(t *testing.T) { t.Fatal("tool call remained blocked after stop") } } + +func TestSessionStdioResumeRefreshesDiscovery(t *testing.T) { + t.Parallel() + ts, marker := sessionTestToolset(t, "discovery") + for _, name := range []string{"before", "after"} { + require.NoError(t, os.WriteFile(marker+".capabilities", []byte(name), 0o600)) + require.NoError(t, ts.Start(t.Context())) + available, err := ts.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, available, 1) + assert.Equal(t, "session_"+name, available[0].Name) + prompts, err := ts.ListPrompts(t.Context()) + require.NoError(t, err) + require.Len(t, prompts, 1) + assert.Equal(t, name, prompts[0].Name) + require.NoError(t, ts.Stop(t.Context())) + } +} + +func TestSessionStdioResumeWaitsForCleanup(t *testing.T) { + t.Parallel() + ts, _ := sessionTestToolset(t, "normal") + require.NoError(t, ts.Start(t.Context())) + cleaning, release := make(chan struct{}), make(chan struct{}) + defer func() { + if release != nil { + close(release) + } + }() + client := ts.mcpClient.(*stdioMCPClient) + client.processMu.Lock() + cancelProcess := client.processCancel + var cleanups atomic.Int32 + client.processCancel = func() { + if cleanups.Add(1) == 1 { + close(cleaning) + <-release + } + cancelProcess() + } + client.processMu.Unlock() + stopped := make(chan error, 1) + go func() { stopped <- ts.Stop(t.Context()) }() + <-cleaning + ctx, cancel := context.WithTimeout(t.Context(), 20*time.Millisecond) + defer cancel() + require.ErrorIs(t, ts.Start(ctx), context.DeadlineExceeded) + close(release) + require.NoError(t, <-stopped) + release = nil + require.NoError(t, ts.Start(t.Context())) + available, err := ts.Tools(t.Context()) + require.NoError(t, err) + for _, tool := range available { + if tool.Name == "session_ping" { + result, err := tool.Handler(t.Context(), tools.ToolCall{Function: tools.FunctionCall{Name: tool.Name, Arguments: "{}"}}, tools.NopRuntime{}) + require.NoError(t, err) + assert.Equal(t, "pong", result.Output) + return + } + } + t.Fatal("ping tool missing after resume") +}