Skip to content

Commit a825685

Browse files
charlesviencharlesvien
committed
Review fixes
1 parent 20cfa8a commit a825685

7 files changed

Lines changed: 217 additions & 57 deletions

File tree

‎apps/code/src/main/services/agent/schemas.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -183,9 +183,9 @@ export const subscribeSessionInput = z.object({
183183
taskRunId: z.string(),
184184
});
185185

186-
// Report activity input — keeps the idle timeout debounce alive for the given task
187-
export const reportActivityInput = z.object({
188-
taskId: z.string().nullable(),
186+
// Record activity input — resets the idle timeout for the given session
187+
export const recordActivityInput = z.object({
188+
taskRunId: z.string(),
189189
});
190190

191191
// Agent events

‎apps/code/src/main/services/agent/service.test.ts‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,13 @@ const mockFetch = vi.hoisted(() => vi.fn());
5151

5252
// --- Module mocks ---
5353

54+
const mockPowerMonitor = vi.hoisted(() => ({
55+
on: vi.fn(),
56+
}));
57+
5458
vi.mock("electron", () => ({
5559
app: mockApp,
60+
powerMonitor: mockPowerMonitor,
5661
}));
5762

5863
vi.mock("../../utils/logger.js", () => ({
@@ -280,4 +285,144 @@ describe("AgentService", () => {
280285
);
281286
});
282287
});
288+
289+
describe("idle timeout", () => {
290+
function injectSession(
291+
svc: AgentService,
292+
taskRunId: string,
293+
overrides: Record<string, unknown> = {},
294+
) {
295+
const sessions = (svc as unknown as { sessions: Map<string, unknown> })
296+
.sessions;
297+
sessions.set(taskRunId, {
298+
taskRunId,
299+
taskId: `task-for-${taskRunId}`,
300+
repoPath: "/mock/repo",
301+
agent: { cleanup: vi.fn().mockResolvedValue(undefined) },
302+
clientSideConnection: {},
303+
channel: `ch-${taskRunId}`,
304+
createdAt: Date.now(),
305+
lastActivityAt: Date.now(),
306+
config: {},
307+
needsRecreation: false,
308+
promptPending: false,
309+
...overrides,
310+
});
311+
}
312+
313+
function getIdleTimeouts(svc: AgentService) {
314+
return (
315+
svc as unknown as {
316+
idleTimeouts: Map<
317+
string,
318+
{ handle: ReturnType<typeof setTimeout>; deadline: number }
319+
>;
320+
}
321+
).idleTimeouts;
322+
}
323+
324+
beforeEach(() => {
325+
vi.useFakeTimers();
326+
});
327+
328+
afterEach(() => {
329+
vi.useRealTimers();
330+
});
331+
332+
it("recordActivity is a no-op for unknown sessions", () => {
333+
service.recordActivity("unknown-run");
334+
expect(getIdleTimeouts(service).size).toBe(0);
335+
});
336+
337+
it("recordActivity sets a timeout for a known session", () => {
338+
injectSession(service, "run-1");
339+
service.recordActivity("run-1");
340+
expect(getIdleTimeouts(service).has("run-1")).toBe(true);
341+
});
342+
343+
it("recordActivity resets the timeout on subsequent calls", () => {
344+
injectSession(service, "run-1");
345+
service.recordActivity("run-1");
346+
const firstDeadline = getIdleTimeouts(service).get("run-1")?.deadline;
347+
348+
vi.advanceTimersByTime(5 * 60 * 1000);
349+
service.recordActivity("run-1");
350+
const secondDeadline = getIdleTimeouts(service).get("run-1")?.deadline;
351+
352+
expect(secondDeadline).toBeGreaterThan(firstDeadline!);
353+
});
354+
355+
it("kills idle session after timeout expires", () => {
356+
injectSession(service, "run-1");
357+
service.recordActivity("run-1");
358+
359+
vi.advanceTimersByTime(15 * 60 * 1000);
360+
361+
expect(service.emit).toHaveBeenCalledWith(
362+
"session-idle-killed",
363+
expect.objectContaining({ taskRunId: "run-1" }),
364+
);
365+
});
366+
367+
it("does not kill session if activity is recorded before timeout", () => {
368+
injectSession(service, "run-1");
369+
service.recordActivity("run-1");
370+
371+
vi.advanceTimersByTime(14 * 60 * 1000);
372+
service.recordActivity("run-1");
373+
vi.advanceTimersByTime(14 * 60 * 1000);
374+
375+
expect(service.emit).not.toHaveBeenCalledWith(
376+
"session-idle-killed",
377+
expect.anything(),
378+
);
379+
});
380+
381+
it("skips kill when promptPending is true", () => {
382+
injectSession(service, "run-1", { promptPending: true });
383+
service.recordActivity("run-1");
384+
385+
vi.advanceTimersByTime(15 * 60 * 1000);
386+
387+
expect(service.emit).not.toHaveBeenCalledWith(
388+
"session-idle-killed",
389+
expect.anything(),
390+
);
391+
});
392+
393+
it("checkIdleDeadlines kills expired sessions on resume", () => {
394+
injectSession(service, "run-1");
395+
service.recordActivity("run-1");
396+
397+
const resumeHandler = mockPowerMonitor.on.mock.calls.find(
398+
([event]: string[]) => event === "resume",
399+
)?.[1] as () => void;
400+
expect(resumeHandler).toBeDefined();
401+
402+
vi.advanceTimersByTime(20 * 60 * 1000);
403+
resumeHandler();
404+
405+
expect(service.emit).toHaveBeenCalledWith(
406+
"session-idle-killed",
407+
expect.objectContaining({ taskRunId: "run-1" }),
408+
);
409+
});
410+
411+
it("checkIdleDeadlines does not kill non-expired sessions", () => {
412+
injectSession(service, "run-1");
413+
service.recordActivity("run-1");
414+
415+
const resumeHandler = mockPowerMonitor.on.mock.calls.find(
416+
([event]: string[]) => event === "resume",
417+
)?.[1] as () => void;
418+
419+
vi.advanceTimersByTime(5 * 60 * 1000);
420+
resumeHandler();
421+
422+
expect(service.emit).not.toHaveBeenCalledWith(
423+
"session-idle-killed",
424+
expect.anything(),
425+
);
426+
});
427+
});
283428
});

‎apps/code/src/main/services/agent/service.ts‎

Lines changed: 45 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import { getLlmGatewayUrl } from "@posthog/agent/posthog-api";
2323
import type { OnLogCallback } from "@posthog/agent/types";
2424
import { isAuthError } from "@shared/errors.js";
2525
import type { AcpMessage } from "@shared/types/session-events.js";
26-
import { app } from "electron";
26+
import { app, powerMonitor } from "electron";
2727
import { inject, injectable, preDestroy } from "inversify";
2828
import { MAIN_TOKENS } from "../../di/tokens.js";
2929
import { isDevBuild } from "../../utils/env.js";
@@ -258,7 +258,10 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
258258
private currentToken: string | null = null;
259259
private pendingPermissions = new Map<string, PendingPermission>();
260260
private mockNodeReady = false;
261-
private idleTimeoutHandles = new Map<string, ReturnType<typeof setTimeout>>();
261+
private idleTimeouts = new Map<
262+
string,
263+
{ handle: ReturnType<typeof setTimeout>; deadline: number }
264+
>();
262265
private processTracking: ProcessTrackingService;
263266
private sleepService: SleepService;
264267
private fsService: FsService;
@@ -279,6 +282,8 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
279282
this.sleepService = sleepService;
280283
this.fsService = fsService;
281284
this.posthogPluginService = posthogPluginService;
285+
286+
powerMonitor.on("resume", () => this.checkIdleDeadlines());
282287
}
283288

284289
public updateToken(newToken: string): void {
@@ -397,34 +402,42 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
397402
return false;
398403
}
399404

400-
public reportActivity(taskId: string | null): void {
401-
if (!taskId) return;
402-
for (const session of this.sessions.values()) {
403-
if (session.taskId === taskId) {
404-
this.recordActivity(session.taskRunId);
405-
}
406-
}
407-
}
405+
public recordActivity(taskRunId: string): void {
406+
if (!this.sessions.has(taskRunId)) return;
408407

409-
private recordActivity(taskRunId: string): void {
410-
const existing = this.idleTimeoutHandles.get(taskRunId);
411-
if (existing) clearTimeout(existing);
408+
const existing = this.idleTimeouts.get(taskRunId);
409+
if (existing) clearTimeout(existing.handle);
412410

411+
const deadline = Date.now() + AgentService.IDLE_TIMEOUT_MS;
413412
const handle = setTimeout(() => {
414-
this.idleTimeoutHandles.delete(taskRunId);
415-
const session = this.sessions.get(taskRunId);
416-
if (!session || session.promptPending) return;
417-
log.info("Killing idle session", { taskRunId, taskId: session.taskId });
418-
this.emit(AgentServiceEvent.SessionIdleKilled, {
419-
taskRunId,
420-
taskId: session.taskId,
421-
});
422-
this.cleanupSession(taskRunId).catch((err) => {
423-
log.error("Failed to cleanup idle session", { taskRunId, err });
424-
});
413+
this.killIdleSession(taskRunId);
425414
}, AgentService.IDLE_TIMEOUT_MS);
426415

427-
this.idleTimeoutHandles.set(taskRunId, handle);
416+
this.idleTimeouts.set(taskRunId, { handle, deadline });
417+
}
418+
419+
private killIdleSession(taskRunId: string): void {
420+
const session = this.sessions.get(taskRunId);
421+
if (!session || session.promptPending) return;
422+
log.info("Killing idle session", { taskRunId, taskId: session.taskId });
423+
this.emit(AgentServiceEvent.SessionIdleKilled, {
424+
taskRunId,
425+
taskId: session.taskId,
426+
});
427+
this.cleanupSession(taskRunId).catch((err) => {
428+
log.error("Failed to cleanup idle session", { taskRunId, err });
429+
});
430+
}
431+
432+
private checkIdleDeadlines(): void {
433+
const now = Date.now();
434+
const expired = [...this.idleTimeouts.entries()].filter(
435+
([, { deadline }]) => now >= deadline,
436+
);
437+
for (const [taskRunId, { handle }] of expired) {
438+
clearTimeout(handle);
439+
this.killIdleSession(taskRunId);
440+
}
428441
}
429442

430443
private getToken(fallback: string): string {
@@ -821,6 +834,7 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
821834
};
822835

823836
this.sessions.set(taskRunId, session);
837+
this.recordActivity(taskRunId);
824838
if (isRetry) {
825839
log.info("Session created after auth retry", { taskRunId });
826840
}
@@ -1176,8 +1190,8 @@ For git operations while detached:
11761190

11771191
@preDestroy()
11781192
async cleanupAll(): Promise<void> {
1179-
for (const handle of this.idleTimeoutHandles.values()) clearTimeout(handle);
1180-
this.idleTimeoutHandles.clear();
1193+
for (const { handle } of this.idleTimeouts.values()) clearTimeout(handle);
1194+
this.idleTimeouts.clear();
11811195
const sessionIds = Array.from(this.sessions.keys());
11821196
log.info("Cleaning up all agent sessions", {
11831197
sessionCount: sessionIds.length,
@@ -1265,10 +1279,10 @@ For git operations while detached:
12651279

12661280
this.sessions.delete(taskRunId);
12671281

1268-
const handle = this.idleTimeoutHandles.get(taskRunId);
1269-
if (handle) {
1270-
clearTimeout(handle);
1271-
this.idleTimeoutHandles.delete(taskRunId);
1282+
const timeout = this.idleTimeouts.get(taskRunId);
1283+
if (timeout) {
1284+
clearTimeout(timeout.handle);
1285+
this.idleTimeouts.delete(taskRunId);
12721286
}
12731287
}
12741288
}

‎apps/code/src/main/trpc/routers/agent.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import {
1414
promptInput,
1515
promptOutput,
1616
reconnectSessionInput,
17-
reportActivityInput,
17+
recordActivityInput,
1818
respondToPermissionInput,
1919
sessionResponseSchema,
2020
setConfigOptionInput,
@@ -184,9 +184,9 @@ export const agentRouter = router({
184184
log.info("All sessions reset successfully");
185185
}),
186186

187-
reportActivity: publicProcedure
188-
.input(reportActivityInput)
189-
.mutation(({ input }) => getService().reportActivity(input.taskId)),
187+
recordActivity: publicProcedure
188+
.input(recordActivityInput)
189+
.mutation(({ input }) => getService().recordActivity(input.taskRunId)),
190190

191191
onSessionIdleKilled: publicProcedure.subscription(async function* (opts) {
192192
const service = getService();

‎apps/code/src/renderer/features/sessions/service/service.ts‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -125,15 +125,13 @@ export class SessionService {
125125

126126
constructor() {
127127
this.idleKilledSubscription =
128-
trpcVanilla.agent.onSessionIdleKilled.subscribe(undefined, {
129-
onData: (event) => {
130-
const { taskRunId } = event as { taskRunId: string; taskId: string };
128+
trpcClient.agent.onSessionIdleKilled.subscribe(undefined, {
129+
onData: (event: { taskRunId: string }) => {
130+
const { taskRunId } = event;
131131
log.info("Session idle-killed by main process", { taskRunId });
132-
this.unsubscribeFromChannel(taskRunId);
133-
sessionStoreSetters.removeSession(taskRunId);
134-
removePersistedConfigOptions(taskRunId);
132+
this.teardownSession(taskRunId);
135133
},
136-
onError: (err) => {
134+
onError: (err: unknown) => {
137135
log.debug("Idle-killed subscription error", { error: err });
138136
},
139137
});
@@ -1617,9 +1615,7 @@ export class SessionService {
16171615
throw new Error("Unable to reach server. Please check your connection.");
16181616
}
16191617

1620-
const prefetchedLogs = logUrl
1621-
? await this.fetchSessionLogs(logUrl, taskRunId)
1622-
: { rawEntries: [] as StoredLogEntry[], adapter: undefined };
1618+
const prefetchedLogs = await this.fetchSessionLogs(logUrl, taskRunId);
16231619

16241620
// Determine sessionId: undefined = use from logs, null = strip (fresh), string = use as-is
16251621
const sessionId =

‎apps/code/src/renderer/features/task-detail/components/TaskLogsPanel.tsx‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import { toast } from "@utils/toast";
3434
import { useCallback, useEffect, useMemo, useRef } from "react";
3535

3636
const log = logger.scope("task-logs-panel");
37+
const ACTIVITY_HEARTBEAT_MS = 5 * 60 * 1000;
3738

3839
interface TaskLogsPanelProps {
3940
taskId: string;
@@ -133,15 +134,14 @@ export function TaskLogsPanel({ taskId, task }: TaskLogsPanelProps) {
133134
}, [taskId, requestFocus]);
134135

135136
useEffect(() => {
136-
trpcVanilla.agent.reportActivity.mutate({ taskId }).catch(() => {});
137-
const heartbeat = setInterval(
138-
() => {
139-
trpcVanilla.agent.reportActivity.mutate({ taskId }).catch(() => {});
140-
},
141-
5 * 60 * 1000,
142-
);
137+
const taskRunId = session?.taskRunId;
138+
if (!taskRunId) return;
139+
trpcClient.agent.recordActivity.mutate({ taskRunId }).catch(() => {});
140+
const heartbeat = setInterval(() => {
141+
trpcClient.agent.recordActivity.mutate({ taskRunId }).catch(() => {});
142+
}, ACTIVITY_HEARTBEAT_MS);
143143
return () => clearInterval(heartbeat);
144-
}, [taskId]);
144+
}, [session?.taskRunId]);
145145

146146
// Keep cloud session title aligned with latest task metadata.
147147
useEffect(() => {

0 commit comments

Comments
 (0)