Skip to content

Commit 240ae1b

Browse files
committed
Fix logout and project switch to properly kill agents, shells, and processes
1 parent 8d8f8a0 commit 240ae1b

4 files changed

Lines changed: 55 additions & 1 deletion

File tree

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import type { ContentBlock } from "@agentclientprotocol/sdk";
22
import { container } from "../../di/container.js";
33
import { MAIN_TOKENS } from "../../di/tokens.js";
4+
import { logger } from "../../lib/logger.js";
45
import {
56
AgentServiceEvent,
67
cancelPermissionInput,
@@ -22,8 +23,13 @@ import {
2223
tokenUpdateInput,
2324
} from "../../services/agent/schemas.js";
2425
import type { AgentService } from "../../services/agent/service.js";
26+
import type { ProcessTrackingService } from "../../services/process-tracking/service.js";
27+
import type { ShellService } from "../../services/shell/service.js";
28+
import type { SleepService } from "../../services/sleep/service.js";
2529
import { publicProcedure, router } from "../trpc.js";
2630

31+
const log = logger.scope("agent-router");
32+
2733
const getService = () => container.get<AgentService>(MAIN_TOKENS.AgentService);
2834

2935
export const agentRouter = router({
@@ -140,6 +146,30 @@ export const agentRouter = router({
140146
getService().markAllSessionsForRecreation(),
141147
),
142148

149+
resetAll: publicProcedure.mutation(async () => {
150+
log.info("Resetting all sessions (logout/project switch)");
151+
152+
// 1. Clean up all agent sessions (flushes logs, stops agents, releases sleep blockers)
153+
const agentService = getService();
154+
await agentService.cleanupAll();
155+
156+
// 2. Destroy all shell PTY sessions
157+
const shellService = container.get<ShellService>(MAIN_TOKENS.ShellService);
158+
shellService.destroyAll();
159+
160+
// 3. Kill any remaining tracked processes (belt and suspenders)
161+
const processTracking = container.get<ProcessTrackingService>(
162+
MAIN_TOKENS.ProcessTrackingService,
163+
);
164+
processTracking.killAll();
165+
166+
// 4. Release any lingering sleep blockers
167+
const sleepService = container.get<SleepService>(MAIN_TOKENS.SleepService);
168+
sleepService.cleanup();
169+
170+
log.info("All sessions reset successfully");
171+
}),
172+
143173
getGatewayModels: publicProcedure
144174
.input(getGatewayModelsInput)
145175
.output(getGatewayModelsOutput)

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -643,6 +643,9 @@ export const useAuthStore = create<AuthState>()(
643643
throw new Error("No access token available");
644644
}
645645

646+
// Clean up all existing sessions before switching projects
647+
resetSessionService();
648+
646649
const apiHost = getCloudUrlFromRegion(cloudRegion);
647650

648651
// Create a new client with the selected project

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,13 +74,27 @@ export function getSessionService(): SessionService {
7474

7575
/**
7676
* Reset the session service singleton.
77-
* Call this on logout or when the app needs to fully reset state.
77+
* Call this on logout or project switch to fully clean up all sessions.
78+
*
79+
* This does three things:
80+
* 1. Unsubscribes from all renderer-side tRPC subscriptions
81+
* 2. Clears all sessions from the renderer session store
82+
* 3. Tells the main process to kill all agents, shells, and tracked processes
7883
*/
7984
export function resetSessionService(): void {
8085
if (serviceInstance) {
8186
serviceInstance.reset();
8287
serviceInstance = null;
8388
}
89+
90+
// Clear all sessions from the renderer store
91+
sessionStoreSetters.clearAll();
92+
93+
// Tell the main process to clean up all agents, shells, and processes.
94+
// Fire-and-forget: the main process handles cleanup independently.
95+
trpcVanilla.agent.resetAll.mutate().catch((err) => {
96+
log.error("Failed to reset all sessions on main process", err);
97+
});
8498
}
8599

86100
export class SessionService {

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -289,4 +289,11 @@ export const sessionStoreSetters = {
289289
getSessions: (): Record<string, AgentSession> => {
290290
return useSessionStore.getState().sessions;
291291
},
292+
293+
clearAll: () => {
294+
useSessionStore.setState((state) => {
295+
state.sessions = {};
296+
state.taskIdIndex = {};
297+
});
298+
},
292299
};

0 commit comments

Comments
 (0)