fix(agent): wait for first-turn cloud attachments

Poll the task run artifact manifest with bounded backoff before building the first cloud prompt. This prevents a just-uploaded pasted-text attachment from being omitted when the initial manifest read is briefly stale.

Generated-By: PostHog Code
Task-Id: 855ffa3f-1804-44a4-8ec4-63cd2bc8b018

Author: tatoalo

packages/agent/src/server/agent-server.test.ts

@@ -3292,6 +3292,7 @@ describe("AgentServer pending user attachments", () => {
   });
 
   afterEach(async () => {
+    vi.useRealTimers();
     await server?.stop();
     server = undefined;
     await rm(tempDir, { recursive: true, force: true });
@@ -3313,6 +3314,7 @@ describe("AgentServer pending user attachments", () => {
   };
 
   it("appends an explicit notice when a pending attachment never reaches the manifest", async () => {
+    vi.useFakeTimers();
     const internals = buildInternals();
     // Refetch still can't see the attachment (truly absent, not just lagging).
     const getTaskRun = vi.fn(async () =>
@@ -3323,17 +3325,18 @@ describe("AgentServer pending user attachments", () => {
     );
     internals.posthogAPI.getTaskRun = getTaskRun;
 
-    const result = await internals.getPendingUserPrompt(
+    const resultPromise = internals.getPendingUserPrompt(
       createTaskRun({
         state: { pending_user_artifact_ids: ["missing-attachment"] },
         artifacts: [],
       }),
     );
+    await vi.runAllTimersAsync();
+    const result = await resultPromise;
 
-    // Refetched once to recover a lagging manifest, then — still missing —
-    // surfaced an explicit notice instead of returning null (which would let the
-    // caller fall back to the misleading "Attached files: …" description).
-    expect(getTaskRun).toHaveBeenCalledTimes(1);
+    // Retried to recover a lagging manifest, then surfaced an explicit notice
+    // instead of falling back to the misleading attachment summary.
+    expect(getTaskRun).toHaveBeenCalledTimes(4);
     expect(result).not.toBeNull();
     expect(result?.prompt).toHaveLength(1);
     const [block] = result?.prompt ?? [];
@@ -3387,6 +3390,117 @@ describe("AgentServer pending user attachments", () => {
     expect(hasNotice).toBe(false);
   });
 
+  it("recovers a pending attachment that only lands in a later manifest refetch", async () => {
+    vi.useFakeTimers();
+    const internals = buildInternals();
+    internals.posthogAPI.getTaskRun = vi
+      .fn()
+      .mockResolvedValueOnce(
+        createTaskRun({
+          state: { pending_user_artifact_ids: ["att-1"] },
+          artifacts: [],
+        }),
+      )
+      .mockResolvedValueOnce(
+        createTaskRun({
+          state: { pending_user_artifact_ids: ["att-1"] },
+          artifacts: [],
+        }),
+      )
+      .mockResolvedValue(
+        createTaskRun({
+          state: { pending_user_artifact_ids: ["att-1"] },
+          artifacts: [
+            {
+              id: "att-1",
+              name: "pasted-text.txt",
+              type: "user_attachment",
+              storage_path: "tasks/artifacts/pasted-text.txt",
+              content_type: "text/plain",
+            },
+          ],
+        }),
+      );
+    internals.posthogAPI.downloadArtifact = vi.fn(async () =>
+      exactArrayBuffer(new TextEncoder().encode("pasted body")),
+    );
+
+    const resultPromise = internals.getPendingUserPrompt(
+      createTaskRun({
+        state: { pending_user_artifact_ids: ["att-1"] },
+        artifacts: [],
+      }),
+    );
+    await vi.runAllTimersAsync();
+    const result = await resultPromise;
+
+    expect(internals.posthogAPI.getTaskRun).toHaveBeenCalledTimes(3);
+    expect(result?.prompt.some((block) => block.type === "resource_link")).toBe(
+      true,
+    );
+    expect(
+      result?.prompt.some(
+        (block) =>
+          block.type === "text" && block.text.includes("could not be loaded"),
+      ),
+    ).toBe(false);
+  });
+
+  it("preserves an initially visible attachment while polling for another", async () => {
+    vi.useFakeTimers();
+    const internals = buildInternals();
+    const firstArtifact = {
+      id: "att-1",
+      name: "first.txt",
+      type: "user_attachment" as const,
+      storage_path: "tasks/artifacts/first.txt",
+      content_type: "text/plain",
+    };
+    const secondArtifact = {
+      id: "att-2",
+      name: "second.txt",
+      type: "user_attachment" as const,
+      storage_path: "tasks/artifacts/second.txt",
+      content_type: "text/plain",
+    };
+    internals.posthogAPI.getTaskRun = vi
+      .fn()
+      .mockResolvedValueOnce(
+        createTaskRun({
+          state: { pending_user_artifact_ids: ["att-1", "att-2"] },
+          artifacts: [],
+        }),
+      )
+      .mockResolvedValue(
+        createTaskRun({
+          state: { pending_user_artifact_ids: ["att-1", "att-2"] },
+          artifacts: [secondArtifact],
+        }),
+      );
+    internals.posthogAPI.downloadArtifact = vi.fn(async () =>
+      exactArrayBuffer(new TextEncoder().encode("body")),
+    );
+
+    const resultPromise = internals.getPendingUserPrompt(
+      createTaskRun({
+        state: { pending_user_artifact_ids: ["att-1", "att-2"] },
+        artifacts: [firstArtifact],
+      }),
+    );
+    await vi.runAllTimersAsync();
+    const result = await resultPromise;
+
+    expect(
+      result?.prompt.filter((block) => block.type === "resource_link"),
+    ).toHaveLength(2);
+    expect(
+      result?.prompt.some(
+        (block) =>
+          block.type === "text" && block.text.includes("could not be loaded"),
+      ),
+    ).toBe(false);
+  });
+
   it("returns null without refetching when no pending artifacts were declared", async () => {
     const internals = buildInternals();
     const getTaskRun = vi.fn();
@@ -3401,6 +3515,7 @@ describe("AgentServer pending user attachments", () => {
   });
 
   it("warns once (not twice) about a missing artifact across the speculative and post-refetch resolves", async () => {
+    vi.useFakeTimers();
     const internals = buildInternals();
     // A non-empty manifest that never lists the requested id — so getArtifactsById
     // reaches its per-id "missing" warning on both the pre- and post-refetch calls
@@ -3425,12 +3540,14 @@ describe("AgentServer pending user attachments", () => {
       .spyOn(loggerHost.logger, "warn")
       .mockImplementation(() => {});
 
-    await internals.getPendingUserPrompt(
+    const resultPromise = internals.getPendingUserPrompt(
       createTaskRun({
         state: { pending_user_artifact_ids: ["missing-attachment"] },
         artifacts: decoyManifest,
       }),
     );
+    await vi.runAllTimersAsync();
+    await resultPromise;
 
     // The speculative pre-refetch resolve stays quiet (a miss there is expected);
     // only the post-refetch resolve emits the per-id "missing" warning.

packages/agent/src/server/agent-server.ts

@@ -131,6 +131,12 @@ export const SSE_KEEPALIVE_INTERVAL_MS = 25_000;
 // cut once, without letting a hard upstream outage loop forever.
 const MAX_UPSTREAM_TURN_RETRIES = 2;
 const UPSTREAM_TURN_RETRY_DELAY_MS = 5_000;
+const PENDING_ARTIFACT_MAX_ATTEMPTS = 4;
+const PENDING_ARTIFACT_RETRY_DELAY_MS = 500;
+
+function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
 
 class NdJsonTap {
   private decoder = new TextDecoder();
@@ -2061,9 +2067,10 @@ export class AgentServer {
 
     // The run's artifact manifest can momentarily lag the pending-artifact ids
     // when a run starts right after the attachments were uploaded. If we were
-    // asked for artifacts the manifest doesn't list yet, refetch the run once so
-    // a transient gap doesn't drop the attachment and send the agent the bare
-    // "Attached files: …" description instead of the file it was promised.
+    // asked for artifacts the manifest doesn't list yet, poll the run with a
+    // short backoff so a transient gap doesn't drop the attachment and send the
+    // agent the bare "Attached files: …" description instead of the file it was
+    // promised.
     let manifest = taskRun.artifacts ?? [];
     let resolvedArtifacts = this.getArtifactsById(manifest, artifactIds, {
       warnOnMissing: false,
@@ -2072,11 +2079,13 @@ export class AgentServer {
       artifactIds.length > 0 &&
       resolvedArtifacts.length < artifactIds.length
     ) {
-      const refreshed = await this.refetchRunArtifacts(taskRun);
-      if (refreshed) {
-        manifest = refreshed;
-        resolvedArtifacts = this.getArtifactsById(manifest, artifactIds);
-      }
+      manifest =
+        (await this.resolvePendingArtifactManifest(
+          taskRun,
+          artifactIds,
+          manifest,
+        )) ?? manifest;
+      resolvedArtifacts = this.getArtifactsById(manifest, artifactIds);
     }
 
     const prompt = await this.buildPromptFromContentAndArtifacts({
@@ -2126,6 +2135,49 @@ export class AgentServer {
     return prompt.prompt.length > 0 ? prompt : null;
   }
 
+  private async resolvePendingArtifactManifest(
+    taskRun: TaskRun,
+    artifactIds: string[],
+    initialManifest: TaskRunArtifact[],
+  ): Promise<TaskRunArtifact[] | null> {
+    let latestManifest = initialManifest;
+
+    for (let attempt = 1; attempt <= PENDING_ARTIFACT_MAX_ATTEMPTS; attempt++) {
+      const refreshed = await this.refetchRunArtifacts(taskRun);
+      if (refreshed) {
+        const mergedManifest = [...latestManifest];
+        for (const artifact of refreshed) {
+          const existingIndex = mergedManifest.findIndex(
+            (existing) =>
+              (artifact.id && existing.id === artifact.id) ||
+              (artifact.storage_path &&
+                existing.storage_path === artifact.storage_path),
+          );
+          if (existingIndex >= 0) {
+            mergedManifest[existingIndex] = artifact;
+          } else {
+            mergedManifest.push(artifact);
+          }
+        }
+        latestManifest = mergedManifest;
+        const resolvedArtifacts = this.getArtifactsById(
+          latestManifest,
+          artifactIds,
+          { warnOnMissing: false },
+        );
+        if (resolvedArtifacts.length === artifactIds.length) {
+          return latestManifest;
+        }
+      }
+
+      if (attempt < PENDING_ARTIFACT_MAX_ATTEMPTS) {
+        await sleep(PENDING_ARTIFACT_RETRY_DELAY_MS * attempt);
+      }
+    }
+
+    return latestManifest.length > 0 ? latestManifest : null;
+  }
+
   // Best-effort refetch of a run's artifact manifest. Returns null on any error
   // so the caller can fall back to the manifest it already has.
   private async refetchRunArtifacts(