diff --git a/pkg/server/author_safety_test.go b/pkg/server/author_safety_test.go index 866e732fc..5535560b2 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 1b77dab61..1bcb614e2 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 595a565cf..0947826f7 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