Preserve current cloud permission retries

Generated-By: PostHog Code
Task-Id: 40c57a59-b4e1-4760-8e56-ecd03e9c2f0f

Author: richardsolomou

packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts

@@ -354,6 +354,10 @@ describe("cloud task update notifications", () => {
       toolCall: { toolCallId: "t1", title: "Run command", kind: "execute" },
       options: [],
     });
+    const session = harness.sessions[RUN_ID];
+    const permission = session?.pendingPermissions.get("t1");
+    if (permission) permission.receivedAt = Date.now() - 1_000;
+    if (session) session.pausedDurationMs = 25;
     harness.sendCommand.mockResolvedValueOnce({
       success: false,
       error: "sandbox rejected response",
@@ -362,6 +366,7 @@ describe("cloud task update notifications", () => {
     await harness.service.respondToPermission(TASK_ID, "t1", "allow");
 
     expect(harness.sessions[RUN_ID]?.pendingPermissions.has("t1")).toBe(true);
+    expect(harness.sessions[RUN_ID]?.pausedDurationMs).toBe(25);
 
     await harness.service.respondToPermission(TASK_ID, "t1", "allow");
     expect(harness.sendCommand).toHaveBeenLastCalledWith(
@@ -371,4 +376,48 @@ describe("cloud task update notifications", () => {
       }),
     );
   });
+
+  it("does not restore a stale permission over a newer request", async () => {
+    const harness = createHarness();
+    harness.sendUpdate({
+      taskId: TASK_ID,
+      runId: RUN_ID,
+      kind: "permission_request",
+      requestId: "r1",
+      toolCall: { toolCallId: "t1", title: "Run command", kind: "execute" },
+      options: [],
+    });
+    let rejectResponse: (error: Error) => void = () => undefined;
+    harness.sendCommand.mockImplementationOnce(
+      () =>
+        new Promise((_, reject) => {
+          rejectResponse = reject;
+        }),
+    );
+
+    const response = harness.service.respondToPermission(
+      TASK_ID,
+      "t1",
+      "allow",
+    );
+    await vi.waitFor(() => expect(harness.sendCommand).toHaveBeenCalled());
+    harness.sendUpdate({
+      taskId: TASK_ID,
+      runId: RUN_ID,
+      kind: "permission_request",
+      requestId: "r2",
+      toolCall: { toolCallId: "t1", title: "Run command", kind: "execute" },
+      options: [],
+    });
+    rejectResponse(new Error("failed"));
+    await response;
+
+    await harness.service.respondToPermission(TASK_ID, "t1", "allow");
+    expect(harness.sendCommand).toHaveBeenLastCalledWith(
+      expect.objectContaining({
+        method: "permission_response",
+        params: { requestId: "r2", optionId: "allow" },
+      }),
+    );
+  });
 });

packages/core/src/sessions/sessionService.ts

@@ -4763,19 +4763,28 @@ export class SessionService {
 
   // --- Permissions ---
 
-  private resolvePermission(session: AgentSession, toolCallId: string): void {
+  private resolvePermission(
+    session: AgentSession,
+    toolCallId: string,
+  ): {
+    previousPausedDurationMs: number;
+    resolvedPausedDurationMs: number;
+  } {
     const permission = session.pendingPermissions.get(toolCallId);
     const newPermissions = new Map(session.pendingPermissions);
     newPermissions.delete(toolCallId);
     this.d.store.setPendingPermissions(session.taskRunId, newPermissions);
 
+    const previousPausedDurationMs = session.pausedDurationMs ?? 0;
+    const resolvedPausedDurationMs = permission?.receivedAt
+      ? previousPausedDurationMs + (Date.now() - permission.receivedAt)
+      : previousPausedDurationMs;
     if (permission?.receivedAt) {
       this.d.store.updateSession(session.taskRunId, {
-        pausedDurationMs:
-          (session.pausedDurationMs ?? 0) +
-          (Date.now() - permission.receivedAt),
+        pausedDurationMs: resolvedPausedDurationMs,
       });
     }
+    return { previousPausedDurationMs, resolvedPausedDurationMs };
   }
 
   /**
@@ -4801,7 +4810,7 @@ export class SessionService {
     });
 
     const cloudRequestId = this.cloudPermissionRequestIds.get(toolCallId);
-    this.resolvePermission(session, toolCallId);
+    const permissionResolution = this.resolvePermission(session, toolCallId);
 
     try {
       const refreshedCloudStatus =
@@ -4877,15 +4886,25 @@ export class SessionService {
       });
     } catch (error) {
       const latestSession = this.d.store.getSessionByTaskId(taskId);
-      if (permission && latestSession) {
+      if (
+        latestSession?.pausedDurationMs ===
+        permissionResolution.resolvedPausedDurationMs
+      ) {
+        this.d.store.updateSession(latestSession.taskRunId, {
+          pausedDurationMs: permissionResolution.previousPausedDurationMs,
+        });
+      }
+      const currentCloudRequestId =
+        this.cloudPermissionRequestIds.get(toolCallId);
+      if (permission && latestSession && !currentCloudRequestId) {
         const restoredPermissions = new Map(latestSession.pendingPermissions);
         restoredPermissions.set(toolCallId, permission);
         this.d.store.setPendingPermissions(
           latestSession.taskRunId,
           restoredPermissions,
         );
       }
-      if (cloudRequestId) {
+      if (cloudRequestId && !currentCloudRequestId) {
         this.cloudPermissionRequestIds.set(toolCallId, cloudRequestId);
       }
       this.d.log.error("Failed to respond to permission", {