diff --git a/apps/code/src/main/trpc/routers/os.ts b/apps/code/src/main/trpc/routers/os.ts index ad9411f21..f56c29a46 100644 --- a/apps/code/src/main/trpc/routers/os.ts +++ b/apps/code/src/main/trpc/routers/os.ts @@ -303,4 +303,27 @@ export const osRouter = router({ return { path: filePath, name: displayName, mimeType }; }), + + /** + * Save arbitrary file bytes to a temp file + * Returns the file path for use as a file attachment + */ + saveClipboardFile: publicProcedure + .input( + z.object({ + base64Data: z.string(), + originalName: z.string().optional(), + }), + ) + .mutation(async ({ input }) => { + const displayName = path.basename(input.originalName ?? "attachment"); + const filePath = await createClipboardTempFilePath(displayName); + + await fsPromises.writeFile( + filePath, + Buffer.from(input.base64Data, "base64"), + ); + + return { path: filePath, name: displayName }; + }), }); diff --git a/apps/code/src/renderer/api/posthogClient.ts b/apps/code/src/renderer/api/posthogClient.ts index e8e1e87f3..6f6e7280b 100644 --- a/apps/code/src/renderer/api/posthogClient.ts +++ b/apps/code/src/renderer/api/posthogClient.ts @@ -98,6 +98,37 @@ export interface ExternalDataSource { schemas?: ExternalDataSourceSchema[] | string; } +export interface TaskArtifactUploadRequest { + name: string; + type: "user_attachment"; + size: number; + content_type?: string; + source?: string; +} + +export interface DirectUploadPresignedPost { + url: string; + fields: Record; +} + +export interface PreparedTaskArtifactUpload extends TaskArtifactUploadRequest { + id: string; + storage_path: string; + expires_in: number; + presigned_post: DirectUploadPresignedPost; +} + +export interface FinalizedTaskArtifactUpload { + id: string; + name: string; + type: string; + source?: string; + size?: number; + content_type?: string; + storage_path: string; + uploaded_at?: string; +} + type CloudRuntimeAdapter = "claude" | "codex"; function isObjectRecord(value: unknown): value is Record { @@ -773,6 +804,7 @@ export class PostHogAPIClient { reasoningLevel?: string; resumeFromRunId?: string; pendingUserMessage?: string; + pendingUserArtifactIds?: string[]; sandboxEnvironmentId?: string; prAuthorshipMode?: PrAuthorshipMode; runSource?: CloudRunSource; @@ -817,6 +849,9 @@ export class PostHogAPIClient { if (options?.pendingUserMessage) { body.pending_user_message = options.pendingUserMessage; } + if (options?.pendingUserArtifactIds?.length) { + body.pending_user_artifact_ids = options.pendingUserArtifactIds; + } if (options?.sandboxEnvironmentId) { body.sandbox_environment_id = options.sandboxEnvironmentId; } @@ -847,6 +882,154 @@ export class PostHogAPIClient { return data as unknown as Task; } + async prepareTaskStagedArtifactUploads( + taskId: string, + artifacts: TaskArtifactUploadRequest[], + ): Promise { + if (!artifacts.length) { + return []; + } + + const teamId = await this.getTeamId(); + const url = new URL( + `${this.api.baseUrl}/api/projects/${teamId}/tasks/${taskId}/staged_artifacts/prepare_upload/`, + ); + const response = await this.api.fetcher.fetch({ + method: "post", + url, + path: `/api/projects/${teamId}/tasks/${taskId}/staged_artifacts/prepare_upload/`, + overrides: { + body: JSON.stringify({ artifacts }), + }, + }); + + if (!response.ok) { + throw new Error( + `Failed to prepare staged uploads: ${response.statusText}`, + ); + } + + const data = (await response.json()) as { + artifacts?: PreparedTaskArtifactUpload[]; + }; + return data.artifacts ?? []; + } + + async finalizeTaskStagedArtifactUploads( + taskId: string, + artifacts: PreparedTaskArtifactUpload[], + ): Promise { + if (!artifacts.length) { + return []; + } + + const teamId = await this.getTeamId(); + const url = new URL( + `${this.api.baseUrl}/api/projects/${teamId}/tasks/${taskId}/staged_artifacts/finalize_upload/`, + ); + const response = await this.api.fetcher.fetch({ + method: "post", + url, + path: `/api/projects/${teamId}/tasks/${taskId}/staged_artifacts/finalize_upload/`, + overrides: { + body: JSON.stringify({ + artifacts: artifacts.map((artifact) => ({ + id: artifact.id, + name: artifact.name, + type: artifact.type, + source: artifact.source, + content_type: artifact.content_type, + storage_path: artifact.storage_path, + })), + }), + }, + }); + + if (!response.ok) { + throw new Error( + `Failed to finalize staged uploads: ${response.statusText}`, + ); + } + + const data = (await response.json()) as { + artifacts?: FinalizedTaskArtifactUpload[]; + }; + return data.artifacts ?? []; + } + + async prepareTaskRunArtifactUploads( + taskId: string, + runId: string, + artifacts: TaskArtifactUploadRequest[], + ): Promise { + if (!artifacts.length) { + return []; + } + + const teamId = await this.getTeamId(); + const url = new URL( + `${this.api.baseUrl}/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/artifacts/prepare_upload/`, + ); + const response = await this.api.fetcher.fetch({ + method: "post", + url, + path: `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/artifacts/prepare_upload/`, + overrides: { + body: JSON.stringify({ artifacts }), + }, + }); + + if (!response.ok) { + throw new Error(`Failed to prepare uploads: ${response.statusText}`); + } + + const data = (await response.json()) as { + artifacts?: PreparedTaskArtifactUpload[]; + }; + return data.artifacts ?? []; + } + + async finalizeTaskRunArtifactUploads( + taskId: string, + runId: string, + artifacts: PreparedTaskArtifactUpload[], + ): Promise { + if (!artifacts.length) { + return []; + } + + const teamId = await this.getTeamId(); + const url = new URL( + `${this.api.baseUrl}/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/artifacts/finalize_upload/`, + ); + const response = await this.api.fetcher.fetch({ + method: "post", + url, + path: `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/artifacts/finalize_upload/`, + overrides: { + body: JSON.stringify({ + artifacts: artifacts.map((artifact) => ({ + id: artifact.id, + name: artifact.name, + type: artifact.type, + source: artifact.source, + content_type: artifact.content_type, + storage_path: artifact.storage_path, + })), + }), + }, + }); + + if (!response.ok) { + throw new Error(`Failed to finalize uploads: ${response.statusText}`); + } + + const data = (await response.json()) as { + artifacts?: FinalizedTaskArtifactUpload[]; + }; + return data.artifacts ?? []; + } + async listTaskRuns(taskId: string): Promise { const teamId = await this.getTeamId(); const url = new URL( diff --git a/apps/code/src/renderer/features/message-editor/utils/persistFile.test.ts b/apps/code/src/renderer/features/message-editor/utils/persistFile.test.ts index 74a121742..8dffe8be0 100644 --- a/apps/code/src/renderer/features/message-editor/utils/persistFile.test.ts +++ b/apps/code/src/renderer/features/message-editor/utils/persistFile.test.ts @@ -2,6 +2,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; const mockSaveClipboardImage = vi.hoisted(() => vi.fn()); const mockSaveClipboardText = vi.hoisted(() => vi.fn()); +const mockSaveClipboardFile = vi.hoisted(() => vi.fn()); vi.mock("@renderer/trpc/client", () => ({ trpcClient: { @@ -12,6 +13,9 @@ vi.mock("@renderer/trpc/client", () => ({ saveClipboardText: { mutate: mockSaveClipboardText, }, + saveClipboardFile: { + mutate: mockSaveClipboardFile, + }, }, }, })); @@ -98,28 +102,47 @@ describe("persistFile", () => { }); }); - it("throws for unsupported file types", async () => { - const file = { name: "archive.zip" } as unknown as File; - await expect(persistBrowserFile(file)).rejects.toThrow(/Unsupported/); + it("persists arbitrary non-image files via saveClipboardFile", async () => { + mockSaveClipboardFile.mockResolvedValue({ + path: "/tmp/posthog-code-clipboard/attachment-def/archive.zip", + name: "archive.zip", + }); + + const file = { + name: "archive.zip", + type: "application/zip", + arrayBuffer: vi.fn().mockResolvedValue(new ArrayBuffer(8)), + } as unknown as File; + + await expect(persistBrowserFile(file)).resolves.toEqual({ + id: "/tmp/posthog-code-clipboard/attachment-def/archive.zip", + label: "archive.zip", + }); + + expect(mockSaveClipboardFile).toHaveBeenCalledWith({ + base64Data: expect.any(String), + originalName: "archive.zip", + }); }); it("returns the preserved filename for browser-selected text files", async () => { - mockSaveClipboardText.mockResolvedValue({ + mockSaveClipboardFile.mockResolvedValue({ path: "/tmp/posthog-code-clipboard/attachment-456/config.json", name: "config.json", }); const file = { name: "config.json", - text: vi.fn().mockResolvedValue('{"ok":true}'), + type: "application/json", + arrayBuffer: vi.fn().mockResolvedValue(new ArrayBuffer(8)), } as unknown as File; await expect(persistBrowserFile(file)).resolves.toEqual({ id: "/tmp/posthog-code-clipboard/attachment-456/config.json", label: "config.json", }); - expect(mockSaveClipboardText).toHaveBeenCalledWith({ - text: '{"ok":true}', + expect(mockSaveClipboardFile).toHaveBeenCalledWith({ + base64Data: expect.any(String), originalName: "config.json", }); }); diff --git a/apps/code/src/renderer/features/message-editor/utils/persistFile.ts b/apps/code/src/renderer/features/message-editor/utils/persistFile.ts index e13ee77a9..c82612a25 100644 --- a/apps/code/src/renderer/features/message-editor/utils/persistFile.ts +++ b/apps/code/src/renderer/features/message-editor/utils/persistFile.ts @@ -1,8 +1,4 @@ import { getImageMimeType } from "@features/code-editor/utils/imageUtils"; -import { - isSupportedCloudImageAttachment, - isSupportedCloudTextAttachment, -} from "@features/editor/utils/cloud-prompt"; import { trpcClient } from "@renderer/trpc/client"; const CHUNK_SIZE = 8192; @@ -46,21 +42,30 @@ export async function persistTextContent( return { path: result.path, name: result.name }; } +export async function persistGenericFile(file: File): Promise { + const arrayBuffer = await file.arrayBuffer(); + const base64Data = arrayBufferToBase64(arrayBuffer); + + const result = await trpcClient.os.saveClipboardFile.mutate({ + base64Data, + originalName: file.name, + }); + + return { + path: result.path, + name: result.name, + mimeType: file.type || undefined, + }; +} + export async function persistBrowserFile( file: File, ): Promise<{ id: string; label: string }> { - if (isSupportedCloudImageAttachment(file.name)) { + if (file.type.startsWith("image/")) { const result = await persistImageFile(file); return { id: result.path, label: file.name }; } - if (isSupportedCloudTextAttachment(file.name)) { - const text = await file.text(); - const result = await persistTextContent(text, file.name); - return { id: result.path, label: result.name }; - } - - throw new Error( - `Unsupported attachment: ${file.name}. Cloud attachments currently support text files and PNG/JPG/GIF/WebP images.`, - ); + const result = await persistGenericFile(file); + return { id: result.path, label: result.name }; } diff --git a/apps/code/src/renderer/features/sessions/hooks/useSessionCallbacks.ts b/apps/code/src/renderer/features/sessions/hooks/useSessionCallbacks.ts index d6d657fbd..3f5431b71 100644 --- a/apps/code/src/renderer/features/sessions/hooks/useSessionCallbacks.ts +++ b/apps/code/src/renderer/features/sessions/hooks/useSessionCallbacks.ts @@ -1,6 +1,5 @@ import { tryExecuteCodeCommand } from "@features/message-editor/commands"; import { useDraftStore } from "@features/message-editor/stores/draftStore"; -import { xmlToContent } from "@features/message-editor/utils/content"; import { useTaskViewed } from "@features/sidebar/hooks/useTaskViewed"; import { trpcClient } from "@renderer/trpc/client"; import type { Task } from "@shared/types"; @@ -11,6 +10,10 @@ import { useCallback, useRef } from "react"; import { getSessionService } from "../service/service"; import type { AgentSession } from "../stores/sessionStore"; import { sessionStoreSetters } from "../stores/sessionStore"; +import { + combineQueuedCloudPrompts, + promptToQueuedEditorContent, +} from "../utils/cloudArtifacts"; const log = logger.scope("session-callbacks"); @@ -73,11 +76,27 @@ export function useSessionCallbacks({ ); const handleCancelPrompt = useCallback(async () => { - const queuedContent = sessionStoreSetters.dequeueMessagesAsText(taskId); - await getSessionService().cancelPrompt(taskId); - - if (queuedContent) { - setPendingContent(taskId, xmlToContent(queuedContent)); + const queuedMessages = sessionStoreSetters.dequeueMessages(taskId); + const result = await getSessionService().cancelPrompt(taskId); + log.info("Prompt cancelled", { success: result }); + + const queuedPrompt = sessionRef.current?.isCloud + ? combineQueuedCloudPrompts(queuedMessages) + : queuedMessages.map((message) => message.content).join("\n\n"); + + if (queuedPrompt) { + const pendingContent = sessionRef.current?.isCloud + ? promptToQueuedEditorContent(queuedPrompt) + : { + segments: [ + { + type: "text" as const, + text: typeof queuedPrompt === "string" ? queuedPrompt : "", + }, + ], + }; + + setPendingContent(taskId, pendingContent); } requestFocus(taskId); }, [taskId, setPendingContent, requestFocus]); diff --git a/apps/code/src/renderer/features/sessions/service/service.test.ts b/apps/code/src/renderer/features/sessions/service/service.test.ts index bd2f64af1..83fbe1c40 100644 --- a/apps/code/src/renderer/features/sessions/service/service.test.ts +++ b/apps/code/src/renderer/features/sessions/service/service.test.ts @@ -39,12 +39,17 @@ const mockTrpcCloudTask = vi.hoisted(() => ({ onUpdate: { subscribe: vi.fn() }, })); +const mockTrpcFs = vi.hoisted(() => ({ + readFileAsBase64: { query: vi.fn() }, +})); + vi.mock("@renderer/trpc/client", () => ({ trpcClient: { agent: mockTrpcAgent, workspace: mockTrpcWorkspace, logs: mockTrpcLogs, cloudTask: mockTrpcCloudTask, + fs: mockTrpcFs, }, })); @@ -57,6 +62,7 @@ const mockSessionStoreSetters = vi.hoisted(() => ({ removeQueuedMessage: vi.fn(), clearMessageQueue: vi.fn(), dequeueMessagesAsText: vi.fn(() => null), + dequeueMessages: vi.fn(() => []), setPendingPermissions: vi.fn(), getSessionByTaskId: vi.fn(), getSessions: vi.fn(() => ({})), @@ -99,6 +105,10 @@ const mockAuthenticatedClient = vi.hoisted(() => ({ getTaskRun: vi.fn(), getTask: vi.fn(), runTaskInCloud: vi.fn(), + prepareTaskRunArtifactUploads: vi.fn(), + finalizeTaskRunArtifactUploads: vi.fn(), + prepareTaskStagedArtifactUploads: vi.fn(), + finalizeTaskStagedArtifactUploads: vi.fn(), })); type MockAuthenticatedClient = typeof mockAuthenticatedClient; @@ -237,6 +247,16 @@ vi.mock("@utils/session", async () => { await vi.importActual("@utils/session"); return { convertStoredEntriesToEvents: vi.fn(() => []), + createUserPromptEvent: vi.fn((prompt, ts) => ({ + type: "acp_message", + ts, + message: { + jsonrpc: "2.0", + id: ts, + method: "session/prompt", + params: { prompt }, + }, + })), createUserMessageEvent: vi.fn((message, ts) => ({ type: "user", ts, @@ -323,6 +343,17 @@ describe("SessionService", () => { mockTrpcCloudTask.onUpdate.subscribe.mockReturnValue({ unsubscribe: vi.fn(), }); + mockTrpcFs.readFileAsBase64.query.mockResolvedValue(null); + mockAuthenticatedClient.prepareTaskRunArtifactUploads.mockResolvedValue([]); + mockAuthenticatedClient.finalizeTaskRunArtifactUploads.mockResolvedValue( + [], + ); + mockAuthenticatedClient.prepareTaskStagedArtifactUploads.mockResolvedValue( + [], + ); + mockAuthenticatedClient.finalizeTaskStagedArtifactUploads.mockResolvedValue( + [], + ); }); describe("singleton management", () => { @@ -1019,6 +1050,36 @@ describe("SessionService", () => { ); }); + it("preserves cloud attachment prompts when queueing a follow-up", async () => { + const service = getSessionService(); + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( + createMockSession({ + isCloud: true, + cloudStatus: "in_progress", + isPromptPending: true, + }), + ); + + const prompt: ContentBlock[] = [ + { type: "text", text: "read this" }, + { + type: "resource_link", + uri: "file:///tmp/test.txt", + name: "test.txt", + mimeType: "text/plain", + }, + ]; + + const result = await service.sendPrompt("task-123", prompt); + + expect(result.stopReason).toBe("queued"); + expect(mockSessionStoreSetters.enqueueMessage).toHaveBeenCalledWith( + "task-123", + "read this\n\nAttached files: test.txt", + prompt, + ); + }); + it("sends prompt via tRPC when session is ready", async () => { const service = getSessionService(); mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( @@ -1035,7 +1096,7 @@ describe("SessionService", () => { }); }); - it("serializes structured prompts before sending cloud follow-ups", async () => { + it("uploads attachments before sending cloud follow-ups", async () => { const service = getSessionService(); mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( createMockSession({ @@ -1047,16 +1108,47 @@ describe("SessionService", () => { success: true, result: { stopReason: "end_turn" }, }); + mockTrpcFs.readFileAsBase64.query.mockResolvedValue("aGVsbG8="); + mockAuthenticatedClient.prepareTaskRunArtifactUploads.mockResolvedValue([ + { + id: "artifact-1", + name: "test.txt", + type: "user_attachment", + source: "posthog_code", + size: 5, + content_type: "text/plain", + storage_path: "tasks/artifacts/test.txt", + expires_in: 3600, + presigned_post: { + url: "https://uploads.example.com", + fields: { key: "tasks/artifacts/test.txt" }, + }, + }, + ]); + mockAuthenticatedClient.finalizeTaskRunArtifactUploads.mockResolvedValue([ + { + id: "artifact-1", + name: "test.txt", + type: "user_attachment", + source: "posthog_code", + size: 5, + content_type: "text/plain", + storage_path: "tasks/artifacts/test.txt", + uploaded_at: "2026-04-16T00:00:00Z", + }, + ]); + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue({ ok: true } as Response), + ); const prompt: ContentBlock[] = [ { type: "text", text: "read this" }, { - type: "resource", - resource: { - uri: "attachment://test.txt", - text: "hello from file", - mimeType: "text/plain", - }, + type: "resource_link", + uri: "file:///tmp/test.txt", + name: "test.txt", + mimeType: "text/plain", }, ]; @@ -1065,17 +1157,13 @@ describe("SessionService", () => { expect(result.stopReason).toBe("end_turn"); expect(mockTrpcCloudTask.sendCommand.mutate).toHaveBeenCalledTimes(1); - const [args] = mockTrpcCloudTask.sendCommand.mutate.mock.calls[0] as [ - { - params?: { content?: unknown }; - }, - ]; - - expect(args.params?.content).toEqual( - expect.stringContaining("__twig_cloud_prompt_v1__:"), - ); - expect(args.params?.content).toEqual( - expect.stringContaining('"type":"resource"'), + expect(mockTrpcCloudTask.sendCommand.mutate).toHaveBeenCalledWith( + expect.objectContaining({ + params: { + content: "read this", + artifact_ids: ["artifact-1"], + }, + }), ); }); @@ -1173,6 +1261,124 @@ describe("SessionService", () => { ); }); + it("preserves attachment blocks in the optimistic resume event", async () => { + const service = getSessionService(); + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( + createMockSession({ + isCloud: true, + cloudStatus: "completed", + cloudBranch: "feature/cloud-run", + }), + ); + mockAuthenticatedClient.getTaskRun.mockResolvedValue({ + id: "run-123", + task: "task-123", + team: 123, + branch: "feature/cloud-run", + runtime_adapter: "claude", + model: "claude-sonnet-4-20250514", + reasoning_effort: null, + environment: "cloud", + status: "completed", + log_url: "https://example.com/logs/run-123", + error_message: null, + output: {}, + state: {}, + created_at: "2026-04-14T00:00:00Z", + updated_at: "2026-04-14T00:00:00Z", + completed_at: "2026-04-14T00:05:00Z", + }); + mockAuthenticatedClient.getTask.mockResolvedValue(createMockTask()); + mockTrpcFs.readFileAsBase64.query.mockResolvedValue("aGVsbG8="); + mockAuthenticatedClient.prepareTaskStagedArtifactUploads.mockResolvedValue( + [ + { + id: "artifact-1", + name: "test.txt", + type: "user_attachment", + source: "posthog_code", + size: 5, + content_type: "text/plain", + storage_path: "tasks/artifacts/test.txt", + expires_in: 3600, + presigned_post: { + url: "https://uploads.example.com", + fields: { key: "tasks/artifacts/test.txt" }, + }, + }, + ], + ); + mockAuthenticatedClient.finalizeTaskStagedArtifactUploads.mockResolvedValue( + [ + { + id: "artifact-1", + name: "test.txt", + type: "user_attachment", + source: "posthog_code", + size: 5, + content_type: "text/plain", + storage_path: "tasks/artifacts/test.txt", + uploaded_at: "2026-04-16T00:00:00Z", + }, + ], + ); + mockAuthenticatedClient.runTaskInCloud.mockResolvedValue( + createMockTask({ + latest_run: { + id: "run-456", + task: "task-123", + team: 123, + branch: "feature/cloud-run", + runtime_adapter: "claude", + model: "claude-sonnet-4-20250514", + reasoning_effort: null, + environment: "cloud", + status: "queued", + log_url: "https://example.com/logs/run-456", + error_message: null, + output: {}, + state: {}, + created_at: "2026-04-14T00:06:00Z", + updated_at: "2026-04-14T00:06:00Z", + completed_at: null, + }, + }), + ); + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue({ ok: true } as Response), + ); + + const prompt: ContentBlock[] = [ + { type: "text", text: "what is this about?" }, + { + type: "resource_link", + uri: "file:///tmp/test.txt", + name: "test.txt", + mimeType: "text/plain", + }, + ]; + + const result = await service.sendPrompt("task-123", prompt); + + expect(result.stopReason).toBe("queued"); + expect(mockSessionStoreSetters.setSession).toHaveBeenCalledWith( + expect.objectContaining({ + events: expect.arrayContaining([ + expect.objectContaining({ + message: expect.objectContaining({ + method: "session/prompt", + params: { + prompt, + }, + }), + }), + ]), + skipPolledPromptCount: 1, + }), + ); + }); + it("attempts automatic recovery on fatal error", async () => { const service = getSessionService(); const mockSession = createMockSession({ diff --git a/apps/code/src/renderer/features/sessions/service/service.ts b/apps/code/src/renderer/features/sessions/service/service.ts index a6f7df6a9..2146a877a 100644 --- a/apps/code/src/renderer/features/sessions/service/service.ts +++ b/apps/code/src/renderer/features/sessions/service/service.ts @@ -8,11 +8,6 @@ import { getAuthenticatedClient, } from "@features/auth/hooks/authClient"; import { fetchAuthState } from "@features/auth/hooks/authQueries"; -import { - buildCloudPromptBlocks, - buildCloudTaskDescription, - serializeCloudPrompt, -} from "@features/editor/utils/cloud-prompt"; import { useSessionAdapterStore } from "@features/sessions/stores/sessionAdapterStore"; import { getPersistedConfigOptions, @@ -68,7 +63,7 @@ import { import { queryClient } from "@utils/queryClient"; import { convertStoredEntriesToEvents, - createUserMessageEvent, + createUserPromptEvent, createUserShellExecuteEvent, extractPromptText, getUserShellExecutesSinceLastPrompt, @@ -76,6 +71,13 @@ import { normalizePromptToBlocks, shellExecutesToContextBlocks, } from "@utils/session"; +import { + cloudPromptToBlocks, + combineQueuedCloudPrompts, + getCloudPromptTransport, + uploadRunAttachments, + uploadTaskStagedAttachments, +} from "../utils/cloudArtifacts"; const log = logger.scope("session-service"); const LOCAL_SESSION_RECONNECT_ATTEMPTS = 3; @@ -129,7 +131,7 @@ function buildCloudDefaultConfigOptions( interface AuthCredentials { apiHost: string; projectId: number; - client: Awaited>; + client: NonNullable>>; } export interface ConnectParams { @@ -1422,41 +1424,26 @@ export class SessionService { // --- Cloud Commands --- - private async prepareCloudPrompt( - prompt: string | ContentBlock[], - ): Promise<{ blocks: ContentBlock[]; promptText: string }> { - const blocks = - typeof prompt === "string" - ? await buildCloudPromptBlocks(prompt) - : prompt; - - if (blocks.length === 0) { - throw new Error("Cloud prompt cannot be empty"); - } - - const promptText = - extractPromptText(blocks).trim() || - (typeof prompt === "string" ? buildCloudTaskDescription(prompt) : ""); - - return { blocks, promptText }; - } - private async sendCloudPrompt( session: AgentSession, prompt: string | ContentBlock[], options?: { skipQueueGuard?: boolean }, ): Promise<{ stopReason: string }> { - const rawPromptText = extractPromptText(prompt); - if (!rawPromptText.trim()) { + const transport = getCloudPromptTransport(prompt); + if (!transport.messageText && transport.filePaths.length === 0) { return { stopReason: "empty" }; } if (isTerminalStatus(session.cloudStatus)) { - return this.resumeCloudRun(session, rawPromptText); + return this.resumeCloudRun(session, prompt); } if (!options?.skipQueueGuard && session.isPromptPending) { - sessionStoreSetters.enqueueMessage(session.taskId, rawPromptText); + sessionStoreSetters.enqueueMessage( + session.taskId, + transport.promptText, + prompt, + ); log.info("Cloud message queued", { taskId: session.taskId, queueLength: session.messageQueue.length + 1, @@ -1464,12 +1451,26 @@ export class SessionService { return { stopReason: "queued" }; } - const auth = await this.getCloudCommandAuth(); - if (!auth) { + const [auth, cloudCommandAuth] = await Promise.all([ + this.getAuthCredentials(), + this.getCloudCommandAuth(), + ]); + if (!auth || !cloudCommandAuth) { throw new Error("Authentication required for cloud commands"); } - - const { blocks, promptText } = await this.prepareCloudPrompt(prompt); + const artifactIds = await uploadRunAttachments( + auth.client, + session.taskId, + session.taskRunId, + transport.filePaths, + ); + const params: Record = {}; + if (transport.messageText) { + params.content = transport.messageText; + } + if (artifactIds.length > 0) { + params.artifact_ids = artifactIds; + } sessionStoreSetters.updateSession(session.taskRunId, { isPromptPending: true, @@ -1479,21 +1480,17 @@ export class SessionService { task_id: session.taskId, is_initial: session.events.length === 0, execution_type: "cloud", - prompt_length_chars: promptText.length, + prompt_length_chars: transport.promptText.length, }); try { const result = await trpcClient.cloudTask.sendCommand.mutate({ taskId: session.taskId, runId: session.taskRunId, - apiHost: auth.apiHost, - teamId: auth.teamId, + apiHost: cloudCommandAuth.apiHost, + teamId: cloudCommandAuth.teamId, method: "user_message", - params: { - // The live /command API still validates user_message content as a - // string, so structured prompts must go through the serialized form. - content: serializeCloudPrompt(blocks), - }, + params, }); sessionStoreSetters.updateSession(session.taskRunId, { @@ -1533,12 +1530,13 @@ export class SessionService { private async sendQueuedCloudMessages( taskId: string, attempt = 0, - pendingText?: string, + pendingPrompt?: string | ContentBlock[], ): Promise<{ stopReason: string }> { - // First attempt: atomically dequeue. Retries reuse the already-dequeued text. - const combinedText = - pendingText ?? sessionStoreSetters.dequeueMessagesAsText(taskId); - if (!combinedText) return { stopReason: "skipped" }; + // First attempt: atomically dequeue. Retries reuse the already-dequeued prompt. + const combinedPrompt = + pendingPrompt ?? + combineQueuedCloudPrompts(sessionStoreSetters.dequeueMessages(taskId)); + if (!combinedPrompt) return { stopReason: "skipped" }; const session = sessionStoreSetters.getSessionByTaskId(taskId); if (!session) { @@ -1550,12 +1548,12 @@ export class SessionService { log.info("Sending queued cloud messages", { taskId, - promptLength: combinedText.length, + promptLength: combinedPrompt.length, attempt, }); try { - return await this.sendCloudPrompt(session, combinedText, { + return await this.sendCloudPrompt(session, combinedPrompt, { skipQueueGuard: true, }); } catch (error) { @@ -1574,7 +1572,7 @@ export class SessionService { this.sendQueuedCloudMessages( taskId, attempt + 1, - combinedText, + combinedPrompt, ).catch((err) => { log.error("Queued cloud message retry failed", { taskId, @@ -1601,8 +1599,8 @@ export class SessionService { session: AgentSession, prompt: string | ContentBlock[], ): Promise<{ stopReason: string }> { - const client = await getAuthenticatedClient(); - if (!client) { + const authCredentials = await this.getAuthCredentials(); + if (!authCredentials) { throw new Error("Authentication required for cloud commands"); } const auth = await this.getCloudCommandAuth(); @@ -1610,11 +1608,19 @@ export class SessionService { throw new Error("Authentication required for cloud commands"); } - const { blocks, promptText } = await this.prepareCloudPrompt(prompt); + const transport = getCloudPromptTransport(prompt); + if (!transport.messageText && transport.filePaths.length === 0) { + return { stopReason: "empty" }; + } + const artifactIds = await uploadTaskStagedAttachments( + authCredentials.client, + session.taskId, + transport.filePaths, + ); const [previousRun, task] = await Promise.all([ - client.getTaskRun(session.taskId, session.taskRunId), - client.getTask(session.taskId), + authCredentials.client.getTaskRun(session.taskId, session.taskRunId), + authCredentials.client.getTask(session.taskId), ]); const hasGitHubRepo = !!task.repository && !!task.github_integration; const previousState = previousRun.state as Record; @@ -1652,7 +1658,7 @@ export class SessionService { // Create a new run WITH resume context — backend validates the previous run, // derives snapshot_external_id server-side, and passes everything as extra_state. // The agent will load conversation history and restore the sandbox snapshot. - const updatedTask = await client.runTaskInCloud( + const updatedTask = await authCredentials.client.runTaskInCloud( session.taskId, previousBaseBranch, { @@ -1660,7 +1666,9 @@ export class SessionService { model: runtimeOptions.model, reasoningLevel: runtimeOptions.reasoningLevel, resumeFromRunId: session.taskRunId, - pendingUserMessage: serializeCloudPrompt(blocks), + pendingUserMessage: transport.messageText, + pendingUserArtifactIds: + artifactIds.length > 0 ? artifactIds : undefined, prAuthorshipMode, runSource: this.getCloudRunSource(previousState), signalReportId: @@ -1688,7 +1696,12 @@ export class SessionService { // Reset processedLineCount to 0 because the new run's log stream starts fresh. newSession.events = [ ...session.events, - createUserMessageEvent(promptText, Date.now()), + createUserPromptEvent( + transport.filePaths.length > 0 + ? cloudPromptToBlocks(prompt) + : [{ type: "text", text: transport.promptText }], + Date.now(), + ), ]; newSession.processedLineCount = 0; // Skip the first session/prompt from polled logs — we already have the @@ -1729,7 +1742,7 @@ export class SessionService { task_id: session.taskId, is_initial: false, execution_type: "cloud", - prompt_length_chars: promptText.length, + prompt_length_chars: transport.promptText.length, }); return { stopReason: "queued" }; diff --git a/apps/code/src/renderer/features/sessions/stores/sessionStore.ts b/apps/code/src/renderer/features/sessions/stores/sessionStore.ts index d5752c845..dc4cfd8fa 100644 --- a/apps/code/src/renderer/features/sessions/stores/sessionStore.ts +++ b/apps/code/src/renderer/features/sessions/stores/sessionStore.ts @@ -19,6 +19,7 @@ export type Adapter = "claude" | "codex"; export interface QueuedMessage { id: string; content: string; + rawPrompt?: string | ContentBlock[]; queuedAt: number; } @@ -292,7 +293,11 @@ export const sessionStoreSetters = { }); }, - enqueueMessage: (taskId: string, content: string) => { + enqueueMessage: ( + taskId: string, + content: string, + rawPrompt?: string | ContentBlock[], + ) => { const id = `queue-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`; useSessionStore.setState((state) => { const taskRunId = state.taskIdIndex[taskId]; @@ -300,7 +305,12 @@ export const sessionStoreSetters = { const session = state.sessions[taskRunId]; if (session) { - session.messageQueue.push({ id, content, queuedAt: Date.now() }); + session.messageQueue.push({ + id, + content, + rawPrompt, + queuedAt: Date.now(), + }); } }); }, @@ -345,6 +355,21 @@ export const sessionStoreSetters = { return result; }, + dequeueMessages: (taskId: string): QueuedMessage[] => { + let queuedMessages: QueuedMessage[] = []; + useSessionStore.setState((state) => { + const taskRunId = state.taskIdIndex[taskId]; + if (!taskRunId) return; + + const session = state.sessions[taskRunId]; + if (!session || session.messageQueue.length === 0) return; + + queuedMessages = [...session.messageQueue]; + session.messageQueue = []; + }); + return queuedMessages; + }, + appendOptimisticItem: ( taskRunId: string, item: Omit, diff --git a/apps/code/src/renderer/features/sessions/utils/cloudArtifacts.test.ts b/apps/code/src/renderer/features/sessions/utils/cloudArtifacts.test.ts new file mode 100644 index 000000000..32804bda5 --- /dev/null +++ b/apps/code/src/renderer/features/sessions/utils/cloudArtifacts.test.ts @@ -0,0 +1,98 @@ +import type { ContentBlock } from "@agentclientprotocol/sdk"; +import { describe, expect, it, vi } from "vitest"; + +vi.mock("@renderer/trpc/client", () => ({ + trpcClient: { + fs: { + readFileAsBase64: { + query: vi.fn(), + }, + }, + }, +})); + +import { trpcClient } from "@renderer/trpc/client"; + +import { + CLOUD_ATTACHMENT_MAX_SIZE_BYTES, + CLOUD_PDF_ATTACHMENT_MAX_SIZE_BYTES, + combineQueuedCloudPrompts, + promptToQueuedEditorContent, + uploadRunAttachments, +} from "./cloudArtifacts"; + +describe("cloudArtifacts", () => { + it("preserves attachment blocks when combining queued cloud prompts", () => { + const prompt: ContentBlock[] = [ + { type: "text", text: "read this" }, + { + type: "resource_link", + uri: "file:///tmp/test.txt", + name: "test.txt", + mimeType: "text/plain", + }, + ]; + + expect( + combineQueuedCloudPrompts([ + { + content: "read this\n\nAttached files: test.txt", + rawPrompt: prompt, + }, + ]), + ).toEqual(prompt); + }); + + it("rejects attachments that exceed the max size", async () => { + const oversizedByteLength = CLOUD_ATTACHMENT_MAX_SIZE_BYTES + 1; + const base64 = btoa("a".repeat(oversizedByteLength)); + vi.mocked(trpcClient.fs.readFileAsBase64.query).mockResolvedValueOnce( + base64, + ); + + const client = { + prepareTaskRunArtifactUploads: vi.fn(), + finalizeTaskRunArtifactUploads: vi.fn(), + } as never; + + await expect( + uploadRunAttachments(client, "task-1", "run-1", ["/tmp/huge.bin"]), + ).rejects.toThrow(/exceeds the 30MB attachment limit/); + }); + + it("rejects PDFs that exceed the stricter cloud limit", async () => { + const oversizedByteLength = CLOUD_PDF_ATTACHMENT_MAX_SIZE_BYTES + 1; + const base64 = btoa("a".repeat(oversizedByteLength)); + vi.mocked(trpcClient.fs.readFileAsBase64.query).mockResolvedValueOnce( + base64, + ); + + const client = { + prepareTaskRunArtifactUploads: vi.fn(), + finalizeTaskRunArtifactUploads: vi.fn(), + } as never; + + await expect( + uploadRunAttachments(client, "task-1", "run-1", ["/tmp/large.pdf"]), + ).rejects.toThrow( + /exceeds the 10MB attachment limit for PDFs in cloud runs/, + ); + }); + + it("restores queued editor content with attachments from prompt blocks", () => { + const prompt: ContentBlock[] = [ + { type: "text", text: "read this" }, + { + type: "resource_link", + uri: "file:///tmp/test.txt", + name: "test.txt", + mimeType: "text/plain", + }, + ]; + + expect(promptToQueuedEditorContent(prompt)).toEqual({ + segments: [{ type: "text", text: "read this" }], + attachments: [{ id: "/tmp/test.txt", label: "test.txt" }], + }); + }); +}); diff --git a/apps/code/src/renderer/features/sessions/utils/cloudArtifacts.ts b/apps/code/src/renderer/features/sessions/utils/cloudArtifacts.ts new file mode 100644 index 000000000..b916c7f60 --- /dev/null +++ b/apps/code/src/renderer/features/sessions/utils/cloudArtifacts.ts @@ -0,0 +1,417 @@ +import type { ContentBlock } from "@agentclientprotocol/sdk"; +import { + buildCloudTaskDescription, + getAbsoluteAttachmentPaths, + stripAbsoluteFileTags, +} from "@features/editor/utils/cloud-prompt"; +import type { + PostHogAPIClient, + PreparedTaskArtifactUpload, + TaskArtifactUploadRequest, +} from "@renderer/api/posthogClient"; +import { trpcClient } from "@renderer/trpc/client"; +import { getFileName } from "@utils/path"; +import type { EditorContent } from "../../message-editor/utils/content"; + +const FILE_URI_PREFIX = "file://"; +const ATTACHMENT_SOURCE = "posthog_code"; +const DEFAULT_CONTENT_TYPE = "application/octet-stream"; +export const CLOUD_ATTACHMENT_MAX_SIZE_BYTES = 30 * 1024 * 1024; +export const CLOUD_PDF_ATTACHMENT_MAX_SIZE_BYTES = 10 * 1024 * 1024; + +const CONTENT_TYPE_BY_EXTENSION: Record = { + bmp: "image/bmp", + c: "text/plain", + cc: "text/plain", + conf: "text/plain", + cpp: "text/plain", + css: "text/css", + csv: "text/csv", + gif: "image/gif", + go: "text/plain", + h: "text/plain", + html: "text/html", + ini: "text/plain", + java: "text/plain", + jpeg: "image/jpeg", + jpg: "image/jpeg", + js: "text/javascript", + json: "application/json", + jsx: "text/javascript", + log: "text/plain", + md: "text/markdown", + pdf: "application/pdf", + png: "image/png", + py: "text/x-python", + rb: "text/plain", + rs: "text/plain", + sh: "text/x-shellscript", + sql: "application/sql", + svg: "image/svg+xml", + toml: "application/toml", + ts: "text/typescript", + tsx: "text/typescript", + txt: "text/plain", + webp: "image/webp", + xml: "application/xml", + yaml: "application/yaml", + yml: "application/yaml", + zip: "application/zip", +}; + +interface LoadedCloudAttachment { + filePath: string; + bytes: Uint8Array; + upload: TaskArtifactUploadRequest; +} + +function pathToFileUri(filePath: string): string { + const encoded = filePath + .split("/") + .map((segment) => encodeURIComponent(segment)) + .join("/"); + return `file://${encoded}`; +} + +export interface CloudPromptTransport { + filePaths: string[]; + messageText?: string; + promptText: string; +} + +export type QueuedCloudPrompt = string | ContentBlock[]; + +function base64ToUint8Array(base64: string): Uint8Array { + const binary = atob(base64); + const bytes = new Uint8Array(new ArrayBuffer(binary.length)); + + for (let index = 0; index < binary.length; index += 1) { + bytes[index] = binary.charCodeAt(index); + } + + return bytes; +} + +function getFileExtension(filePath: string): string { + const parts = getFileName(filePath).split("."); + return parts.length > 1 ? (parts.at(-1)?.toLowerCase() ?? "") : ""; +} + +function inferContentType(filePath: string): string { + return ( + CONTENT_TYPE_BY_EXTENSION[getFileExtension(filePath)] ?? + DEFAULT_CONTENT_TYPE + ); +} + +function getCloudAttachmentMaxSizeBytes( + filePath: string, + contentType: string, +): number { + const extension = getFileExtension(filePath); + const normalizedContentType = + contentType.split(";")[0]?.trim().toLowerCase() ?? ""; + + if (extension === "pdf" || normalizedContentType === "application/pdf") { + return CLOUD_PDF_ATTACHMENT_MAX_SIZE_BYTES; + } + + return CLOUD_ATTACHMENT_MAX_SIZE_BYTES; +} + +function getCloudAttachmentSizeError( + filePath: string, + maxSizeBytes: number, +): string { + const maxMb = Math.floor(maxSizeBytes / (1024 * 1024)); + + if (getFileExtension(filePath) === "pdf") { + return `${getFileName(filePath)} exceeds the ${maxMb}MB attachment limit for PDFs in cloud runs`; + } + + return `${getFileName(filePath)} exceeds the ${maxMb}MB attachment limit`; +} + +function decodeFileUri(uri: string): string | null { + if (!uri.startsWith(FILE_URI_PREFIX)) { + return null; + } + + const encodedPath = uri.slice(FILE_URI_PREFIX.length); + const normalizedPath = encodedPath.startsWith("/") + ? encodedPath + : `/${encodedPath}`; + + try { + return normalizedPath + .split("/") + .map((segment, index) => + index === 0 && segment === "" ? segment : decodeURIComponent(segment), + ) + .join("/"); + } catch { + return null; + } +} + +function collectBlockAttachmentPaths(prompt: ContentBlock[]): string[] { + const filePaths = prompt + .map((block) => { + if (block.type === "resource_link") { + return decodeFileUri(block.uri); + } + + if (block.type === "resource") { + return block.resource.uri ? decodeFileUri(block.resource.uri) : null; + } + + if (block.type === "image") { + return block.uri ? decodeFileUri(block.uri) : null; + } + + return null; + }) + .filter((value): value is string => Boolean(value)); + + return Array.from(new Set(filePaths)); +} + +function summarizePrompt(text: string, filePaths: string[]): string { + if (filePaths.length === 0) { + return text.trim(); + } + + const attachmentSummary = `Attached files: ${filePaths.map(getFileName).join(", ")}`; + return text.trim() + ? `${text.trim()}\n\n${attachmentSummary}` + : attachmentSummary; +} + +async function loadCloudAttachments( + filePaths: string[], +): Promise { + return Promise.all( + filePaths.map(async (filePath) => { + const base64 = await trpcClient.fs.readFileAsBase64.query({ filePath }); + if (!base64) { + throw new Error( + `Unable to read attached file ${getFileName(filePath)}`, + ); + } + + const bytes = base64ToUint8Array(base64); + const contentType = inferContentType(filePath); + const maxSizeBytes = getCloudAttachmentMaxSizeBytes( + filePath, + contentType, + ); + if (bytes.byteLength > maxSizeBytes) { + throw new Error(getCloudAttachmentSizeError(filePath, maxSizeBytes)); + } + return { + filePath, + bytes, + upload: { + name: getFileName(filePath), + type: "user_attachment", + source: ATTACHMENT_SOURCE, + size: bytes.byteLength, + content_type: contentType, + }, + }; + }), + ); +} + +async function uploadPreparedArtifacts( + attachments: LoadedCloudAttachment[], + preparedArtifacts: PreparedTaskArtifactUpload[], +): Promise { + if (attachments.length !== preparedArtifacts.length) { + throw new Error("Prepared uploads do not match the selected attachments"); + } + + await Promise.all( + preparedArtifacts.map(async (preparedArtifact, index) => { + const attachment = attachments[index]; + const formData = new FormData(); + + for (const [key, value] of Object.entries( + preparedArtifact.presigned_post.fields, + )) { + formData.append(key, value); + } + + formData.append( + "file", + new Blob([attachment.bytes], { + type: attachment.upload.content_type || DEFAULT_CONTENT_TYPE, + }), + attachment.upload.name, + ); + + const response = await fetch(preparedArtifact.presigned_post.url, { + method: "POST", + body: formData, + }); + + if (!response.ok) { + throw new Error(`Failed to upload ${attachment.upload.name}`); + } + }), + ); +} + +export function getCloudPromptTransport( + prompt: string | ContentBlock[], + filePaths: string[] = [], +): CloudPromptTransport { + if (typeof prompt === "string") { + const attachmentPaths = getAbsoluteAttachmentPaths(prompt, filePaths); + const messageText = stripAbsoluteFileTags(prompt).trim(); + + return { + filePaths: attachmentPaths, + messageText: messageText || undefined, + promptText: buildCloudTaskDescription(prompt, filePaths).trim(), + }; + } + + const promptText = prompt + .filter( + (block): block is Extract => + block.type === "text", + ) + .map((block) => block.text) + .join("") + .trim(); + const attachmentPaths = collectBlockAttachmentPaths(prompt); + + return { + filePaths: attachmentPaths, + messageText: promptText || undefined, + promptText: summarizePrompt(promptText, attachmentPaths), + }; +} + +export function cloudPromptToBlocks(prompt: QueuedCloudPrompt): ContentBlock[] { + if (typeof prompt !== "string") { + return prompt; + } + + const transport = getCloudPromptTransport(prompt); + const blocks: ContentBlock[] = []; + + if (transport.messageText) { + blocks.push({ type: "text", text: transport.messageText }); + } + + for (const filePath of transport.filePaths) { + blocks.push({ + type: "resource_link", + uri: pathToFileUri(filePath), + name: getFileName(filePath), + }); + } + + return blocks; +} + +export async function uploadTaskStagedAttachments( + client: PostHogAPIClient, + taskId: string, + filePaths: string[], +): Promise { + if (!filePaths.length) { + return []; + } + + const attachments = await loadCloudAttachments(filePaths); + const preparedArtifacts = await client.prepareTaskStagedArtifactUploads( + taskId, + attachments.map((attachment) => attachment.upload), + ); + + await uploadPreparedArtifacts(attachments, preparedArtifacts); + + const finalizedArtifacts = await client.finalizeTaskStagedArtifactUploads( + taskId, + preparedArtifacts, + ); + + return finalizedArtifacts.map((artifact) => artifact.id); +} + +export async function uploadRunAttachments( + client: PostHogAPIClient, + taskId: string, + runId: string, + filePaths: string[], +): Promise { + if (!filePaths.length) { + return []; + } + + const attachments = await loadCloudAttachments(filePaths); + const preparedArtifacts = await client.prepareTaskRunArtifactUploads( + taskId, + runId, + attachments.map((attachment) => attachment.upload), + ); + + await uploadPreparedArtifacts(attachments, preparedArtifacts); + + const finalizedArtifacts = await client.finalizeTaskRunArtifactUploads( + taskId, + runId, + preparedArtifacts, + ); + + return finalizedArtifacts.map((artifact) => artifact.id); +} + +export function promptToQueuedEditorContent( + prompt: QueuedCloudPrompt, +): EditorContent { + const transport = getCloudPromptTransport(prompt); + const attachments = transport.filePaths.map((filePath) => ({ + id: filePath, + label: getFileName(filePath), + })); + const text = + typeof prompt === "string" + ? stripAbsoluteFileTags(prompt) + : (transport.messageText ?? ""); + + return { + segments: [{ type: "text", text }], + ...(attachments.length > 0 ? { attachments } : {}), + }; +} + +export function combineQueuedCloudPrompts( + queuedPrompts: Array<{ content: string; rawPrompt?: QueuedCloudPrompt }>, +): QueuedCloudPrompt | null { + if (queuedPrompts.length === 0) { + return null; + } + + const blocks: ContentBlock[] = []; + + for (const [index, queuedPrompt] of queuedPrompts.entries()) { + const promptBlocks = cloudPromptToBlocks( + queuedPrompt.rawPrompt ?? queuedPrompt.content, + ); + if (promptBlocks.length === 0) { + continue; + } + + if (index > 0 && blocks.length > 0) { + blocks.push({ type: "text", text: "\n\n" }); + } + + blocks.push(...promptBlocks); + } + + return blocks.length > 0 ? blocks : null; +} diff --git a/apps/code/src/renderer/sagas/task/task-creation.test.ts b/apps/code/src/renderer/sagas/task/task-creation.test.ts index 5570eb927..6e3d1b499 100644 --- a/apps/code/src/renderer/sagas/task/task-creation.test.ts +++ b/apps/code/src/renderer/sagas/task/task-creation.test.ts @@ -255,15 +255,44 @@ describe("TaskCreationSaga", () => { }); }); - it("sends initial cloud prompts with attachments as pending user messages", async () => { + it("uploads initial cloud attachments before starting the run", async () => { const createdTask = createTask(); const startedTask = createTask({ latest_run: createRun() }); const createTaskMock = vi.fn().mockResolvedValue(createdTask); const runTaskInCloudMock = vi.fn().mockResolvedValue(startedTask); + const prepareTaskStagedArtifactUploadsMock = vi.fn().mockResolvedValue([ + { + id: "artifact-1", + name: "test.txt", + type: "user_attachment", + size: 5, + source: "posthog_code", + content_type: "text/plain", + storage_path: "tasks/artifacts/test.txt", + expires_in: 3600, + presigned_post: { + url: "https://uploads.example.com", + fields: { key: "tasks/artifacts/test.txt" }, + }, + }, + ]); + const finalizeTaskStagedArtifactUploadsMock = vi.fn().mockResolvedValue([ + { + id: "artifact-1", + name: "test.txt", + type: "user_attachment", + size: 5, + source: "posthog_code", + content_type: "text/plain", + storage_path: "tasks/artifacts/test.txt", + uploaded_at: "2026-04-16T00:00:00Z", + }, + ]); const sendRunCommandMock = vi.fn(); const onTaskReady = vi.fn(); - mockReadAbsoluteFile.mockResolvedValue("hello from attachment"); + mockReadFileAsBase64.mockResolvedValue("aGVsbG8="); + vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: true } as Response)); const saga = new TaskCreationSaga({ posthogClient: { @@ -271,6 +300,9 @@ describe("TaskCreationSaga", () => { deleteTask: vi.fn(), getTask: vi.fn(), runTaskInCloud: runTaskInCloudMock, + prepareTaskStagedArtifactUploads: prepareTaskStagedArtifactUploadsMock, + finalizeTaskStagedArtifactUploads: + finalizeTaskStagedArtifactUploadsMock, sendRunCommand: sendRunCommandMock, updateTask: vi.fn(), } as never, @@ -306,9 +338,8 @@ describe("TaskCreationSaga", () => { adapter: "codex", model: "gpt-5.4", reasoningLevel: "medium", - pendingUserMessage: expect.stringContaining( - "__twig_cloud_prompt_v1__:", - ), + pendingUserMessage: "read this file", + pendingUserArtifactIds: ["artifact-1"], sandboxEnvironmentId: undefined, prAuthorshipMode: "bot", runSource: "manual", diff --git a/apps/code/src/renderer/sagas/task/task-creation.ts b/apps/code/src/renderer/sagas/task/task-creation.ts index 20052070d..72d9f5b70 100644 --- a/apps/code/src/renderer/sagas/task/task-creation.ts +++ b/apps/code/src/renderer/sagas/task/task-creation.ts @@ -1,7 +1,3 @@ -import { - buildCloudPromptBlocks, - serializeCloudPrompt, -} from "@features/editor/utils/cloud-prompt"; import { buildPromptBlocks } from "@features/editor/utils/prompt-builder"; import { DEFAULT_PANEL_IDS } from "@features/panels/constants/panelConstants"; import { usePanelLayoutStore } from "@features/panels/store/panelLayoutStore"; @@ -10,6 +6,10 @@ import { type ConnectParams, getSessionService, } from "@features/sessions/service/service"; +import { + getCloudPromptTransport, + uploadTaskStagedAttachments, +} from "@features/sessions/utils/cloudArtifacts"; import { getTaskDirectory } from "@hooks/useRepositoryDirectory"; import type { Workspace, @@ -117,13 +117,6 @@ export class TaskCreationSaga extends Saga< protected async execute( input: TaskCreationInput, ): Promise { - const initialCloudPrompt = - input.workspaceMode === "cloud" && !input.taskId && input.content - ? await this.readOnlyStep("cloud_prompt_preparation", () => - buildCloudPromptBlocks(input.content ?? "", input.filePaths), - ) - : null; - // Step 1: Get or create task // For new tasks, start folder registration in parallel with task creation // since folder_registration only needs repoPath (from input), not task.id @@ -290,13 +283,28 @@ export class TaskCreationSaga extends Saga< githubUserToken = await getGhUserTokenOrThrow(); } + const transport = + (input.content || input.filePaths?.length) && + workspaceMode === "cloud" + ? getCloudPromptTransport(input.content ?? "", input.filePaths) + : null; + const pendingUserArtifactIds = transport + ? await uploadTaskStagedAttachments( + this.deps.posthogClient, + task.id, + transport.filePaths, + ) + : []; + return this.deps.posthogClient.runTaskInCloud(task.id, branch, { adapter: input.adapter, model: input.model, reasoningLevel: input.reasoningLevel, - pendingUserMessage: initialCloudPrompt - ? serializeCloudPrompt(initialCloudPrompt) - : undefined, + pendingUserMessage: transport?.messageText, + pendingUserArtifactIds: + pendingUserArtifactIds.length > 0 + ? pendingUserArtifactIds + : undefined, sandboxEnvironmentId: input.sandboxEnvironmentId, prAuthorshipMode, runSource: input.cloudRunSource ?? "manual", diff --git a/apps/code/src/renderer/utils/session.ts b/apps/code/src/renderer/utils/session.ts index 6809b35a1..fc846c568 100644 --- a/apps/code/src/renderer/utils/session.ts +++ b/apps/code/src/renderer/utils/session.ts @@ -49,6 +49,27 @@ export function createUserMessageEvent(text: string, ts: number): AcpMessage { }; } +/** + * Create a user prompt event from structured content blocks for display. + */ +export function createUserPromptEvent( + prompt: ContentBlock[], + ts: number, +): AcpMessage { + return { + type: "acp_message", + ts, + message: { + jsonrpc: "2.0", + id: ts, + method: "session/prompt", + params: { + prompt, + }, + } as JsonRpcRequest, + }; +} + /** * Create a user shell execute event. * When id is provided, it's used to track async execution (start/complete).