From 68e6493ff35aca8bd953a696040b6775f0076efb 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, once the session's stream has drained, call StopToolSets from both delete paths, bounding it by the originating request's deadline. Nil for attached runtimes (AttachRuntime), whose toolset lifecycle belongs to the embedder, so it is a no-op there. --- pkg/server/session_manager.go | 85 ++++++++++++++++++++++++++---- pkg/server/session_manager_test.go | 4 +- 2 files changed, 78 insertions(+), 11 deletions(-) diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index c4f8053da..7924067c5 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -39,6 +39,13 @@ 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 + streaming sync.Mutex // Held while a RunStream is in progress; serialises concurrent requests } @@ -682,6 +689,35 @@ func (sm *SessionManager) GetSessions(ctx context.Context) ([]*session.Session, return sessions, nil } +// sessionDrainTimeout bounds the whole post-delete teardown of a session: +// waiting for its stream to drain plus stopping its toolsets. Detached from +// the delete request's context (which is cancelled when the handler returns) +// so teardown runs to completion in the background. +const sessionDrainTimeout = 5 * time.Minute + +// stopSessionToolSets releases the toolset-owned resources (e.g. stdio MCP +// subprocesses) of a per-session team the manager built via teamloader. 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. +// +// The caller passes a detached, deadline-bounded context (the drain goroutine's +// own budget, see sessionDrainTimeout): detached because this outlives the +// request that triggered the delete, bounded so StopToolSets cannot block +// forever. +func (sm *SessionManager) stopSessionToolSets(ctx context.Context, rs *activeRuntimes) { + if rs == nil || rs.team == nil { + return + } + sid := "" + if rs.session != nil { + sid = rs.session.ID + } + if err := rs.team.StopToolSets(ctx); err != nil { + slog.Error("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. @@ -718,11 +754,16 @@ func (sm *SessionManager) DeleteSession(ctx context.Context, sessionID string) e // 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. + // the caller does not use ?wait=true. It also releases the session's + // toolset-owned resources (stdio MCP subprocesses etc.) once the + // stream has drained — otherwise a per-session team leaks its + // subprocesses until the server process exits. go func() { + drainCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), sessionDrainTimeout) + defer cancel() + defer sm.stopSessionToolSets(drainCtx, sessionRuntime) ticker := time.NewTicker(100 * time.Millisecond) defer ticker.Stop() - deadline := time.After(5 * time.Minute) for { if sessionRuntime.streaming.TryLock() { sessionRuntime.streaming.Unlock() @@ -730,7 +771,7 @@ func (sm *SessionManager) DeleteSession(ctx context.Context, sessionID string) e return } select { - case <-deadline: + case <-drainCtx.Done(): sm.deletedSessions.Delete(sess.ID) return case <-ticker.C: @@ -800,13 +841,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, @@ -1274,7 +1317,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 @@ -1300,14 +1343,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() @@ -1344,7 +1387,7 @@ func (sm *SessionManager) runtimeForSession(ctx context.Context, sess *session.S } run, err := runtime.New(ctx, t, opts...) if err != nil { - return nil, nil, err + return nil, nil, nil, err } // Give this session an out-of-band, session-scoped route for @@ -1374,7 +1417,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 } func (sm *SessionManager) loadTeam(ctx context.Context, agentFilename string, runConfig *config.RuntimeConfig) (*team.Team, error) { @@ -1898,6 +1941,30 @@ func (sm *SessionManager) BatchDeleteSessions(ctx context.Context, sessionIDs [] if sessionRuntime.cancel != nil { sessionRuntime.cancel() } + // Release the session's toolset-owned resources (stdio MCP + // subprocesses etc.) once its stream drains, mirroring + // DeleteSession — otherwise the per-session team leaks its + // subprocesses until the server process exits. + if sessionRuntime.team != nil { + go func(rs *activeRuntimes) { + drainCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), sessionDrainTimeout) + defer cancel() + defer sm.stopSessionToolSets(drainCtx, rs) + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + if rs.streaming.TryLock() { + rs.streaming.Unlock() + return + } + select { + case <-drainCtx.Done(): + return + case <-ticker.C: + } + } + }(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 6604a4eb5..332a940e6 100644 --- a/pkg/server/session_manager_test.go +++ b/pkg/server/session_manager_test.go @@ -1226,7 +1226,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() }) @@ -1752,7 +1752,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)