Review fixes

Author: charlesvien

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

@@ -183,9 +183,9 @@ export const subscribeSessionInput = z.object({
   taskRunId: z.string(),
 });
 
-// Report activity input — keeps the idle timeout debounce alive for the given task
-export const reportActivityInput = z.object({
-  taskId: z.string().nullable(),
+// Record activity input — resets the idle timeout for the given session
+export const recordActivityInput = z.object({
+  taskRunId: z.string(),
 });
 
 // Agent events

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

@@ -51,8 +51,13 @@ const mockFetch = vi.hoisted(() => vi.fn());
 
 // --- Module mocks ---
 
+const mockPowerMonitor = vi.hoisted(() => ({
+  on: vi.fn(),
+}));
+
 vi.mock("electron", () => ({
   app: mockApp,
+  powerMonitor: mockPowerMonitor,
 }));
 
 vi.mock("../../utils/logger.js", () => ({

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

@@ -23,7 +23,7 @@ import { getLlmGatewayUrl } from "@posthog/agent/posthog-api";
 import type { OnLogCallback } from "@posthog/agent/types";
 import { isAuthError } from "@shared/errors.js";
 import type { AcpMessage } from "@shared/types/session-events.js";
-import { app } from "electron";
+import { app, powerMonitor } from "electron";
 import { inject, injectable, preDestroy } from "inversify";
 import { MAIN_TOKENS } from "../../di/tokens.js";
 import { isDevBuild } from "../../utils/env.js";
@@ -258,7 +258,10 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
   private currentToken: string | null = null;
   private pendingPermissions = new Map<string, PendingPermission>();
   private mockNodeReady = false;
-  private idleTimeoutHandles = new Map<string, ReturnType<typeof setTimeout>>();
+  private idleTimeouts = new Map<
+    string,
+    { handle: ReturnType<typeof setTimeout>; deadline: number }
+  >();
   private processTracking: ProcessTrackingService;
   private sleepService: SleepService;
   private fsService: FsService;
@@ -279,6 +282,8 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
     this.sleepService = sleepService;
     this.fsService = fsService;
     this.posthogPluginService = posthogPluginService;
+
+    powerMonitor.on("resume", () => this.checkIdleDeadlines());
   }
 
   public updateToken(newToken: string): void {
@@ -397,34 +402,43 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
     return false;
   }
 
-  public reportActivity(taskId: string | null): void {
-    if (!taskId) return;
-    for (const session of this.sessions.values()) {
-      if (session.taskId === taskId) {
-        this.recordActivity(session.taskRunId);
-      }
-    }
-  }
+  public recordActivity(taskRunId: string): void {
+    if (!this.sessions.has(taskRunId)) return;
 
-  private recordActivity(taskRunId: string): void {
-    const existing = this.idleTimeoutHandles.get(taskRunId);
-    if (existing) clearTimeout(existing);
+    const existing = this.idleTimeouts.get(taskRunId);
+    if (existing) clearTimeout(existing.handle);
 
+    const deadline = Date.now() + AgentService.IDLE_TIMEOUT_MS;
     const handle = setTimeout(() => {
-      this.idleTimeoutHandles.delete(taskRunId);
-      const session = this.sessions.get(taskRunId);
-      if (!session || session.promptPending) return;
-      log.info("Killing idle session", { taskRunId, taskId: session.taskId });
-      this.emit(AgentServiceEvent.SessionIdleKilled, {
-        taskRunId,
-        taskId: session.taskId,
-      });
-      this.cleanupSession(taskRunId).catch((err) => {
-        log.error("Failed to cleanup idle session", { taskRunId, err });
-      });
+      this.killIdleSession(taskRunId);
     }, AgentService.IDLE_TIMEOUT_MS);
 
-    this.idleTimeoutHandles.set(taskRunId, handle);
+    this.idleTimeouts.set(taskRunId, { handle, deadline });
+  }
+
+  private killIdleSession(taskRunId: string): void {
+    this.idleTimeouts.delete(taskRunId);
+    const session = this.sessions.get(taskRunId);
+    if (!session || session.promptPending) return;
+    log.info("Killing idle session", { taskRunId, taskId: session.taskId });
+    this.emit(AgentServiceEvent.SessionIdleKilled, {
+      taskRunId,
+      taskId: session.taskId,
+    });
+    this.cleanupSession(taskRunId).catch((err) => {
+      log.error("Failed to cleanup idle session", { taskRunId, err });
+    });
+  }
+
+  private checkIdleDeadlines(): void {
+    const now = Date.now();
+    const expired = [...this.idleTimeouts.entries()].filter(
+      ([, { deadline }]) => now >= deadline,
+    );
+    for (const [taskRunId, { handle }] of expired) {
+      clearTimeout(handle);
+      this.killIdleSession(taskRunId);
+    }
   }
 
   private getToken(fallback: string): string {
@@ -821,6 +835,7 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
       };
 
       this.sessions.set(taskRunId, session);
+      this.recordActivity(taskRunId);
       if (isRetry) {
         log.info("Session created after auth retry", { taskRunId });
       }
@@ -1176,8 +1191,8 @@ For git operations while detached:
 
   @preDestroy()
   async cleanupAll(): Promise<void> {
-    for (const handle of this.idleTimeoutHandles.values()) clearTimeout(handle);
-    this.idleTimeoutHandles.clear();
+    for (const { handle } of this.idleTimeouts.values()) clearTimeout(handle);
+    this.idleTimeouts.clear();
     const sessionIds = Array.from(this.sessions.keys());
     log.info("Cleaning up all agent sessions", {
       sessionCount: sessionIds.length,
@@ -1265,10 +1280,10 @@ For git operations while detached:
 
       this.sessions.delete(taskRunId);
 
-      const handle = this.idleTimeoutHandles.get(taskRunId);
-      if (handle) {
-        clearTimeout(handle);
-        this.idleTimeoutHandles.delete(taskRunId);
+      const timeout = this.idleTimeouts.get(taskRunId);
+      if (timeout) {
+        clearTimeout(timeout.handle);
+        this.idleTimeouts.delete(taskRunId);
       }
     }
   }

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

@@ -14,7 +14,7 @@ import {
   promptInput,
   promptOutput,
   reconnectSessionInput,
-  reportActivityInput,
+  recordActivityInput,
   respondToPermissionInput,
   sessionResponseSchema,
   setConfigOptionInput,
@@ -184,9 +184,9 @@ export const agentRouter = router({
     log.info("All sessions reset successfully");
   }),
 
-  reportActivity: publicProcedure
-    .input(reportActivityInput)
-    .mutation(({ input }) => getService().reportActivity(input.taskId)),
+  recordActivity: publicProcedure
+    .input(recordActivityInput)
+    .mutation(({ input }) => getService().recordActivity(input.taskRunId)),
 
   onSessionIdleKilled: publicProcedure.subscription(async function* (opts) {
     const service = getService();

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

@@ -125,15 +125,13 @@ export class SessionService {
 
   constructor() {
     this.idleKilledSubscription =
-      trpcVanilla.agent.onSessionIdleKilled.subscribe(undefined, {
-        onData: (event) => {
-          const { taskRunId } = event as { taskRunId: string; taskId: string };
+      trpcClient.agent.onSessionIdleKilled.subscribe(undefined, {
+        onData: (event: { taskRunId: string }) => {
+          const { taskRunId } = event;
           log.info("Session idle-killed by main process", { taskRunId });
-          this.unsubscribeFromChannel(taskRunId);
-          sessionStoreSetters.removeSession(taskRunId);
-          removePersistedConfigOptions(taskRunId);
+          this.teardownSession(taskRunId);
         },
-        onError: (err) => {
+        onError: (err: unknown) => {
           log.debug("Idle-killed subscription error", { error: err });
         },
       });
@@ -1617,9 +1615,7 @@ export class SessionService {
       throw new Error("Unable to reach server. Please check your connection.");
     }
 
-    const prefetchedLogs = logUrl
-      ? await this.fetchSessionLogs(logUrl, taskRunId)
-      : { rawEntries: [] as StoredLogEntry[], adapter: undefined };
+    const prefetchedLogs = await this.fetchSessionLogs(logUrl, taskRunId);
 
     // Determine sessionId: undefined = use from logs, null = strip (fresh), string = use as-is
     const sessionId =

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

@@ -34,6 +34,7 @@ import { toast } from "@utils/toast";
 import { useCallback, useEffect, useMemo, useRef } from "react";
 
 const log = logger.scope("task-logs-panel");
+const ACTIVITY_HEARTBEAT_MS = 5 * 60 * 1000;
 
 interface TaskLogsPanelProps {
   taskId: string;
@@ -133,15 +134,14 @@ export function TaskLogsPanel({ taskId, task }: TaskLogsPanelProps) {
   }, [taskId, requestFocus]);
 
   useEffect(() => {
-    trpcVanilla.agent.reportActivity.mutate({ taskId }).catch(() => {});
-    const heartbeat = setInterval(
-      () => {
-        trpcVanilla.agent.reportActivity.mutate({ taskId }).catch(() => {});
-      },
-      5 * 60 * 1000,
-    );
+    const taskRunId = session?.taskRunId;
+    if (!taskRunId) return;
+    trpcClient.agent.recordActivity.mutate({ taskRunId }).catch(() => {});
+    const heartbeat = setInterval(() => {
+      trpcClient.agent.recordActivity.mutate({ taskRunId }).catch(() => {});
+    }, ACTIVITY_HEARTBEAT_MS);
     return () => clearInterval(heartbeat);
-  }, [taskId]);
+  }, [session?.taskRunId]);
 
   // Keep cloud session title aligned with latest task metadata.
   useEffect(() => {

packages/agent/src/adapters/claude/tools.ts

@@ -26,7 +26,12 @@ export const SEARCH_TOOLS: Set<string> = new Set(["Glob", "Grep", "LS"]);
 
 export const WEB_TOOLS: Set<string> = new Set(["WebSearch", "WebFetch"]);
 
-export const AGENT_TOOLS: Set<string> = new Set(["Task", "TodoWrite", "Skill"]);
+export const AGENT_TOOLS: Set<string> = new Set([
+  "Task",
+  "Agent",
+  "TodoWrite",
+  "Skill",
+]);
 
 const BASE_ALLOWED_TOOLS = [
   ...READ_TOOLS,