From e0785f33f663687cf3b5991b045d332511a864df Mon Sep 17 00:00:00 2001 From: Eron Wright Date: Sat, 1 Aug 2026 16:29:34 -0700 Subject: [PATCH] fix(server): stop per-session toolsets on session delete MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A session materialised via teamloader owns a team whose toolsets may hold external resources — notably stdio MCP subprocesses. DeleteSession cancelled the runtime context but never called team.StopToolSets, so those subprocesses leaked until the server process exited; BatchDeleteSessions had the same gap. Track the per-session team on activeRuntimes and route both delete paths through one scheduleTeardown helper: it keeps the entry in deletedSessions, waits for the session's stream to drain, stops the toolsets, then drops the entry. The drain and StopToolSets carry separate budgets, so a session that exhausts the drain still gets a live context to stop its toolsets with. BatchDeleteSessions now registers the same teardown, so a batch-deleted session is observable through WaitStopped just like a singly-deleted one — it previously never stored into deletedSessions, making WaitStopped return nil while teardown was still running. WaitStopped in turn waits on a completion channel closed after StopToolSets rather than polling the streaming mutex, so it means "fully torn down" and can no longer return before teardown has run. Attached runtimes (AttachRuntime) leave team nil, so stopping is a no-op there — their toolset lifecycle belongs to the embedder. Claude-Session: https://claude.ai/code/session_01BiXE2cM4DxDQNYhKfsPn5K --- pkg/server/author_safety_test.go | 5 +- pkg/server/session_manager.go | 176 ++++++++++++++++++++--------- pkg/server/session_manager_test.go | 150 +++++++++++++++++++++++- 3 files changed, 276 insertions(+), 55 deletions(-) diff --git a/pkg/server/author_safety_test.go b/pkg/server/author_safety_test.go index 866e732fc1..5535560b20 100644 --- a/pkg/server/author_safety_test.go +++ b/pkg/server/author_safety_test.go @@ -79,7 +79,7 @@ func newAuthorSafetySessionManager(t *testing.T) (*SessionManager, session.Store func buildRuntime(t *testing.T, sm *SessionManager, sess *session.Session, agentFilename, currentAgent string) { t.Helper() - run, _, err := sm.runtimeForSession(t.Context(), sess, agentFilename, currentAgent, &config.RuntimeConfig{}) + run, _, _, err := sm.runtimeForSession(t.Context(), sess, agentFilename, currentAgent, &config.RuntimeConfig{}) require.NoError(t, err) t.Cleanup(func() { _ = run.Close() }) } @@ -266,8 +266,9 @@ func TestAuthorSafetyDefault_FailedBuildKeepsMarkerForRetry(t *testing.T) { sm.newRuntime = func(context.Context, *team.Team, ...runtime.Opt) (runtime.Runtime, error) { return nil, buildErr } - _, _, err = sm.runtimeForSession(ctx, sess, "agent.yaml", "root", &config.RuntimeConfig{}) + run, _, _, err := sm.runtimeForSession(ctx, sess, "agent.yaml", "root", &config.RuntimeConfig{}) require.ErrorIs(t, err, buildErr) + require.Nil(t, run, "a failed build must not hand back a runtime") assert.Equal(t, session.SafetyPolicy(""), sess.GetSafetyPolicy(), "a failed build must not seed the in-memory session") diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 1b77dab610..1bcb614e2f 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -39,6 +39,18 @@ type activeRuntimes struct { session *session.Session // The actual session object used by the runtime titleGen *sessiontitle.Generator // Title generator (includes fallback models) + // team is the per-session team this manager built via teamloader in + // runtimeForSession. It is stopped (StopToolSets) on session teardown so + // stdio MCP subprocesses and other toolset-owned resources are released, + // not leaked until process exit. Nil for attached runtimes (AttachRuntime), + // whose team and toolset lifecycle belong to the external embedder. + team *team.Team + + // stopped is closed once teardown has finished: the stream has drained + // and the toolsets have been stopped. Set by scheduleTeardown, which is + // the only writer; nil until the session is deleted. + stopped chan struct{} + streaming sync.Mutex // Held while a RunStream is in progress; serialises concurrent requests } @@ -823,6 +835,91 @@ func (sm *SessionManager) GetSessions(ctx context.Context) ([]*session.Session, return sessions, nil } +const ( + // sessionDrainTimeout bounds how long teardown waits for a deleted + // session's stream to drain before giving up and stopping its toolsets + // anyway. + sessionDrainTimeout = 5 * time.Minute + // toolSetStopTimeout bounds StopToolSets. It is a separate budget from + // sessionDrainTimeout on purpose: a session that exhausted the drain + // budget must still get a live context to stop its toolsets with. + toolSetStopTimeout = 30 * time.Second +) + +// scheduleTeardown detaches the post-delete teardown of a session's runtime and +// runs it in the background: wait for the stream to drain, release the session's +// toolset-owned resources, then drop the deletedSessions entry. Both delete +// paths (DeleteSession and BatchDeleteSessions) go through here so they share +// one teardown contract; callers hold sm.mux. +// +// The entry stays in deletedSessions for the duration so WaitStopped can observe +// teardown after the runtime has been deregistered, and rs.stopped is closed +// before the entry is dropped so a waiter never misses the completion signal. +// +// ctx only supplies values (trace and logging context): teardown outlives the +// request that triggered the delete, so cancellation is stripped and each phase +// carries its own deadline instead. +func (sm *SessionManager) scheduleTeardown(ctx context.Context, sessionID string, rs *activeRuntimes) { + rs.stopped = make(chan struct{}) + sm.deletedSessions.Store(sessionID, rs) + + detached := context.WithoutCancel(ctx) + go func() { + // Deferred LIFO: close(stopped) first, so a WaitStopped that still + // holds the entry is released before the entry disappears. + defer sm.deletedSessions.Delete(sessionID) + defer close(rs.stopped) + + drainSessionStream(detached, rs) + sm.stopSessionToolSets(detached, rs) + }() +} + +// drainSessionStream blocks until the session's RunStream has released the +// streaming mutex, or until sessionDrainTimeout expires. Teardown continues +// either way — the timeout exists so a wedged stream cannot pin the session's +// toolsets forever. +func drainSessionStream(ctx context.Context, rs *activeRuntimes) { + ctx, cancel := context.WithTimeout(ctx, sessionDrainTimeout) + defer cancel() + + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + if rs.streaming.TryLock() { + rs.streaming.Unlock() + return + } + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } +} + +// stopSessionToolSets releases the toolset-owned resources (e.g. stdio MCP +// subprocesses) of a per-session team the manager built via teamloader — +// otherwise they leak until the server process exits. It is a no-op for +// attached runtimes, whose team is nil and whose toolset lifecycle belongs to +// the external embedder. Call it only after the session's stream has drained, +// so a toolset is never stopped mid-turn. +func (sm *SessionManager) stopSessionToolSets(ctx context.Context, rs *activeRuntimes) { + if rs == nil || rs.team == nil { + return + } + ctx, cancel := context.WithTimeout(ctx, toolSetStopTimeout) + defer cancel() + + sid := "" + if rs.session != nil { + sid = rs.session.ID + } + if err := rs.team.StopToolSets(ctx); err != nil { + slog.ErrorContext(ctx, "Failed to stop session tool sets", "session_id", sid, "error", err) + } +} + // DeleteSession deletes a session by ID. It cancels the runtime context and // removes the session from all registries. Callers that need to wait for // the stream to fully stop should call WaitStopped afterwards. @@ -852,32 +949,8 @@ func (sm *SessionManager) DeleteSession(ctx context.Context, sessionID string) e if sessionRuntime.cancel != nil { sessionRuntime.cancel() } - // Keep the entry in deletedSessions so WaitStopped can probe the - // streaming mutex after the runtime is deregistered. - sm.deletedSessions.Store(sess.ID, sessionRuntime) + sm.scheduleTeardown(ctx, sess.ID, sessionRuntime) sm.runtimeSessions.Delete(sess.ID) - - // Background cleanup: remove the deletedSessions entry once the - // stream goroutine has exited. This prevents a memory leak when - // the caller does not use ?wait=true. - go func() { - ticker := time.NewTicker(100 * time.Millisecond) - defer ticker.Stop() - deadline := time.After(5 * time.Minute) - for { - if sessionRuntime.streaming.TryLock() { - sessionRuntime.streaming.Unlock() - sm.deletedSessions.Delete(sess.ID) - return - } - select { - case <-deadline: - sm.deletedSessions.Delete(sess.ID) - return - case <-ticker.C: - } - } - }() } sm.dropEventLog(sess.ID) sm.followUpInjectors.Delete(sess.ID) @@ -887,33 +960,27 @@ func (sm *SessionManager) DeleteSession(ctx context.Context, sessionID string) e return nil } -// WaitStopped blocks until the session's runtime stream goroutine has fully -// exited (streaming mutex released), the timeout fires, or ctx is cancelled -// (e.g. client disconnect). It should be called after DeleteSession. -// Returns nil when the stream has stopped. +// WaitStopped blocks until the session's teardown has fully completed — its +// runtime stream goroutine has exited and its toolsets have been stopped — or +// the timeout fires, or ctx is cancelled (e.g. client disconnect). It should be +// called after DeleteSession or BatchDeleteSessions. Returns nil when teardown +// has finished, including when there is nothing left to wait for. func (sm *SessionManager) WaitStopped(ctx context.Context, sessionID string, timeout time.Duration) error { rs, ok := sm.deletedSessions.Load(sessionID) if !ok { return nil // already cleaned up } - deadline := time.After(timeout) - ticker := time.NewTicker(50 * time.Millisecond) - defer ticker.Stop() + deadline := time.NewTimer(timeout) + defer deadline.Stop() - for { - if rs.streaming.TryLock() { - rs.streaming.Unlock() - sm.deletedSessions.Delete(sessionID) - return nil - } - select { - case <-ctx.Done(): - return ctx.Err() - case <-deadline: - return fmt.Errorf("timeout waiting for session %s to stop", sessionID) - case <-ticker.C: - } + select { + case <-rs.stopped: + return nil + case <-ctx.Done(): + return ctx.Err() + case <-deadline.C: + return fmt.Errorf("timeout waiting for session %s to stop", sessionID) } } @@ -942,13 +1009,15 @@ func (sm *SessionManager) RunSession(ctx context.Context, sessionID, agentFilena var titleGen *sessiontitle.Generator if !exists { var rt runtime.Runtime - rt, titleGen, err = sm.runtimeForSession(ctx, sess, agentFilename, currentAgent, sm.runConfig) + var tm *team.Team + rt, titleGen, tm, err = sm.runtimeForSession(ctx, sess, agentFilename, currentAgent, sm.runConfig) if err != nil { cancel() return nil, err } runtimeSession = &activeRuntimes{ runtime: rt, + team: tm, cancel: cancel, session: sess, titleGen: titleGen, @@ -1416,7 +1485,7 @@ func (sm *SessionManager) generateTitle(ctx context.Context, sess *session.Sessi } } -func (sm *SessionManager) runtimeForSession(ctx context.Context, sess *session.Session, agentFilename, currentAgent string, rc *config.RuntimeConfig) (_ runtime.Runtime, _ *sessiontitle.Generator, err error) { +func (sm *SessionManager) runtimeForSession(ctx context.Context, sess *session.Session, agentFilename, currentAgent string, rc *config.RuntimeConfig) (_ runtime.Runtime, _ *sessiontitle.Generator, _ *team.Team, err error) { // Caller (RunSession) holds sm.mux and has already verified that no // active runtime exists for this session. This function is purely a // constructor: it must not touch sm.runtimeSessions, otherwise it would @@ -1442,14 +1511,14 @@ func (sm *SessionManager) runtimeForSession(ctx context.Context, sess *session.S loadResult, err := sm.loadTeamWithConfig(ctx, agentFilename, rc, teamloader.WithWorkingDir(sess.WorkingDir)) if err != nil { - return nil, nil, err + return nil, nil, nil, err } t := loadResult.Team // Resolve the team's default agent when no specific agent was requested. agt, err := t.AgentOrDefault(currentAgent) if err != nil { - return nil, nil, err + return nil, nil, nil, err } currentAgent = agt.Name() sess.MaxIterations = agt.MaxIterations() @@ -1502,7 +1571,7 @@ func (sm *SessionManager) runtimeForSession(ctx context.Context, sess *session.S } run, err := newRuntime(ctx, t, opts...) if err != nil { - return nil, nil, err + return nil, nil, nil, err } // If any later construction step fails, close the runtime before // returning: the caller only ever sees the error, so an unclosed @@ -1546,7 +1615,7 @@ func (sm *SessionManager) runtimeForSession(ctx context.Context, sess *session.S slog.DebugContext(ctx, "Runtime created for session", "session_id", sess.ID) - return run, titleGen, nil + return run, titleGen, t, nil } // applyAuthorSafetyDefault seeds an API-created session that carries no @@ -2113,6 +2182,11 @@ func (sm *SessionManager) BatchDeleteSessions(ctx context.Context, sessionIDs [] if sessionRuntime.cancel != nil { sessionRuntime.cancel() } + // Same teardown as DeleteSession, so a batch-deleted session + // releases its toolset-owned resources (stdio MCP subprocesses + // etc.) and is observable through WaitStopped just like a + // singly-deleted one. + sm.scheduleTeardown(ctx, sessionID, sessionRuntime) sm.runtimeSessions.Delete(sessionID) } sm.dropEventLog(sessionID) diff --git a/pkg/server/session_manager_test.go b/pkg/server/session_manager_test.go index 595a565cf1..0947826f7a 100644 --- a/pkg/server/session_manager_test.go +++ b/pkg/server/session_manager_test.go @@ -29,6 +29,7 @@ import ( "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/session/sqlitestore" "github.com/docker/docker-agent/pkg/sessiontitle" + "github.com/docker/docker-agent/pkg/team" "github.com/docker/docker-agent/pkg/tools" ) @@ -1227,7 +1228,7 @@ func TestRuntimeForSession_RegistersSessionScopedElicitationSink(t *testing.T) { require.False(t, sm.HasEventSource(sess.ID)) - run, _, err := sm.runtimeForSession(ctx, sess, "agent.yaml", "", &config.RuntimeConfig{}) + run, _, _, err := sm.runtimeForSession(ctx, sess, "agent.yaml", "", &config.RuntimeConfig{}) require.NoError(t, err) t.Cleanup(func() { _ = run.Close() }) @@ -1753,7 +1754,7 @@ func TestDeleteSession_SilencesLiveRuntimeElicitationDelivery(t *testing.T) { sources := config.Sources{"agent.yaml": config.NewBytesSource("agent.yaml", cfg)} sm := NewSessionManager(ctx, sources, store, 0, &config.RuntimeConfig{}) - run, _, err := sm.runtimeForSession(ctx, sess, "agent.yaml", "", &config.RuntimeConfig{}) + run, _, _, err := sm.runtimeForSession(ctx, sess, "agent.yaml", "", &config.RuntimeConfig{}) require.NoError(t, err) t.Cleanup(func() { _ = run.Close() }) lr, ok := run.(*runtime.LocalRuntime) @@ -1985,6 +1986,151 @@ func TestBatchDeleteSessions_ToleratesNilRuntimeCancel(t *testing.T) { assert.False(t, ok, "the runtime entry must still be deregistered") } +// stopCountingToolSet is a startable toolset that records how often it was +// stopped, so teardown tests can assert that deleting a session actually +// releases its toolset-owned resources (the stand-in for a stdio MCP +// subprocess) instead of leaking them until process exit. +type stopCountingToolSet struct { + stops atomic.Int32 +} + +func (s *stopCountingToolSet) Tools(context.Context) ([]tools.Tool, error) { return nil, nil } + +func (s *stopCountingToolSet) Start(context.Context) error { return nil } + +func (s *stopCountingToolSet) Stop(context.Context) error { + s.stops.Add(1) + return nil +} + +// newToolSetSession registers a session whose runtime carries a per-session +// team with one started toolset, the shape runtimeForSession produces. It +// returns the toolset so the caller can assert on teardown. +func newToolSetSession(t *testing.T, sm *SessionManager, sess *session.Session) *stopCountingToolSet { + t.Helper() + + ts := &stopCountingToolSet{} + tm := team.New(team.WithAgents(agent.New("root", "", agent.WithToolSets(ts)))) + + // Start the toolsets the way the runtime does; Agent.StopToolSets skips + // any that never started, so an unstarted toolset would pass vacuously. + agt, err := tm.AgentOrDefault("root") + require.NoError(t, err) + _, err = agt.Tools(t.Context()) + require.NoError(t, err) + + sm.runtimeSessions.Store(sess.ID, &activeRuntimes{ + runtime: &fakeRuntime{}, + session: sess, + team: tm, + }) + return ts +} + +// TestDeleteSession_StopsSessionToolSets covers the leak this fix targets: a +// per-session team's toolsets (stdio MCP subprocesses and the like) must be +// stopped once the session is deleted and its stream has drained. +func TestDeleteSession_StopsSessionToolSets(t *testing.T) { + t.Parallel() + + ctx := t.Context() + store := session.NewInMemorySessionStore() + sm := NewSessionManager(ctx, config.Sources{}, store, 0, &config.RuntimeConfig{}) + + sess := session.New() + require.NoError(t, store.AddSession(ctx, sess)) + ts := newToolSetSession(t, sm, sess) + + require.NoError(t, sm.DeleteSession(ctx, sess.ID)) + require.NoError(t, sm.WaitStopped(ctx, sess.ID, 5*time.Second)) + + assert.Equal(t, int32(1), ts.stops.Load(), "delete must stop the session's tool sets") +} + +// TestBatchDeleteSessions_StopsSessionToolSets is the batch counterpart: the +// batch path must tear a session down exactly like DeleteSession does, and +// WaitStopped must observe that teardown rather than reporting an immediate +// false "already stopped". +func TestBatchDeleteSessions_StopsSessionToolSets(t *testing.T) { + t.Parallel() + + ctx := t.Context() + store := session.NewInMemorySessionStore() + sm := NewSessionManager(ctx, config.Sources{}, store, 0, &config.RuntimeConfig{}) + + sess := session.New() + require.NoError(t, store.AddSession(ctx, sess)) + ts := newToolSetSession(t, sm, sess) + + deleted, failed := sm.BatchDeleteSessions(ctx, []string{sess.ID}) + require.Equal(t, 1, deleted) + require.Empty(t, failed) + + require.NoError(t, sm.WaitStopped(ctx, sess.ID, 5*time.Second)) + assert.Equal(t, int32(1), ts.stops.Load(), "batch delete must stop the session's tool sets") +} + +// TestWaitStopped_WaitsForToolSetTeardown pins the contract WaitStopped +// advertises: it returns only once teardown has completed, tool sets included. +// A stream still holding the streaming mutex parks it; releasing the stream +// lets teardown finish, and the toolset is stopped by the time WaitStopped +// returns. +func TestWaitStopped_WaitsForToolSetTeardown(t *testing.T) { + t.Parallel() + + ctx := t.Context() + store := session.NewInMemorySessionStore() + sm := NewSessionManager(ctx, config.Sources{}, store, 0, &config.RuntimeConfig{}) + + sess := session.New() + require.NoError(t, store.AddSession(ctx, sess)) + ts := newToolSetSession(t, sm, sess) + + rs, ok := sm.runtimeSessions.Load(sess.ID) + require.True(t, ok) + rs.streaming.Lock() // stand in for an in-flight RunStream + + require.NoError(t, sm.DeleteSession(ctx, sess.ID)) + + // Teardown is parked on the drain, so nothing has been stopped yet and a + // short wait must time out rather than report success. + assert.Equal(t, int32(0), ts.stops.Load()) + require.Error(t, sm.WaitStopped(ctx, sess.ID, 100*time.Millisecond)) + + rs.streaming.Unlock() + + require.NoError(t, sm.WaitStopped(ctx, sess.ID, 5*time.Second)) + assert.Equal(t, int32(1), ts.stops.Load()) +} + +// TestDeleteSession_StopsSessionToolSetsOnlyOnce pins the idempotency the two +// delete paths rely on: a second teardown of the same session must not stop an +// already-stopped tool set again. +func TestDeleteSession_StopsSessionToolSetsOnlyOnce(t *testing.T) { + t.Parallel() + + ctx := t.Context() + store := session.NewInMemorySessionStore() + sm := NewSessionManager(ctx, config.Sources{}, store, 0, &config.RuntimeConfig{}) + + sess := session.New() + require.NoError(t, store.AddSession(ctx, sess)) + ts := newToolSetSession(t, sm, sess) + rs, ok := sm.runtimeSessions.Load(sess.ID) + require.True(t, ok) + + require.NoError(t, sm.DeleteSession(ctx, sess.ID)) + require.NoError(t, sm.WaitStopped(ctx, sess.ID, 5*time.Second)) + require.Equal(t, int32(1), ts.stops.Load()) + + // A repeated delete never reaches teardown — the session is already gone + // from the store — so drive the second teardown directly. + require.Error(t, sm.DeleteSession(ctx, sess.ID)) + sm.stopSessionToolSets(ctx, rs) + + assert.Equal(t, int32(1), ts.stops.Load(), "a stopped tool set must not be stopped twice") +} + type commandFakeRuntime struct { fakeRuntime