fix: Serve session manager shutdown can hang forever after the bounded HTTP shutdown returns

cmd/serve.go

@@ -493,7 +493,11 @@ func runServeLegacy(parentCtx context.Context, cmd *cobra.Command, args []string
 		return runtimeFactory(ctx, "", "")
 	}
 	sessionMgr := newServeSessionManager(serveSessionTTL, serveSessionMax, factory)
-	defer sessionMgr.Close()
+	defer func() {
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		sessionMgr.CloseContext(shutdownCtx)
+	}()
 
 	if hasJobs && strings.TrimSpace(serveAgent) != "" {
 		fmt.Fprintln(cmd.ErrOrStderr(), "warning: --agent is ignored for --platform jobs; set llm runner_config.agent_name per job definition")

cmd/serve_runtime.go

@@ -218,15 +218,54 @@ func (rt *serveRuntime) configureContextManagementForRequest(req llm.Request) {
 }
 
 func (rt *serveRuntime) Close() {
+	rt.CloseContext(context.Background())
+}
+
+func (rt *serveRuntime) CloseContext(ctx context.Context) {
 	rt.interruptMu.Lock()
 	state := rt.activeInterrupt
 	rt.interruptMu.Unlock()
 	if state != nil && state.cancel != nil {
 		state.cancel()
 	}
 
-	rt.mu.Lock()
+	if !rt.lockForClose(ctx) {
+		return
+	}
 	defer rt.mu.Unlock()
+	rt.closeLocked()
+}
+
+// lockForClose waits for the runtime mutex, but lets CloseContext abandon the
+// wait when a provider/tool run keeps holding rt.mu after cancellation.
+func (rt *serveRuntime) lockForClose(ctx context.Context) bool {
+	if ctx == nil {
+		rt.mu.Lock()
+		return true
+	}
+	if rt.mu.TryLock() {
+		return true
+	}
+	done := ctx.Done()
+	if done == nil {
+		rt.mu.Lock()
+		return true
+	}
+	ticker := time.NewTicker(10 * time.Millisecond)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-done:
+			return false
+		case <-ticker.C:
+			if rt.mu.TryLock() {
+				return true
+			}
+		}
+	}
+}
+
+func (rt *serveRuntime) closeLocked() {
 	rt.clearPendingAskUsers()
 	rt.clearPendingApprovals()
 	if rt.mcpManager != nil {

cmd/serve_session.go

@@ -546,6 +546,10 @@ func (m *serveSessionManager) ActiveSessionIDs() map[string]bool {
 }
 
 func (m *serveSessionManager) Close() {
+	m.CloseContext(context.Background())
+}
+
+func (m *serveSessionManager) CloseContext(ctx context.Context) {
 	m.mu.Lock()
 	if m.closed {
 		m.mu.Unlock()
@@ -560,10 +564,13 @@ func (m *serveSessionManager) Close() {
 	m.sessions = map[string]*serveRuntime{}
 	m.mu.Unlock()
 
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	for _, rt := range sessions {
 		if m.onEvict != nil {
 			m.onEvict(rt)
 		}
-		rt.Close()
+		rt.CloseContext(ctx)
 	}
 }

cmd/serve_test.go

@@ -1902,6 +1902,59 @@ func TestServeSessionManager_EvictExpiredSkipsActiveRun(t *testing.T) {
 	}
 }
 
+func TestServeSessionManager_CloseContextDoesNotBlockOnStuckRuntime(t *testing.T) {
+	manager := newServeSessionManager(time.Minute, 10, nil)
+	t.Cleanup(manager.Close)
+
+	runCtx, runCancel := context.WithCancel(context.Background())
+	rt := &serveRuntime{}
+	state := &runtimeInterruptState{cancel: runCancel, done: make(chan struct{})}
+	rt.setActiveInterrupt(state)
+	rt.mu.Lock()
+	t.Cleanup(func() {
+		rt.mu.Unlock()
+		rt.clearActiveInterrupt(state)
+		runCancel()
+	})
+	putTestSession(manager, "busy", rt)
+
+	closeCtx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
+	defer cancel()
+
+	done := make(chan struct{})
+	started := time.Now()
+	go func() {
+		manager.CloseContext(closeCtx)
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		if elapsed := time.Since(started); elapsed > 500*time.Millisecond {
+			t.Fatalf("CloseContext took %s with an expired close context", elapsed)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("CloseContext blocked on a runtime that never released its mutex")
+	}
+
+	select {
+	case <-runCtx.Done():
+	default:
+		t.Fatal("CloseContext did not cancel the active runtime before waiting")
+	}
+
+	manager.mu.Lock()
+	closed := manager.closed
+	sessionCount := len(manager.sessions)
+	manager.mu.Unlock()
+	if !closed {
+		t.Fatal("expected manager to be marked closed")
+	}
+	if sessionCount != 0 {
+		t.Fatalf("sessions after CloseContext = %d, want 0", sessionCount)
+	}
+}
+
 func TestServeSessionManager_GetOrCreateSkipsEvictingActiveRunAtCapacity(t *testing.T) {
 	manager := newServeSessionManager(time.Minute, 2, func(ctx context.Context) (*serveRuntime, error) {
 		rt := &serveRuntime{}