fix: Fix logout and project switch to properly kill agents, shells, and processes (#832)

Author: charlesvien

apps/twig/src/main/trpc/routers/agent.ts

@@ -1,6 +1,7 @@
 import type { ContentBlock } from "@agentclientprotocol/sdk";
 import { container } from "../../di/container.js";
 import { MAIN_TOKENS } from "../../di/tokens.js";
+import { logger } from "../../lib/logger.js";
 import {
   AgentServiceEvent,
   cancelPermissionInput,
@@ -22,8 +23,13 @@ import {
   tokenUpdateInput,
 } from "../../services/agent/schemas.js";
 import type { AgentService } from "../../services/agent/service.js";
+import type { ProcessTrackingService } from "../../services/process-tracking/service.js";
+import type { ShellService } from "../../services/shell/service.js";
+import type { SleepService } from "../../services/sleep/service.js";
 import { publicProcedure, router } from "../trpc.js";
 
+const log = logger.scope("agent-router");
+
 const getService = () => container.get<AgentService>(MAIN_TOKENS.AgentService);
 
 export const agentRouter = router({
@@ -140,6 +146,30 @@ export const agentRouter = router({
     getService().markAllSessionsForRecreation(),
   ),
 
+  resetAll: publicProcedure.mutation(async () => {
+    log.info("Resetting all sessions (logout/project switch)");
+
+    // Clean up all agent sessions (flushes logs, stops agents, releases sleep blockers)
+    const agentService = getService();
+    await agentService.cleanupAll();
+
+    // Destroy all shell PTY sessions
+    const shellService = container.get<ShellService>(MAIN_TOKENS.ShellService);
+    shellService.destroyAll();
+
+    // Kill any remaining tracked processes (belt and suspenders)
+    const processTracking = container.get<ProcessTrackingService>(
+      MAIN_TOKENS.ProcessTrackingService,
+    );
+    processTracking.killAll();
+
+    // Release any lingering sleep blockers
+    const sleepService = container.get<SleepService>(MAIN_TOKENS.SleepService);
+    sleepService.cleanup();
+
+    log.info("All sessions reset successfully");
+  }),
+
   getGatewayModels: publicProcedure
     .input(getGatewayModelsInput)
     .output(getGatewayModelsOutput)

apps/twig/src/renderer/features/auth/stores/authStore.ts

@@ -643,6 +643,9 @@ export const useAuthStore = create<AuthState>()(
             throw new Error("No access token available");
           }
 
+          // Clean up all existing sessions before switching projects
+          resetSessionService();
+
           const apiHost = getCloudUrlFromRegion(cloudRegion);
 
           // Create a new client with the selected project

apps/twig/src/renderer/features/sessions/service/service.test.ts

@@ -15,6 +15,7 @@ const mockTrpcAgent = vi.hoisted(() => ({
   cancelPermission: { mutate: vi.fn() },
   onSessionEvent: { subscribe: vi.fn() },
   onPermissionRequest: { subscribe: vi.fn() },
+  resetAll: { mutate: vi.fn().mockResolvedValue(undefined) },
 }));
 
 const mockTrpcWorkspace = vi.hoisted(() => ({
@@ -44,6 +45,7 @@ const mockSessionStoreSetters = vi.hoisted(() => ({
   setPendingPermissions: vi.fn(),
   getSessionByTaskId: vi.fn(),
   getSessions: vi.fn(() => ({})),
+  clearAll: vi.fn(),
 }));
 
 vi.mock("@features/sessions/stores/sessionStore", () => ({

apps/twig/src/renderer/features/sessions/service/service.ts

@@ -73,15 +73,17 @@ export function getSessionService(): SessionService {
   return serviceInstance;
 }
 
-/**
- * Reset the session service singleton.
- * Call this on logout or when the app needs to fully reset state.
- */
 export function resetSessionService(): void {
   if (serviceInstance) {
     serviceInstance.reset();
     serviceInstance = null;
   }
+
+  sessionStoreSetters.clearAll();
+
+  trpcVanilla.agent.resetAll.mutate().catch((err) => {
+    log.error("Failed to reset all sessions on main process", err);
+  });
 }
 
 export class SessionService {

apps/twig/src/renderer/features/sessions/stores/sessionStore.ts

@@ -289,4 +289,11 @@ export const sessionStoreSetters = {
   getSessions: (): Record<string, AgentSession> => {
     return useSessionStore.getState().sessions;
   },
+
+  clearAll: () => {
+    useSessionStore.setState((state) => {
+      state.sessions = {};
+      state.taskIdIndex = {};
+    });
+  },
 };