Skip to content

Commit e19656c

Browse files
charlesviencharlesvien
committed
Review fixes
1 parent 4585a0f commit e19656c

7 files changed

Lines changed: 226 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: 146 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,145 @@ 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("reschedules when promptPending is true at timeout", () => {
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+
expect(getIdleTimeouts(service).has("run-1")).toBe(true);
392+
});
393+
394+
it("checkIdleDeadlines kills expired sessions on resume", () => {
395+
injectSession(service, "run-1");
396+
service.recordActivity("run-1");
397+
398+
const resumeHandler = mockPowerMonitor.on.mock.calls.find(
399+
([event]: string[]) => event === "resume",
400+
)?.[1] as () => void;
401+
expect(resumeHandler).toBeDefined();
402+
403+
vi.advanceTimersByTime(20 * 60 * 1000);
404+
resumeHandler();
405+
406+
expect(service.emit).toHaveBeenCalledWith(
407+
"session-idle-killed",
408+
expect.objectContaining({ taskRunId: "run-1" }),
409+
);
410+
});
411+
412+
it("checkIdleDeadlines does not kill non-expired sessions", () => {
413+
injectSession(service, "run-1");
414+
service.recordActivity("run-1");
415+
416+
const resumeHandler = mockPowerMonitor.on.mock.calls.find(
417+
([event]: string[]) => event === "resume",
418+
)?.[1] as () => void;
419+
420+
vi.advanceTimersByTime(5 * 60 * 1000);
421+
resumeHandler();
422+
423+
expect(service.emit).not.toHaveBeenCalledWith(
424+
"session-idle-killed",
425+
expect.anything(),
426+
);
427+
});
428+
});
283429
});

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

Lines changed: 49 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,46 @@ 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) return;
422+
if (session.promptPending) {
423+
this.recordActivity(taskRunId);
424+
return;
425+
}
426+
log.info("Killing idle session", { taskRunId, taskId: session.taskId });
427+
this.emit(AgentServiceEvent.SessionIdleKilled, {
428+
taskRunId,
429+
taskId: session.taskId,
430+
});
431+
this.cleanupSession(taskRunId).catch((err) => {
432+
log.error("Failed to cleanup idle session", { taskRunId, err });
433+
});
434+
}
435+
436+
private checkIdleDeadlines(): void {
437+
const now = Date.now();
438+
const expired = [...this.idleTimeouts.entries()].filter(
439+
([, { deadline }]) => now >= deadline,
440+
);
441+
for (const [taskRunId, { handle }] of expired) {
442+
clearTimeout(handle);
443+
this.killIdleSession(taskRunId);
444+
}
428445
}
429446

430447
private getToken(fallback: string): string {
@@ -821,6 +838,7 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
821838
};
822839

823840
this.sessions.set(taskRunId, session);
841+
this.recordActivity(taskRunId);
824842
if (isRetry) {
825843
log.info("Session created after auth retry", { taskRunId });
826844
}
@@ -1176,8 +1194,8 @@ For git operations while detached:
11761194

11771195
@preDestroy()
11781196
async cleanupAll(): Promise<void> {
1179-
for (const handle of this.idleTimeoutHandles.values()) clearTimeout(handle);
1180-
this.idleTimeoutHandles.clear();
1197+
for (const { handle } of this.idleTimeouts.values()) clearTimeout(handle);
1198+
this.idleTimeouts.clear();
11811199
const sessionIds = Array.from(this.sessions.keys());
11821200
log.info("Cleaning up all agent sessions", {
11831201
sessionCount: sessionIds.length,
@@ -1265,10 +1283,10 @@ For git operations while detached:
12651283

12661284
this.sessions.delete(taskRunId);
12671285

1268-
const handle = this.idleTimeoutHandles.get(taskRunId);
1269-
if (handle) {
1270-
clearTimeout(handle);
1271-
this.idleTimeoutHandles.delete(taskRunId);
1286+
const timeout = this.idleTimeouts.get(taskRunId);
1287+
if (timeout) {
1288+
clearTimeout(timeout.handle);
1289+
this.idleTimeouts.delete(taskRunId);
12721290
}
12731291
}
12741292
}

‎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
@@ -127,15 +127,13 @@ export class SessionService {
127127

128128
constructor() {
129129
this.idleKilledSubscription =
130-
trpcVanilla.agent.onSessionIdleKilled.subscribe(undefined, {
131-
onData: (event) => {
132-
const { taskRunId } = event as { taskRunId: string; taskId: string };
130+
trpcClient.agent.onSessionIdleKilled.subscribe(undefined, {
131+
onData: (event: { taskRunId: string }) => {
132+
const { taskRunId } = event;
133133
log.info("Session idle-killed by main process", { taskRunId });
134-
this.unsubscribeFromChannel(taskRunId);
135-
sessionStoreSetters.removeSession(taskRunId);
136-
removePersistedConfigOptions(taskRunId);
134+
this.teardownSession(taskRunId);
137135
},
138-
onError: (err) => {
136+
onError: (err: unknown) => {
139137
log.debug("Idle-killed subscription error", { error: err });
140138
},
141139
});
@@ -1736,9 +1734,7 @@ export class SessionService {
17361734
throw new Error("Unable to reach server. Please check your connection.");
17371735
}
17381736

1739-
const prefetchedLogs = logUrl
1740-
? await this.fetchSessionLogs(logUrl, taskRunId)
1741-
: { rawEntries: [] as StoredLogEntry[], adapter: undefined };
1737+
const prefetchedLogs = await this.fetchSessionLogs(logUrl, taskRunId);
17421738

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

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

Lines changed: 12 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;
@@ -130,15 +131,18 @@ export function TaskLogsPanel({ taskId, task }: TaskLogsPanelProps) {
130131
}, [taskId, requestFocus]);
131132

132133
useEffect(() => {
133-
trpcVanilla.agent.reportActivity.mutate({ taskId }).catch(() => {});
134-
const heartbeat = setInterval(
135-
() => {
136-
trpcVanilla.agent.reportActivity.mutate({ taskId }).catch(() => {});
137-
},
138-
5 * 60 * 1000,
139-
);
134+
const taskRunId = session?.taskRunId;
135+
if (!taskRunId) return;
136+
trpcClient.agent.recordActivity
137+
.mutate({ taskRunId })
138+
.catch((err) => log.debug("Failed to record activity", { err }));
139+
const heartbeat = setInterval(() => {
140+
trpcClient.agent.recordActivity
141+
.mutate({ taskRunId })
142+
.catch((err) => log.debug("Failed to record activity", { err }));
143+
}, ACTIVITY_HEARTBEAT_MS);
140144
return () => clearInterval(heartbeat);
141-
}, [taskId]);
145+
}, [session?.taskRunId]);
142146

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

0 commit comments

Comments
 (0)