Prepare prompt file attachments for future UI (#1376)

* Support file uploads to agents * Harden prompt file upload handling * Fix duplicate upload retry handling * Update relay reconnect tests for queued messages * Make upload timeout track idle progress * Clean failed upload directories
2026-06-06 20:37:38 +08:00
parent 9a8912b3ef
commit db4376d17a
16 changed files with 917 additions and 21 deletions
@@ -2,6 +2,7 @@ import { afterEach, expect, expectTypeOf, test, vi } from "vitest";
 import { z } from "zod";
 import { DaemonClient, type DaemonTransport } from "./daemon-client";
 import {
+  decodeFileTransferFrame,
  encodeFileTransferFrame,
  FileTransferOpcode,
 } from "@getpaseo/protocol/binary-frames/index";
@@ -103,6 +104,12 @@ function assertStr(data: string | Uint8Array | ArrayBuffer | undefined): string
  return data;
 }

+function assertUint8Array(data: string | Uint8Array | ArrayBuffer | undefined): Uint8Array {
+  if (data instanceof Uint8Array) return data;
+  if (data instanceof ArrayBuffer) return new Uint8Array(data);
+  throw new Error("Expected binary frame");
+}
+
 function parseSentFrame(
  data: string | Uint8Array | ArrayBuffer | undefined,
 ): Record<string, unknown> {
@@ -567,6 +574,110 @@ test("readFile resolves from binary file frames when the daemon supports them",
  expect(new TextDecoder().decode(result.bytes)).toBe("hello");
 });

+test("uploadFile sends metadata request and file bytes as binary chunks", async () => {
+  const logger = createMockLogger();
+  const mock = createMockTransport();
+
+  const client = new DaemonClient({
+    url: "ws://test",
+    clientId: "clsk_unit_test",
+    logger,
+    reconnect: { enabled: false },
+    transportFactory: () => mock.transport,
+  });
+  clients.push(client);
+
+  const connectPromise = client.connect();
+  mock.triggerOpen();
+  await connectPromise;
+
+  const responsePromise = client.uploadFile({
+    fileName: "notes.txt",
+    mimeType: "text/plain",
+    bytes: new TextEncoder().encode("hello world"),
+    modifiedAt: "2026-05-02T00:00:00.000Z",
+    requestId: "req-upload",
+    chunkSize: 5,
+  });
+
+  expect(JSON.parse(assertStr(mock.sent[0]))).toEqual({
+    type: "session",
+    message: {
+      type: "file.upload.request",
+      fileName: "notes.txt",
+      mimeType: "text/plain",
+      size: 11,
+      modifiedAt: "2026-05-02T00:00:00.000Z",
+      requestId: "req-upload",
+    },
+  });
+  expect(mock.sent.slice(1).map(assertUint8Array).map(decodeFileTransferFrame)).toEqual([
+    {
+      opcode: FileTransferOpcode.FileBegin,
+      requestId: "req-upload",
+      metadata: {
+        mime: "text/plain",
+        size: 11,
+        encoding: "binary",
+        modifiedAt: "2026-05-02T00:00:00.000Z",
+        fileName: "notes.txt",
+      },
+      payload: new Uint8Array(),
+    },
+    {
+      opcode: FileTransferOpcode.FileChunk,
+      requestId: "req-upload",
+      payload: new TextEncoder().encode("hello"),
+    },
+    {
+      opcode: FileTransferOpcode.FileChunk,
+      requestId: "req-upload",
+      payload: new TextEncoder().encode(" worl"),
+    },
+    {
+      opcode: FileTransferOpcode.FileChunk,
+      requestId: "req-upload",
+      payload: new TextEncoder().encode("d"),
+    },
+    {
+      opcode: FileTransferOpcode.FileEnd,
+      requestId: "req-upload",
+      payload: new Uint8Array(),
+    },
+  ]);
+
+  mock.triggerMessage(
+    wrapSessionMessage({
+      type: "file.upload.response",
+      payload: {
+        requestId: "req-upload",
+        file: {
+          type: "uploaded_file",
+          id: "upload_req-upload",
+          fileName: "notes.txt",
+          mimeType: "text/plain",
+          size: 11,
+          path: "/tmp/paseo-uploads/upload_req-upload/notes.txt",
+        },
+        error: null,
+      },
+    }),
+  );
+
+  await expect(responsePromise).resolves.toEqual({
+    requestId: "req-upload",
+    file: {
+      type: "uploaded_file",
+      id: "upload_req-upload",
+      fileName: "notes.txt",
+      mimeType: "text/plain",
+      size: 11,
+      path: "/tmp/paseo-uploads/upload_req-upload/notes.txt",
+    },
+    error: null,
+  });
+});
+
 test("normalizes workspace_setup_progress into a workspace-scoped daemon event", async () => {
  const logger = createMockLogger();
  const mock = createMockTransport();
@@ -22,6 +22,7 @@ import type {
  CreateAgentRequestMessage,
  CreatePaseoWorktreeRequest,
  FileDownloadTokenResponse,
+  FileUploadResponse,
  FileExplorerResponse,
  FetchAgentTimelineResponseMessage,
  GitSetupOptions,
@@ -90,6 +91,7 @@ import { isRelayClientWebSocketUrl } from "@getpaseo/protocol/daemon-endpoints";
 import {
  asUint8Array,
  decodeFileTransferFrame,
+  encodeFileTransferFrame,
  decodeTerminalStreamFrame,
  FileTransferOpcode,
  TerminalStreamOpcode,
@@ -318,6 +320,15 @@ export interface FileReadResult {
  kind: LegacyFileExplorerFilePayload["kind"];
  modifiedAt: string;
 }
+export interface FileUploadInput {
+  fileName: string;
+  mimeType: string;
+  bytes: Uint8Array | ArrayBuffer;
+  modifiedAt?: string;
+  requestId?: string;
+  chunkSize?: number;
+}
+export type FileUploadResult = FileUploadResponse["payload"];
 type FileDownloadTokenPayload = FileDownloadTokenResponse["payload"];
 type ListProviderFeaturesPayload = ListProviderFeaturesResponseMessage["payload"];
 type ListProviderModelsPayload = ListProviderModelsResponseMessage["payload"];
@@ -3316,6 +3327,63 @@ export class DaemonClient {
    }
  }

+  async uploadFile(input: FileUploadInput): Promise<FileUploadResult> {
+    const bytes = asUint8Array(input.bytes);
+    if (!bytes) {
+      throw new Error("File bytes are required.");
+    }
+    const resolvedRequestId = this.createRequestId(input.requestId);
+    const modifiedAt = input.modifiedAt ?? new Date().toISOString();
+    const responsePromise = this.sendCorrelatedRequest({
+      requestId: resolvedRequestId,
+      message: {
+        type: "file.upload.request",
+        fileName: input.fileName,
+        mimeType: input.mimeType,
+        size: bytes.byteLength,
+        modifiedAt,
+        requestId: resolvedRequestId,
+      },
+      responseType: "file.upload.response",
+      timeout: 60000,
+      options: { skipQueue: true },
+    });
+
+    this.sendBinaryFrame(
+      encodeFileTransferFrame({
+        opcode: FileTransferOpcode.FileBegin,
+        requestId: resolvedRequestId,
+        metadata: {
+          mime: input.mimeType,
+          size: bytes.byteLength,
+          encoding: "binary",
+          modifiedAt,
+          fileName: input.fileName,
+        },
+      }),
+    );
+
+    const chunkSize = input.chunkSize ?? 1024 * 1024;
+    for (let offset = 0; offset < bytes.byteLength; offset += chunkSize) {
+      this.sendBinaryFrame(
+        encodeFileTransferFrame({
+          opcode: FileTransferOpcode.FileChunk,
+          requestId: resolvedRequestId,
+          payload: bytes.subarray(offset, Math.min(offset + chunkSize, bytes.byteLength)),
+        }),
+      );
+    }
+
+    this.sendBinaryFrame(
+      encodeFileTransferFrame({
+        opcode: FileTransferOpcode.FileEnd,
+        requestId: resolvedRequestId,
+      }),
+    );
+
+    return responsePromise;
+  }
+
  async requestDownloadToken(
    cwd: string,
    path: string,