diff --git a/src/features/workspaces/ai/ai-thread-runtime.ts b/src/features/workspaces/ai/ai-thread-runtime.ts index adf78306..b0134c73 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 00000000..cfc0c0b8 --- /dev/null +++ b/src/features/workspaces/ai/workspace-read-file-fallback.ts @@ -0,0 +1,74 @@ +import type { Tool } from "ai"; + +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"; + +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; + workspaceId: string; +}): Promise { + const parsedOutput = workspaceReadItemsOutputSchema.safeParse(input.toolOutput); + if (!parsedOutput.success) { + return null; + } + + const [result] = parsedOutput.data.results; + if ( + parsedOutput.data.results.length !== 1 || + !result || + result.status !== "pending" || + !result.itemId + ) { + return null; + } + + try { + const kernel = await getWorkspaceKernelFromEnv(input.env, input.workspaceId); + const { object, source } = await getWorkspaceFileSourceObject({ + env: input.env, + itemId: result.itemId, + kernel, + }); + if (source.contentType !== "application/pdf" || object.size > maxPendingPdfBytes) { + return null; + } + + return { + type: "content", + value: [ + { + type: "text", + 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", + 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 00000000..8003689b --- /dev/null +++ b/src/features/workspaces/ai/workspace-read-file-fallback.worker.test.ts @@ -0,0 +1,109 @@ +import { describe, expect, it } from "vitest"; + +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", () => { + it("attaches only a small pending PDF", async () => { + const bytes = new Uint8Array([1, 2, 3]); + const output = await createPendingPdfModelOutput({ + env: createEnv(bytes), + 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", + }, + ], + }); + const text = output?.type === "content" ? output.value[0] : null; + expect(text).toMatchObject({ + text: expect.stringContaining('"retryAfterSeconds":15'), + }); + expect(text).toMatchObject({ + text: expect.stringContaining("The original PDF is attached temporarily"), + }); + + const largeBytes = new Uint8Array(3.5 * 1024 * 1024 + 1); + await expect( + createPendingPdfModelOutput({ + env: createEnv(largeBytes), + toolOutput: pendingOutput(), + workspaceId: "workspace-1", + }), + ).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() { + return { + references: [], + results: [ + { + elapsedSeconds: 8, + itemId: "file-1", + path: "/Report.pdf", + phase: "extracting", + retryAfterSeconds: 15, + status: "pending", + type: "file", + }, + ], + }; +} + +function createEnv(bytes: Uint8Array): Cloudflare.Env { + const kernel = { + getFileSource: async () => ({ + contentType: "application/pdf", + fileName: "Report.pdf", + objectKey: "sources/file-1.pdf", + sizeBytes: bytes.byteLength, + }), + } as unknown as WorkspaceKernelClient; + + return { + WORKSPACE_KERNEL_FILES: { + get: async () => ({ + arrayBuffer: async () => bytes.buffer, + size: bytes.byteLength, + }), + }, + 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 5fc774ea..0ad155bf 100644 --- a/src/features/workspaces/ai/workspace-tools.ts +++ b/src/features/workspaces/ai/workspace-tools.ts @@ -2,6 +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, + 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 { @@ -16,6 +20,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 +29,9 @@ type WorkspaceThreadToolConfig = { function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { const { definition } = input; const resultAdapter = getWorkspaceToolResultAdapter(definition.name); + // 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, @@ -32,10 +40,24 @@ function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { outputSchema: definition.outputSchema, ...(resultAdapter ? { - toModelOutput: ({ output }) => ({ - type: "json" as const, - value: resultAdapter.projectOutput(output), - }), + toModelOutput: async ({ output, toolCallId }) => { + const workspaceId = freshWorkspaceIds?.get(toolCallId); + if (workspaceId) { + const fileOutput = await createPendingPdfModelOutput({ + env: input.env, + toolOutput: output, + workspaceId, + }); + if (fileOutput) { + return fileOutput; + } + } + + return { + type: "json" as const, + value: resultAdapter.projectOutput(output), + }; + }, } : {}), execute: async (args, context) => { @@ -55,6 +77,9 @@ function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { if (references.length > 0 && input.onWorkspaceReferences) { input.onWorkspaceReferences(references); } + if (freshWorkspaceIds && isPendingPdfFallbackInput(args)) { + freshWorkspaceIds.set(context.invocationId, thread.workspaceId); + } return output; }, @@ -62,6 +87,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 +97,7 @@ export function createAIThreadWorkspaceTools(input: { definition.name, createWorkspaceThreadTool({ definition, + env: input.env, getThreadContext: input.getThreadContext, onWorkspaceReferences: input.onWorkspaceReferences, resolveWorkspaceReferences: input.resolveWorkspaceReferences, diff --git a/src/features/workspaces/content/workspace-content-contract.ts b/src/features/workspaces/content/workspace-content-contract.ts index 08860029..c3b8f868 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 3d5050e9..f368f49f 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 efe0a846..f1493d25 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 8ce627d0..5950fa6a 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;