From e0b29de43f5dcc8b7d58c8ade28630e7a399ed26 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Sun, 9 Aug 2026 18:16:59 -0400 Subject: [PATCH 1/3] feat(ai): attach pending PDFs to model --- .../workspaces/ai/ai-thread-runtime.ts | 1 + .../ai/workspace-read-file-fallback.ts | 80 ++++++++++++ ...orkspace-read-file-fallback.worker.test.ts | 118 ++++++++++++++++++ src/features/workspaces/ai/workspace-tools.ts | 31 ++++- 4 files changed, 226 insertions(+), 4 deletions(-) create mode 100644 src/features/workspaces/ai/workspace-read-file-fallback.ts create mode 100644 src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts diff --git a/src/features/workspaces/ai/ai-thread-runtime.ts b/src/features/workspaces/ai/ai-thread-runtime.ts index adf78306d..b0134c73c 100644 --- a/src/features/workspaces/ai/ai-thread-runtime.ts +++ b/src/features/workspaces/ai/ai-thread-runtime.ts @@ -158,6 +158,7 @@ function createAIThreadToolCatalog(input: { defaultTimeZone: input.timeZone, }); const workspaceTools = createAIThreadWorkspaceTools({ + env: input.env, getThreadContext: input.getThreadContext, onWorkspaceReferences: input.onWorkspaceReferences, resolveWorkspaceReferences: input.resolveWorkspaceReferences, diff --git a/src/features/workspaces/ai/workspace-read-file-fallback.ts b/src/features/workspaces/ai/workspace-read-file-fallback.ts new file mode 100644 index 000000000..6e3eda6a3 --- /dev/null +++ b/src/features/workspaces/ai/workspace-read-file-fallback.ts @@ -0,0 +1,80 @@ +import type { Tool } from "ai"; + +import { + workspaceReadItemsInputSchema, + workspaceReadItemsOutputSchema, +} from "#/features/workspaces/content/workspace-content-contract"; +import { getWorkspaceFileSourceObject } from "#/features/workspaces/extraction/workspace-file-source"; +import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; +import { resolveWorkspaceFileTypeFromItem } from "#/features/workspaces/model/workspace-file"; + +const maxPendingPdfBytes = 3.5 * 1024 * 1024; +type ModelToolOutput = Awaited>>; + +export async function createPendingPdfModelOutput(input: { + env: Cloudflare.Env; + toolInput: unknown; + toolOutput: unknown; + workspaceId: string; +}): Promise { + const parsedInput = workspaceReadItemsInputSchema.safeParse(input.toolInput); + const parsedOutput = workspaceReadItemsOutputSchema.safeParse(input.toolOutput); + if (!parsedInput.success || !parsedOutput.success) { + return null; + } + + const [request] = parsedInput.data.requests; + const [result] = parsedOutput.data.results; + if ( + parsedInput.data.requests.length !== 1 || + parsedOutput.data.results.length !== 1 || + !request || + !result || + result.status !== "pending" || + result.path !== request.path + ) { + return null; + } + + try { + const kernel = await getWorkspaceKernelFromEnv(input.env, input.workspaceId); + const [resolution] = await kernel.resolvePaths({ paths: [result.path] }); + if ( + resolution?.status !== "item" || + resolveWorkspaceFileTypeFromItem(resolution.item)?.assetKind !== "pdf" + ) { + return null; + } + + const { object, source } = await getWorkspaceFileSourceObject({ + env: input.env, + itemId: resolution.item.id, + kernel, + }); + if (source.contentType !== "application/pdf" || object.size > maxPendingPdfBytes) { + return null; + } + + return { + type: "content", + value: [ + { + type: "text", + text: `The original PDF at ${result.path} is attached temporarily because its indexed extraction is still running. Use it for this response. Do not claim extraction is complete or invent ThinkEx page citations. Do not read this path again in this response. On a later turn, call workspace_read_items again so extracted content and citations are used when ready. If the raw PDF is not enough, tell the user extraction is still running and ask them to try again in about ${result.retryAfterSeconds} seconds.`, + }, + { + type: "file", + data: { + data: new Uint8Array(await object.arrayBuffer()), + type: "data", + }, + filename: source.fileName, + mediaType: "application/pdf", + }, + ], + }; + } catch { + // This is an optional bridge while extraction runs; preserve the normal pending result on failure. + return null; + } +} diff --git a/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts b/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts new file mode 100644 index 000000000..e731e657d --- /dev/null +++ b/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts @@ -0,0 +1,118 @@ +import { describe, expect, it, vi } from "vitest"; + +import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; +import { createPendingPdfModelOutput } from "#/features/workspaces/ai/workspace-read-file-fallback"; +import type { WorkspaceKernelClient } from "#/features/workspaces/kernel/workspace-kernel-access"; + +const fileItem: WorkspaceItemSummary = { + color: null, + createdAt: "2026-01-01T00:00:00.000Z", + deletedAt: null, + id: "file-1", + meta: "PDF", + metadataJson: { assetKind: "pdf" }, + name: "Report.pdf", + parentId: null, + sortOrder: 1, + title: "Report", + type: "file", + updatedAt: "2026-01-01T00:00:00.000Z", + workspaceId: "workspace-1", +}; + +describe("pending PDF model output", () => { + it("attaches one small pending PDF with extraction guidance", async () => { + const bytes = new Uint8Array([1, 2, 3]); + const output = await createPendingPdfModelOutput({ + env: createEnv(bytes), + toolInput: { requests: [{ mode: "start", path: "/Report.pdf" }] }, + toolOutput: pendingOutput(), + workspaceId: "workspace-1", + }); + + expect(output).toMatchObject({ + type: "content", + value: [ + { type: "text" }, + { + data: { data: bytes, type: "data" }, + filename: "Report.pdf", + mediaType: "application/pdf", + type: "file", + }, + ], + }); + expect(output?.type === "content" ? output.value[0] : null).toMatchObject({ + text: expect.stringContaining("indexed extraction is still running"), + }); + expect(output?.type === "content" ? output.value[0] : null).toMatchObject({ + text: expect.stringContaining("about 15 seconds"), + }); + }); + + it("keeps the normal pending result for large or batched reads", async () => { + const largeBytes = new Uint8Array(3.5 * 1024 * 1024 + 1); + await expect( + createPendingPdfModelOutput({ + env: createEnv(largeBytes), + toolInput: { requests: [{ mode: "start", path: "/Report.pdf" }] }, + toolOutput: pendingOutput(), + workspaceId: "workspace-1", + }), + ).resolves.toBeNull(); + + await expect( + createPendingPdfModelOutput({ + env: createEnv(new Uint8Array([1, 2, 3])), + toolInput: { + requests: [ + { mode: "start", path: "/Report.pdf" }, + { mode: "start", path: "/Other.pdf" }, + ], + }, + toolOutput: pendingOutput(), + workspaceId: "workspace-1", + }), + ).resolves.toBeNull(); + }); +}); + +function pendingOutput() { + return { + references: [], + results: [ + { + elapsedSeconds: 8, + path: "/Report.pdf", + phase: "extracting", + retryAfterSeconds: 15, + status: "pending", + type: "file", + }, + ], + }; +} + +function createEnv(bytes: Uint8Array): Cloudflare.Env { + const kernel = { + getFileSource: vi.fn(async () => ({ + contentType: "application/pdf", + fileName: "Report.pdf", + objectKey: "sources/file-1.pdf", + sizeBytes: bytes.byteLength, + })), + resolvePaths: vi.fn(async () => [ + { item: fileItem, path: "/Report.pdf", status: "item" as const }, + ]), + } as unknown as WorkspaceKernelClient; + + return { + WORKSPACE_KERNEL_FILES: { + get: vi.fn(async () => ({ + arrayBuffer: vi.fn(async () => bytes.buffer), + size: bytes.byteLength, + })), + }, + WorkspaceKernel: { getByName: vi.fn(() => kernel) }, + } as unknown as Cloudflare.Env; +} diff --git a/src/features/workspaces/ai/workspace-tools.ts b/src/features/workspaces/ai/workspace-tools.ts index 5fc774eaa..2b360f87c 100644 --- a/src/features/workspaces/ai/workspace-tools.ts +++ b/src/features/workspaces/ai/workspace-tools.ts @@ -2,6 +2,7 @@ import type { ToolSet } from "ai"; import type { AIThreadContext } from "#/features/workspaces/ai/ai-thread-metadata"; import { defineAIThreadTool } from "#/features/workspaces/ai/ai-thread-tool"; +import { createPendingPdfModelOutput } from "#/features/workspaces/ai/workspace-read-file-fallback"; import { getWorkspaceToolResultAdapter } from "#/features/workspaces/ai/workspace-tool-result-adapters"; import type { WorkspaceReferenceRecord } from "#/features/workspaces/locations/workspace-location"; import { @@ -16,6 +17,7 @@ import { type WorkspaceThreadToolConfig = { definition: (typeof workspaceToolDefinitions)[number]; + env: Cloudflare.Env; getThreadContext: () => Promise; onWorkspaceReferences?: (records: readonly WorkspaceReferenceRecord[]) => void; resolveWorkspaceReferences?: (refs: readonly string[]) => Promise; @@ -24,6 +26,7 @@ type WorkspaceThreadToolConfig = { function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { const { definition } = input; const resultAdapter = getWorkspaceToolResultAdapter(definition.name); + const freshWorkspaceIds = new Map(); return defineAIThreadTool({ description: definition.description, @@ -32,10 +35,25 @@ function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { outputSchema: definition.outputSchema, ...(resultAdapter ? { - toModelOutput: ({ output }) => ({ - type: "json" as const, - value: resultAdapter.projectOutput(output), - }), + toModelOutput: async ({ input: toolInput, output, toolCallId }) => { + const workspaceId = freshWorkspaceIds.get(toolCallId); + if (definition.name === "workspace_read_items" && workspaceId) { + const fileOutput = await createPendingPdfModelOutput({ + env: input.env, + toolInput, + toolOutput: output, + workspaceId, + }); + if (fileOutput) { + return fileOutput; + } + } + + return { + type: "json" as const, + value: resultAdapter.projectOutput(output), + }; + }, } : {}), execute: async (args, context) => { @@ -55,6 +73,9 @@ function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { if (references.length > 0 && input.onWorkspaceReferences) { input.onWorkspaceReferences(references); } + if (definition.name === "workspace_read_items") { + freshWorkspaceIds.set(context.invocationId, thread.workspaceId); + } return output; }, @@ -62,6 +83,7 @@ function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { } export function createAIThreadWorkspaceTools(input: { + env: Cloudflare.Env; getThreadContext: () => Promise; onWorkspaceReferences?: (records: readonly WorkspaceReferenceRecord[]) => void; resolveWorkspaceReferences?: (refs: readonly string[]) => Promise; @@ -71,6 +93,7 @@ export function createAIThreadWorkspaceTools(input: { definition.name, createWorkspaceThreadTool({ definition, + env: input.env, getThreadContext: input.getThreadContext, onWorkspaceReferences: input.onWorkspaceReferences, resolveWorkspaceReferences: input.resolveWorkspaceReferences, From dedb1aa347c8f159cf3b0f44a6ff73bfeeaf2c7b Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Sun, 9 Aug 2026 18:25:42 -0400 Subject: [PATCH 2/3] refactor(ai): simplify pending PDF fallback --- .../ai/workspace-read-file-fallback.ts | 28 ++------ ...orkspace-read-file-fallback.worker.test.ts | 64 ++++--------------- src/features/workspaces/ai/workspace-tools.ts | 15 ++--- .../content/workspace-content-contract.ts | 1 + .../content/workspace-content-reader.ts | 4 +- .../content/workspace-read-references.test.ts | 5 +- .../content/workspace-read-references.ts | 10 ++- 7 files changed, 42 insertions(+), 85 deletions(-) diff --git a/src/features/workspaces/ai/workspace-read-file-fallback.ts b/src/features/workspaces/ai/workspace-read-file-fallback.ts index 6e3eda6a3..41a5716d5 100644 --- a/src/features/workspaces/ai/workspace-read-file-fallback.ts +++ b/src/features/workspaces/ai/workspace-read-file-fallback.ts @@ -1,54 +1,38 @@ import type { Tool } from "ai"; -import { - workspaceReadItemsInputSchema, - workspaceReadItemsOutputSchema, -} from "#/features/workspaces/content/workspace-content-contract"; +import { workspaceReadItemsOutputSchema } from "#/features/workspaces/content/workspace-content-contract"; +import { createWorkspaceReadItemsModelOutput } from "#/features/workspaces/content/workspace-read-references"; import { getWorkspaceFileSourceObject } from "#/features/workspaces/extraction/workspace-file-source"; import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; -import { resolveWorkspaceFileTypeFromItem } from "#/features/workspaces/model/workspace-file"; const maxPendingPdfBytes = 3.5 * 1024 * 1024; type ModelToolOutput = Awaited>>; export async function createPendingPdfModelOutput(input: { env: Cloudflare.Env; - toolInput: unknown; toolOutput: unknown; workspaceId: string; }): Promise { - const parsedInput = workspaceReadItemsInputSchema.safeParse(input.toolInput); const parsedOutput = workspaceReadItemsOutputSchema.safeParse(input.toolOutput); - if (!parsedInput.success || !parsedOutput.success) { + if (!parsedOutput.success) { return null; } - const [request] = parsedInput.data.requests; const [result] = parsedOutput.data.results; if ( - parsedInput.data.requests.length !== 1 || parsedOutput.data.results.length !== 1 || - !request || !result || result.status !== "pending" || - result.path !== request.path + !result.itemId ) { return null; } try { const kernel = await getWorkspaceKernelFromEnv(input.env, input.workspaceId); - const [resolution] = await kernel.resolvePaths({ paths: [result.path] }); - if ( - resolution?.status !== "item" || - resolveWorkspaceFileTypeFromItem(resolution.item)?.assetKind !== "pdf" - ) { - return null; - } - const { object, source } = await getWorkspaceFileSourceObject({ env: input.env, - itemId: resolution.item.id, + itemId: result.itemId, kernel, }); if (source.contentType !== "application/pdf" || object.size > maxPendingPdfBytes) { @@ -60,7 +44,7 @@ export async function createPendingPdfModelOutput(input: { value: [ { type: "text", - text: `The original PDF at ${result.path} is attached temporarily because its indexed extraction is still running. Use it for this response. Do not claim extraction is complete or invent ThinkEx page citations. Do not read this path again in this response. On a later turn, call workspace_read_items again so extracted content and citations are used when ready. If the raw PDF is not enough, tell the user extraction is still running and ask them to try again in about ${result.retryAfterSeconds} seconds.`, + text: `${JSON.stringify(createWorkspaceReadItemsModelOutput(parsedOutput.data))}\n\nThe original PDF is attached temporarily for this response. Do not claim extraction is complete or invent ThinkEx page citations. On a later turn, call workspace_read_items again for extracted, citable content.`, }, { type: "file", diff --git a/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts b/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts index e731e657d..5a9bfff7f 100644 --- a/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts +++ b/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts @@ -1,31 +1,13 @@ -import { describe, expect, it, vi } from "vitest"; +import { describe, expect, it } from "vitest"; -import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; import { createPendingPdfModelOutput } from "#/features/workspaces/ai/workspace-read-file-fallback"; import type { WorkspaceKernelClient } from "#/features/workspaces/kernel/workspace-kernel-access"; -const fileItem: WorkspaceItemSummary = { - color: null, - createdAt: "2026-01-01T00:00:00.000Z", - deletedAt: null, - id: "file-1", - meta: "PDF", - metadataJson: { assetKind: "pdf" }, - name: "Report.pdf", - parentId: null, - sortOrder: 1, - title: "Report", - type: "file", - updatedAt: "2026-01-01T00:00:00.000Z", - workspaceId: "workspace-1", -}; - describe("pending PDF model output", () => { - it("attaches one small pending PDF with extraction guidance", async () => { + it("attaches only a small pending PDF", async () => { const bytes = new Uint8Array([1, 2, 3]); const output = await createPendingPdfModelOutput({ env: createEnv(bytes), - toolInput: { requests: [{ mode: "start", path: "/Report.pdf" }] }, toolOutput: pendingOutput(), workspaceId: "workspace-1", }); @@ -42,34 +24,18 @@ describe("pending PDF model output", () => { }, ], }); - expect(output?.type === "content" ? output.value[0] : null).toMatchObject({ - text: expect.stringContaining("indexed extraction is still running"), + const text = output?.type === "content" ? output.value[0] : null; + expect(text).toMatchObject({ + text: expect.stringContaining('"retryAfterSeconds":15'), }); - expect(output?.type === "content" ? output.value[0] : null).toMatchObject({ - text: expect.stringContaining("about 15 seconds"), + expect(text).toMatchObject({ + text: expect.stringContaining("The original PDF is attached temporarily"), }); - }); - it("keeps the normal pending result for large or batched reads", async () => { const largeBytes = new Uint8Array(3.5 * 1024 * 1024 + 1); await expect( createPendingPdfModelOutput({ env: createEnv(largeBytes), - toolInput: { requests: [{ mode: "start", path: "/Report.pdf" }] }, - toolOutput: pendingOutput(), - workspaceId: "workspace-1", - }), - ).resolves.toBeNull(); - - await expect( - createPendingPdfModelOutput({ - env: createEnv(new Uint8Array([1, 2, 3])), - toolInput: { - requests: [ - { mode: "start", path: "/Report.pdf" }, - { mode: "start", path: "/Other.pdf" }, - ], - }, toolOutput: pendingOutput(), workspaceId: "workspace-1", }), @@ -83,6 +49,7 @@ function pendingOutput() { results: [ { elapsedSeconds: 8, + itemId: "file-1", path: "/Report.pdf", phase: "extracting", retryAfterSeconds: 15, @@ -95,24 +62,21 @@ function pendingOutput() { function createEnv(bytes: Uint8Array): Cloudflare.Env { const kernel = { - getFileSource: vi.fn(async () => ({ + getFileSource: async () => ({ contentType: "application/pdf", fileName: "Report.pdf", objectKey: "sources/file-1.pdf", sizeBytes: bytes.byteLength, - })), - resolvePaths: vi.fn(async () => [ - { item: fileItem, path: "/Report.pdf", status: "item" as const }, - ]), + }), } as unknown as WorkspaceKernelClient; return { WORKSPACE_KERNEL_FILES: { - get: vi.fn(async () => ({ - arrayBuffer: vi.fn(async () => bytes.buffer), + get: async () => ({ + arrayBuffer: async () => bytes.buffer, size: bytes.byteLength, - })), + }), }, - WorkspaceKernel: { getByName: vi.fn(() => kernel) }, + WorkspaceKernel: { getByName: () => kernel }, } as unknown as Cloudflare.Env; } diff --git a/src/features/workspaces/ai/workspace-tools.ts b/src/features/workspaces/ai/workspace-tools.ts index 2b360f87c..8905fcb0f 100644 --- a/src/features/workspaces/ai/workspace-tools.ts +++ b/src/features/workspaces/ai/workspace-tools.ts @@ -26,7 +26,9 @@ type WorkspaceThreadToolConfig = { function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { const { definition } = input; const resultAdapter = getWorkspaceToolResultAdapter(definition.name); - const freshWorkspaceIds = new Map(); + // Projection also runs for history; only fresh calls may hydrate raw bytes. + const freshWorkspaceIds = + definition.name === "workspace_read_items" ? new Map() : null; return defineAIThreadTool({ description: definition.description, @@ -35,12 +37,11 @@ function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { outputSchema: definition.outputSchema, ...(resultAdapter ? { - toModelOutput: async ({ input: toolInput, output, toolCallId }) => { - const workspaceId = freshWorkspaceIds.get(toolCallId); - if (definition.name === "workspace_read_items" && workspaceId) { + toModelOutput: async ({ output, toolCallId }) => { + const workspaceId = freshWorkspaceIds?.get(toolCallId); + if (workspaceId) { const fileOutput = await createPendingPdfModelOutput({ env: input.env, - toolInput, toolOutput: output, workspaceId, }); @@ -73,9 +74,7 @@ function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { if (references.length > 0 && input.onWorkspaceReferences) { input.onWorkspaceReferences(references); } - if (definition.name === "workspace_read_items") { - freshWorkspaceIds.set(context.invocationId, thread.workspaceId); - } + freshWorkspaceIds?.set(context.invocationId, thread.workspaceId); return output; }, diff --git a/src/features/workspaces/content/workspace-content-contract.ts b/src/features/workspaces/content/workspace-content-contract.ts index 08860029c..c3b8f8686 100644 --- a/src/features/workspaces/content/workspace-content-contract.ts +++ b/src/features/workspaces/content/workspace-content-contract.ts @@ -131,6 +131,7 @@ const workspaceContentReadResultSchema = z.union([ .int() .nonnegative() .describe("How long extraction has been running."), + itemId: z.string().min(1).optional(), path: workspacePathSchema, phase: z .enum(["queued", "extracting"]) diff --git a/src/features/workspaces/content/workspace-content-reader.ts b/src/features/workspaces/content/workspace-content-reader.ts index 3d5050e94..f368f49f7 100644 --- a/src/features/workspaces/content/workspace-content-reader.ts +++ b/src/features/workspaces/content/workspace-content-reader.ts @@ -246,7 +246,7 @@ async function readFile(input: { Date.now(), ); if (projection.state !== "ready") { - return describeUnreadableProjection(projection, input.path); + return describeUnreadableProjection(projection, input.path, input.item.id); } const encodedCursor = input.request.mode === "continue" ? input.request.cursor : undefined; @@ -312,10 +312,12 @@ async function readFile(input: { function describeUnreadableProjection( projection: Exclude, path: string, + itemId: string, ): WorkspaceContentReadResult { if (projection.state === "pending") { return { elapsedSeconds: projection.elapsedSeconds, + itemId, path, phase: projection.phase, retryAfterSeconds: projection.retryAfterSeconds, diff --git a/src/features/workspaces/content/workspace-read-references.test.ts b/src/features/workspaces/content/workspace-read-references.test.ts index efe0a846b..f1493d254 100644 --- a/src/features/workspaces/content/workspace-read-references.test.ts +++ b/src/features/workspaces/content/workspace-read-references.test.ts @@ -101,6 +101,7 @@ describe("workspace read references", () => { const results = [ { elapsedSeconds: 4, + itemId: "file-a", path: "/A.pdf", phase: "extracting", retryAfterSeconds: 15, @@ -116,10 +117,12 @@ describe("workspace read references", () => { type: "file", }, ] satisfies WorkspaceContentReadResult[]; - const guidance = createWorkspaceReadItemsModelOutput({ references: [], results }).guidance; + const modelOutput = createWorkspaceReadItemsModelOutput({ references: [], results }); + const { guidance } = modelOutput; expect(guidance).toHaveLength(1); expect(guidance?.[0]).toContain("Never sleep"); + expect(JSON.stringify(modelOutput)).not.toContain("file-a"); }); it("separates failures that will never resolve from transient ones", () => { diff --git a/src/features/workspaces/content/workspace-read-references.ts b/src/features/workspaces/content/workspace-read-references.ts index 8ce627d0e..5950fa6a3 100644 --- a/src/features/workspaces/content/workspace-read-references.ts +++ b/src/features/workspaces/content/workspace-read-references.ts @@ -60,7 +60,7 @@ export function createWorkspaceReadReferences( */ const workspaceReadGuidance = { pending: - "Some paths are still extracting. Never sleep, poll, or otherwise stall waiting for them, including inside compute, sandbox_bash, or orchestrate. Either do other work and read those paths again later in this reply, or tell the user they are still processing and to ask again in about retryAfterSeconds. Never read the same pending path more than twice in one reply.", + "Some paths are still extracting. Never sleep, poll, or otherwise stall waiting for them, including inside compute, sandbox_bash, or orchestrate. If an original file is attached, use it and do not read that path again in this reply. Otherwise, either do other work and read pending paths again later in this reply, or tell the user they are still processing and to ask again in about retryAfterSeconds. Never read the same pending path more than twice in one reply.", unrecoverable: "Extraction will not finish for some paths. Report the code and any message to the user; do not retry those reads and do not suggest re-uploading the file.", transient: @@ -128,7 +128,11 @@ export function createWorkspaceReadItemsModelOutput(output: WorkspaceReadItemsOu return { ...(guidance.length > 0 ? { guidance } : {}), results: output.results.map((result) => { - if (result.status !== "ready") { + if (result.status === "pending") { + return omitWorkspaceReadItemId(result); + } + + if (result.status === "failed") { return result; } @@ -168,7 +172,7 @@ export function createWorkspaceReadItemsModelOutput(output: WorkspaceReadItemsOu }; } -function omitWorkspaceReadItemId( +function omitWorkspaceReadItemId( result: T, ): Omit { const { itemId: _itemId, ...modelResult } = result; From 618a82e934cf2e94c2e0c1788bee9164b4c5e572 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:03:03 -0400 Subject: [PATCH 3/3] fix(ai): limit pending PDF fallback to initial reads --- .../ai/workspace-read-file-fallback.ts | 12 +++++++- ...orkspace-read-file-fallback.worker.test.ts | 29 ++++++++++++++++++- src/features/workspaces/ai/workspace-tools.ts | 9 ++++-- 3 files changed, 46 insertions(+), 4 deletions(-) diff --git a/src/features/workspaces/ai/workspace-read-file-fallback.ts b/src/features/workspaces/ai/workspace-read-file-fallback.ts index 41a5716d5..cfc0c0b82 100644 --- a/src/features/workspaces/ai/workspace-read-file-fallback.ts +++ b/src/features/workspaces/ai/workspace-read-file-fallback.ts @@ -1,6 +1,9 @@ import type { Tool } from "ai"; -import { workspaceReadItemsOutputSchema } from "#/features/workspaces/content/workspace-content-contract"; +import { + workspaceReadItemsInputSchema, + workspaceReadItemsOutputSchema, +} from "#/features/workspaces/content/workspace-content-contract"; import { createWorkspaceReadItemsModelOutput } from "#/features/workspaces/content/workspace-read-references"; import { getWorkspaceFileSourceObject } from "#/features/workspaces/extraction/workspace-file-source"; import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; @@ -8,6 +11,13 @@ import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspac const maxPendingPdfBytes = 3.5 * 1024 * 1024; type ModelToolOutput = Awaited>>; +export function isPendingPdfFallbackInput(input: unknown) { + const parsed = workspaceReadItemsInputSchema.safeParse(input); + return ( + parsed.success && parsed.data.requests.length === 1 && parsed.data.requests[0]?.mode === "start" + ); +} + export async function createPendingPdfModelOutput(input: { env: Cloudflare.Env; toolOutput: unknown; diff --git a/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts b/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts index 5a9bfff7f..8003689ba 100644 --- a/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts +++ b/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts @@ -1,6 +1,9 @@ import { describe, expect, it } from "vitest"; -import { createPendingPdfModelOutput } from "#/features/workspaces/ai/workspace-read-file-fallback"; +import { + createPendingPdfModelOutput, + isPendingPdfFallbackInput, +} from "#/features/workspaces/ai/workspace-read-file-fallback"; import type { WorkspaceKernelClient } from "#/features/workspaces/kernel/workspace-kernel-access"; describe("pending PDF model output", () => { @@ -41,6 +44,30 @@ describe("pending PDF model output", () => { }), ).resolves.toBeNull(); }); + + it("allows the fallback only for one initial read", () => { + expect(isPendingPdfFallbackInput({ requests: [{ mode: "start", path: "/Report.pdf" }] })).toBe( + true, + ); + expect( + isPendingPdfFallbackInput({ + requests: [{ cursor: "opaque", mode: "continue", path: "/Report.pdf" }], + }), + ).toBe(false); + expect( + isPendingPdfFallbackInput({ + requests: [{ mode: "pages", path: "/Report.pdf", range: "2" }], + }), + ).toBe(false); + expect( + isPendingPdfFallbackInput({ + requests: [ + { mode: "start", path: "/Report.pdf" }, + { mode: "start", path: "/Appendix.pdf" }, + ], + }), + ).toBe(false); + }); }); function pendingOutput() { diff --git a/src/features/workspaces/ai/workspace-tools.ts b/src/features/workspaces/ai/workspace-tools.ts index 8905fcb0f..0ad155bf9 100644 --- a/src/features/workspaces/ai/workspace-tools.ts +++ b/src/features/workspaces/ai/workspace-tools.ts @@ -2,7 +2,10 @@ import type { ToolSet } from "ai"; import type { AIThreadContext } from "#/features/workspaces/ai/ai-thread-metadata"; import { defineAIThreadTool } from "#/features/workspaces/ai/ai-thread-tool"; -import { createPendingPdfModelOutput } from "#/features/workspaces/ai/workspace-read-file-fallback"; +import { + createPendingPdfModelOutput, + isPendingPdfFallbackInput, +} from "#/features/workspaces/ai/workspace-read-file-fallback"; import { getWorkspaceToolResultAdapter } from "#/features/workspaces/ai/workspace-tool-result-adapters"; import type { WorkspaceReferenceRecord } from "#/features/workspaces/locations/workspace-location"; import { @@ -74,7 +77,9 @@ function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { if (references.length > 0 && input.onWorkspaceReferences) { input.onWorkspaceReferences(references); } - freshWorkspaceIds?.set(context.invocationId, thread.workspaceId); + if (freshWorkspaceIds && isPendingPdfFallbackInput(args)) { + freshWorkspaceIds.set(context.invocationId, thread.workspaceId); + } return output; },