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)