Skip to content

fix(server): stop per-session toolsets on session delete #3884

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
EronWright wants to merge 1 commit into
base: main
Choose a base branch
from
+78 −11
Open

fix(server): stop per-session toolsets on session delete #3884

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 76 additions & 9 deletions pkg/server/session_manager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -718,19 +754,24 @@ 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()
sm.deletedSessions.Delete(sess.ID)
return
}
select {
case <-deadline:
case <-drainCtx.Done():
sm.deletedSessions.Delete(sess.ID)
return
case <-ticker.C:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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)

@docker-agent docker-agent Aug 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[medium] BatchDeleteSessions new drain goroutine cannot be awaited via WaitStopped

The PR adds an asynchronous goroutine in BatchDeleteSessions that waits for the session stream to drain and then calls sm.stopSessionToolSets — the same teardown work that DeleteSession does. However, unlike DeleteSession (which stores the runtime in deletedSessions before returning), BatchDeleteSessions never calls sm.deletedSessions.Store(sessionID, sessionRuntime).

WaitStopped reads exclusively from deletedSessions:

rs, ok := sm.deletedSessions.Load(sessionID)
if !ok {
    return nil // already cleaned up
}

So a caller that does BatchDeleteSessions(...) followed by WaitStopped(sessionID, ...) will receive an immediate nil even while the new drain goroutine is still running (waiting to drain the stream and call StopToolSets). Before this PR there was nothing to wait for in BatchDeleteSessions, but now there is — and the wait contract is silently broken.

Impact: Any consumer that uses WaitStopped to synchronise on full teardown (e.g. to confirm MCP subprocesses have exited) after a batch delete will get a false-complete signal. Consider mirroring DeleteSession's deletedSessions.Store call in BatchDeleteSessions (and having the drain goroutine call deletedSessions.Delete when done), or documenting that WaitStopped is not supported after batch deletes.

Confidence Score
🟡 moderate 65/100

}
sm.dropEventLog(sessionID)
Expand Down
4 changes: 2 additions & 2 deletions pkg/server/session_manager_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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() })

Expand Down Expand Up @@ -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)
Expand Down
Loading