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
This commit is contained in:
Mohamed Boudra
2026-06-06 20:37:38 +08:00
committed by GitHub
parent 9a8912b3ef
commit db4376d17a
16 changed files with 917 additions and 21 deletions
+111
View File
@@ -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();
+68
View File
@@ -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,
@@ -0,0 +1,53 @@
import { describe, expect, it } from "vitest";
import {
decodeBinaryFrame,
encodeFileTransferFrame,
encodeTerminalStreamFrame,
FileTransferOpcode,
TerminalStreamOpcode,
} from "./index.js";
describe("binary frame demux", () => {
it("routes terminal frames by opcode", () => {
expect(
decodeBinaryFrame(
encodeTerminalStreamFrame({
opcode: TerminalStreamOpcode.Input,
slot: 7,
payload: "ls",
}),
),
).toEqual({
kind: "terminal",
frame: {
opcode: TerminalStreamOpcode.Input,
slot: 7,
payload: new TextEncoder().encode("ls"),
},
});
});
it("routes file-transfer frames by opcode", () => {
expect(
decodeBinaryFrame(
encodeFileTransferFrame({
opcode: FileTransferOpcode.FileChunk,
requestId: "req-upload",
payload: new TextEncoder().encode("hello"),
}),
),
).toEqual({
kind: "file_transfer",
frame: {
opcode: FileTransferOpcode.FileChunk,
requestId: "req-upload",
payload: new TextEncoder().encode("hello"),
},
});
});
it("rejects unknown binary opcodes", () => {
expect(decodeBinaryFrame(new Uint8Array([0xff, 0]))).toBeNull();
});
});
@@ -0,0 +1,35 @@
import {
decodeFileTransferFrame,
FileTransferOpcode,
type FileTransferFrame,
} from "./file-transfer.js";
import {
decodeTerminalStreamFrame,
TerminalStreamOpcode,
type TerminalStreamFrame,
} from "./terminal.js";
export type BinaryFrame =
| { kind: "terminal"; frame: TerminalStreamFrame }
| { kind: "file_transfer"; frame: FileTransferFrame };
export function decodeBinaryFrame(bytes: Uint8Array): BinaryFrame | null {
switch (bytes[0]) {
case TerminalStreamOpcode.Output:
case TerminalStreamOpcode.Input:
case TerminalStreamOpcode.Resize:
case TerminalStreamOpcode.Snapshot:
case TerminalStreamOpcode.Restore: {
const frame = decodeTerminalStreamFrame(bytes);
return frame ? { kind: "terminal", frame } : null;
}
case FileTransferOpcode.FileBegin:
case FileTransferOpcode.FileChunk:
case FileTransferOpcode.FileEnd: {
const frame = decodeFileTransferFrame(bytes);
return frame ? { kind: "file_transfer", frame } : null;
}
default:
return null;
}
}
@@ -14,6 +14,7 @@ export const FileBeginMetadataSchema = z.object({
size: z.number().int().nonnegative(),
encoding: z.enum(["utf-8", "binary"]),
modifiedAt: z.string(),
fileName: z.string().optional(),
});
export interface FileBegin {
@@ -1,2 +1,3 @@
export * from "./demux.js";
export * from "./file-transfer.js";
export * from "./terminal.js";
+32
View File
@@ -839,11 +839,21 @@ export const ReviewAttachmentSchema = z.object({
comments: z.array(ReviewAttachmentCommentSchema),
});
export const UploadedFileAttachmentSchema = z.object({
type: z.literal("uploaded_file"),
id: z.string(),
fileName: z.string(),
mimeType: z.string(),
size: z.number().int().nonnegative(),
path: z.string(),
});
export const AgentAttachmentSchema = z.discriminatedUnion("type", [
GitHubPrAttachmentSchema,
GitHubIssueAttachmentSchema,
TextAttachmentSchema,
ReviewAttachmentSchema,
UploadedFileAttachmentSchema,
]);
function normalizeAgentAttachments(input: unknown): AgentAttachment[] {
@@ -1680,6 +1690,15 @@ export const FileDownloadTokenRequestSchema = z.object({
requestId: z.string(),
});
export const FileUploadRequestSchema = z.object({
type: z.literal("file.upload.request"),
fileName: z.string().min(1),
mimeType: z.string().min(1),
size: z.number().int().nonnegative(),
modifiedAt: z.string(),
requestId: z.string(),
});
export const ClearAgentAttentionMessageSchema = z.object({
type: z.literal("clear_agent_attention"),
agentId: z.union([z.string(), z.array(z.string())]),
@@ -1911,6 +1930,7 @@ export const SessionInboundMessageSchema = z.discriminatedUnion("type", [
FileExplorerRequestSchema,
ProjectIconRequestSchema,
FileDownloadTokenRequestSchema,
FileUploadRequestSchema,
ClearAgentAttentionMessageSchema,
ClientHeartbeatMessageSchema,
PingMessageSchema,
@@ -3436,6 +3456,15 @@ export const FileDownloadTokenResponseSchema = z.object({
}),
});
export const FileUploadResponseSchema = z.object({
type: z.literal("file.upload.response"),
payload: z.object({
requestId: z.string(),
file: UploadedFileAttachmentSchema.nullable(),
error: z.string().nullable(),
}),
});
export const ListProviderModelsResponseMessageSchema = z.object({
type: z.literal("list_provider_models_response"),
payload: z.object({
@@ -3756,6 +3785,7 @@ export const SessionOutboundMessageSchema = z.discriminatedUnion("type", [
FileExplorerResponseSchema,
ProjectIconResponseSchema,
FileDownloadTokenResponseSchema,
FileUploadResponseSchema,
ListProviderModelsResponseMessageSchema,
ListProviderModesResponseMessageSchema,
ListProviderFeaturesResponseMessageSchema,
@@ -4056,6 +4086,8 @@ export type ProjectIconResponse = z.infer<typeof ProjectIconResponseSchema>;
export type ProjectIcon = z.infer<typeof ProjectIconSchema>;
export type FileDownloadTokenRequest = z.infer<typeof FileDownloadTokenRequestSchema>;
export type FileDownloadTokenResponse = z.infer<typeof FileDownloadTokenResponseSchema>;
export type FileUploadRequest = z.infer<typeof FileUploadRequestSchema>;
export type FileUploadResponse = z.infer<typeof FileUploadResponseSchema>;
export type RestartServerRequestMessage = z.infer<typeof RestartServerRequestMessageSchema>;
export type ShutdownServerRequestMessage = z.infer<typeof ShutdownServerRequestMessageSchema>;
export type ClearAgentAttentionMessage = z.infer<typeof ClearAgentAttentionMessageSchema>;
@@ -96,6 +96,26 @@ describe("prompt attachments", () => {
).toBe("<browser-element>button.primary</browser-element>");
});
it("renders uploaded file attachments as local file references", () => {
expect(
renderPromptAttachmentAsText({
type: "uploaded_file",
id: "upload_req-upload",
fileName: "notes.txt",
mimeType: "text/plain",
size: 11,
path: "/tmp/paseo/uploads/upload_req-upload/notes.txt",
}),
).toBe(
[
"Uploaded file: notes.txt",
"Path: /tmp/paseo/uploads/upload_req-upload/notes.txt",
"MIME: text/plain",
"Size: 11 bytes",
].join("\n"),
);
});
it("returns undefined when firstAgentContext is empty", () => {
expect(buildAgentBranchNameSeed(undefined)).toBeUndefined();
expect(buildAgentBranchNameSeed({})).toBeUndefined();
@@ -54,6 +54,14 @@ export function renderPromptAttachmentAsText(attachment: AgentAttachment): strin
});
return lines.join("\n");
}
case "uploaded_file": {
return [
`Uploaded file: ${attachment.fileName}`,
`Path: ${attachment.path}`,
`MIME: ${attachment.mimeType}`,
`Size: ${attachment.size} bytes`,
].join("\n");
}
default:
throw new Error("unreachable");
}
@@ -137,6 +137,44 @@ test("createAgent without an initial prompt returns an idle snapshot", async ()
}
});
test("DaemonClient uploads file bytes to daemon temp storage", async () => {
const daemon = await createTestPaseoDaemon();
const client = new DaemonClient({
url: `ws://127.0.0.1:${daemon.port}/ws`,
appVersion: "0.1.82",
});
try {
await client.connect();
const result = await client.uploadFile({
fileName: "notes.txt",
mimeType: "text/plain",
bytes: new TextEncoder().encode("hello world"),
modifiedAt: "2026-05-02T00:00:00.000Z",
requestId: "req-upload-e2e",
chunkSize: 5,
});
expect(result).toEqual({
requestId: "req-upload-e2e",
file: {
type: "uploaded_file",
id: "upload_req-upload-e2e",
fileName: "notes.txt",
mimeType: "text/plain",
size: 11,
path: path.join(daemon.paseoHome, "uploads", "upload_req-upload-e2e", "notes.txt"),
},
error: null,
});
await expect(readFile(result.file?.path ?? "", "utf8")).resolves.toBe("hello world");
} finally {
await client.close();
await daemon.close();
}
});
test("createAgent with background initialPrompt returns a running snapshot before turn completion", async () => {
const daemon = await createTestPaseoDaemon();
const client = new DaemonClient({
@@ -0,0 +1,254 @@
import { existsSync, mkdtempSync, readFileSync, realpathSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
decodeFileTransferFrame,
encodeFileTransferFrame,
FileTransferOpcode,
type FileTransferFrame,
} from "@getpaseo/protocol/binary-frames/index";
import { FileUploadStore } from "./index.js";
const tempDirs: string[] = [];
describe("file uploads", () => {
afterEach(() => {
vi.useRealTimers();
for (const dir of tempDirs.splice(0)) {
rmSync(dir, { recursive: true, force: true });
}
});
it("stores chunked upload bytes and returns an uploaded-file attachment", async () => {
const paseoHome = makePaseoHome();
const uploads = new FileUploadStore({ paseoHome });
uploads.beginUpload({
type: "file.upload.request",
fileName: "notes.txt",
mimeType: "text/plain",
size: 11,
modifiedAt: "2026-05-02T00:00:00.000Z",
requestId: "req-upload",
});
await expect(uploads.receiveFrame(uploadBegins("req-upload"))).resolves.toBeNull();
await expect(uploads.receiveFrame(uploadChunk("req-upload", "hello"))).resolves.toBeNull();
await expect(uploads.receiveFrame(uploadChunk("req-upload", " world"))).resolves.toBeNull();
const path = join(paseoHome, "uploads", "upload_req-upload", "notes.txt");
await expect(uploads.receiveFrame(uploadEnds("req-upload"))).resolves.toEqual({
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,
},
error: null,
},
});
expect(readFileSync(path, "utf8")).toBe("hello world");
});
it("rejects chunks beyond the declared size and removes the partial file", async () => {
const paseoHome = makePaseoHome();
const uploads = new FileUploadStore({ paseoHome });
uploads.beginUpload({
type: "file.upload.request",
fileName: "notes.txt",
mimeType: "text/plain",
size: 5,
modifiedAt: "2026-05-02T00:00:00.000Z",
requestId: "req-overflow",
});
await expect(uploads.receiveFrame(uploadBegins("req-overflow"))).resolves.toBeNull();
const uploadDir = join(paseoHome, "uploads", "upload_req-overflow");
const path = join(uploadDir, "notes.txt");
await expect(uploads.receiveFrame(uploadChunk("req-overflow", "hello!"))).resolves.toEqual({
type: "file.upload.response",
payload: {
requestId: "req-overflow",
file: null,
error: "Upload exceeded declared size: expected 5, received 6.",
},
});
expect(existsSync(path)).toBe(false);
expect(existsSync(uploadDir)).toBe(false);
});
it("preserves chunk order when frames arrive before earlier disk writes finish", async () => {
const paseoHome = makePaseoHome();
const uploads = new FileUploadStore({ paseoHome });
uploads.beginUpload({
type: "file.upload.request",
fileName: "notes.txt",
mimeType: "text/plain",
size: 11,
modifiedAt: "2026-05-02T00:00:00.000Z",
requestId: "req-queued",
});
const results = await Promise.all([
uploads.receiveFrame(uploadBegins("req-queued")),
uploads.receiveFrame(uploadChunk("req-queued", "hello")),
uploads.receiveFrame(uploadChunk("req-queued", " world")),
uploads.receiveFrame(uploadEnds("req-queued")),
]);
expect(results.slice(0, 3)).toEqual([null, null, null]);
expect(results[3]?.payload.error).toBeNull();
expect(readFileSync(join(paseoHome, "uploads", "upload_req-queued", "notes.txt"), "utf8")).toBe(
"hello world",
);
});
it("replaces duplicate upload starts without letting the old stale timeout evict the replacement", async () => {
vi.useFakeTimers();
const paseoHome = makePaseoHome();
const uploads = new FileUploadStore({ paseoHome, staleUploadTimeoutMs: 50 });
uploads.beginUpload({
type: "file.upload.request",
fileName: "old.txt",
mimeType: "text/plain",
size: 3,
modifiedAt: "2026-05-02T00:00:00.000Z",
requestId: "req-duplicate",
});
await expect(uploads.receiveFrame(uploadBegins("req-duplicate"))).resolves.toBeNull();
await expect(uploads.receiveFrame(uploadChunk("req-duplicate", "old"))).resolves.toBeNull();
await vi.advanceTimersByTimeAsync(25);
uploads.beginUpload({
type: "file.upload.request",
fileName: "new.txt",
mimeType: "text/plain",
size: 3,
modifiedAt: "2026-05-02T00:00:00.000Z",
requestId: "req-duplicate",
});
await vi.advanceTimersByTimeAsync(30);
const path = join(paseoHome, "uploads", "upload_req-duplicate_2", "new.txt");
await expect(uploads.receiveFrame(uploadBegins("req-duplicate"))).resolves.toBeNull();
await expect(uploads.receiveFrame(uploadChunk("req-duplicate", "new"))).resolves.toBeNull();
await expect(uploads.receiveFrame(uploadEnds("req-duplicate"))).resolves.toEqual({
type: "file.upload.response",
payload: {
requestId: "req-duplicate",
file: {
type: "uploaded_file",
id: "upload_req-duplicate_2",
fileName: "new.txt",
mimeType: "text/plain",
size: 3,
path,
},
error: null,
},
});
expect(readFileSync(path, "utf8")).toBe("new");
});
it("keeps an active upload alive beyond the initial stale timeout", async () => {
vi.useFakeTimers();
const paseoHome = makePaseoHome();
const uploads = new FileUploadStore({ paseoHome, staleUploadTimeoutMs: 50 });
uploads.beginUpload({
type: "file.upload.request",
fileName: "notes.txt",
mimeType: "text/plain",
size: 11,
modifiedAt: "2026-05-02T00:00:00.000Z",
requestId: "req-slow-active",
});
await vi.advanceTimersByTimeAsync(25);
await expect(uploads.receiveFrame(uploadBegins("req-slow-active"))).resolves.toBeNull();
await vi.advanceTimersByTimeAsync(30);
await expect(uploads.receiveFrame(uploadChunk("req-slow-active", "hello"))).resolves.toBeNull();
await vi.advanceTimersByTimeAsync(30);
await expect(
uploads.receiveFrame(uploadChunk("req-slow-active", " world")),
).resolves.toBeNull();
const path = join(paseoHome, "uploads", "upload_req-slow-active", "notes.txt");
await expect(uploads.receiveFrame(uploadEnds("req-slow-active"))).resolves.toEqual({
type: "file.upload.response",
payload: {
requestId: "req-slow-active",
file: {
type: "uploaded_file",
id: "upload_req-slow-active",
fileName: "notes.txt",
mimeType: "text/plain",
size: 11,
path,
},
error: null,
},
});
expect(readFileSync(path, "utf8")).toBe("hello world");
});
});
function makePaseoHome(): string {
const root = realpathSync(mkdtempSync(join(tmpdir(), "file-upload-test-")));
tempDirs.push(root);
return root;
}
function uploadBegins(requestId: string): FileTransferFrame {
return decodeUploadFrame(
encodeFileTransferFrame({
opcode: FileTransferOpcode.FileBegin,
requestId,
metadata: {
mime: "text/plain",
size: 11,
encoding: "binary",
modifiedAt: "2026-05-02T00:00:00.000Z",
fileName: "notes.txt",
},
}),
);
}
function uploadChunk(requestId: string, text: string): FileTransferFrame {
return decodeUploadFrame(
encodeFileTransferFrame({
opcode: FileTransferOpcode.FileChunk,
requestId,
payload: new TextEncoder().encode(text),
}),
);
}
function uploadEnds(requestId: string): FileTransferFrame {
return decodeUploadFrame(
encodeFileTransferFrame({
opcode: FileTransferOpcode.FileEnd,
requestId,
}),
);
}
function decodeUploadFrame(bytes: Uint8Array): FileTransferFrame {
const frame = decodeFileTransferFrame(bytes);
if (!frame) {
throw new Error("Expected file transfer frame");
}
return frame;
}
@@ -0,0 +1,220 @@
import { appendFile, mkdir, rm, writeFile } from "node:fs/promises";
import { basename, join } from "node:path";
import { FileTransferOpcode, type FileTransferFrame } from "@getpaseo/protocol/binary-frames/index";
import { getErrorMessage } from "@getpaseo/protocol/error-utils";
import type { FileUploadRequest, FileUploadResponse } from "../messages.js";
interface FileUploadStoreOptions {
paseoHome: string;
staleUploadTimeoutMs?: number;
}
interface PendingUpload {
requestId: string;
id: string;
attempt: number;
fileName: string;
mimeType: string;
size: number;
path: string;
receivedBytes: number;
started: boolean;
staleTimeout: ReturnType<typeof setTimeout>;
queue: Promise<void>;
}
export class FileUploadStore {
private static readonly defaultStaleUploadTimeoutMs = 10 * 60 * 1000;
private readonly paseoHome: string;
private readonly staleUploadTimeoutMs: number;
private readonly pending = new Map<string, PendingUpload>();
constructor(options: FileUploadStoreOptions) {
this.paseoHome = options.paseoHome;
this.staleUploadTimeoutMs =
options.staleUploadTimeoutMs ?? FileUploadStore.defaultStaleUploadTimeoutMs;
}
beginUpload(request: FileUploadRequest): void {
const existingUpload = this.pending.get(request.requestId);
if (existingUpload) {
this.clearPendingUpload(existingUpload);
void existingUpload.queue.then(() => this.removeUploadDirectory(existingUpload));
}
const fileName = sanitizeFileName(request.fileName);
const attempt = existingUpload ? existingUpload.attempt + 1 : 1;
const id = buildUploadId(request.requestId, attempt);
const uploadDir = join(this.paseoHome, "uploads", id);
const upload: PendingUpload = {
requestId: request.requestId,
id,
attempt,
fileName,
mimeType: request.mimeType,
size: request.size,
path: join(uploadDir, fileName),
receivedBytes: 0,
started: false,
staleTimeout: this.createStaleUploadTimeout(request.requestId),
queue: Promise.resolve(),
};
this.pending.set(request.requestId, upload);
}
async receiveFrame(frame: FileTransferFrame): Promise<FileUploadResponse | null> {
const upload = this.pending.get(frame.requestId);
if (!upload) {
return null;
}
this.refreshStaleUploadTimeout(upload);
const operation = upload.queue.then(() => this.applyFrame(upload, frame));
upload.queue = operation.then(
() => undefined,
() => undefined,
);
return operation;
}
private async applyFrame(
upload: PendingUpload,
frame: FileTransferFrame,
): Promise<FileUploadResponse | null> {
if (this.pending.get(upload.requestId) !== upload) {
return null;
}
try {
if (frame.opcode === FileTransferOpcode.FileBegin) {
await this.startWriting(upload);
return null;
}
if (frame.opcode === FileTransferOpcode.FileChunk) {
await this.writeChunk(upload, frame.payload);
return null;
}
return await this.completeUpload(upload);
} catch (error) {
await this.removeFailedUpload(upload);
return buildUploadResponse(upload, getErrorMessage(error));
}
}
private async startWriting(upload: PendingUpload): Promise<void> {
await mkdir(join(this.paseoHome, "uploads", upload.id), { recursive: true });
await writeFile(upload.path, new Uint8Array());
upload.started = true;
}
private async writeChunk(upload: PendingUpload, bytes: Uint8Array): Promise<void> {
if (!upload.started) {
throw new Error("Upload chunks arrived before file begin.");
}
const nextReceivedBytes = upload.receivedBytes + bytes.byteLength;
if (nextReceivedBytes > upload.size) {
throw new Error(
`Upload exceeded declared size: expected ${upload.size}, received ${nextReceivedBytes}.`,
);
}
await appendFile(upload.path, bytes);
upload.receivedBytes += bytes.byteLength;
}
private async completeUpload(upload: PendingUpload): Promise<FileUploadResponse> {
this.clearPendingUpload(upload);
if (upload.receivedBytes !== upload.size) {
await this.removeUploadDirectory(upload);
return buildUploadResponse(
upload,
`Upload size mismatch: expected ${upload.size}, received ${upload.receivedBytes}.`,
);
}
return buildUploadResponse(upload, null);
}
private createStaleUploadTimeout(requestId: string): ReturnType<typeof setTimeout> {
const timeout = setTimeout(() => {
this.expireStaleUpload(requestId);
}, this.staleUploadTimeoutMs);
timeout.unref?.();
return timeout;
}
private refreshStaleUploadTimeout(upload: PendingUpload): void {
clearTimeout(upload.staleTimeout);
upload.staleTimeout = this.createStaleUploadTimeout(upload.requestId);
}
private expireStaleUpload(requestId: string): void {
const upload = this.pending.get(requestId);
if (!upload) {
return;
}
this.clearPendingUpload(upload);
const cleanup = upload.queue.then(
() => this.removeUploadDirectory(upload),
() => this.removeUploadDirectory(upload),
);
upload.queue = cleanup.then(
() => undefined,
() => undefined,
);
}
private clearPendingUpload(upload: PendingUpload): void {
clearTimeout(upload.staleTimeout);
if (this.pending.get(upload.requestId) === upload) {
this.pending.delete(upload.requestId);
}
}
private async removeFailedUpload(upload: PendingUpload): Promise<void> {
this.clearPendingUpload(upload);
await this.removeUploadDirectory(upload);
}
private async removeUploadDirectory(upload: PendingUpload): Promise<void> {
await rm(join(this.paseoHome, "uploads", upload.id), { recursive: true, force: true }).catch(
() => undefined,
);
}
}
function buildUploadResponse(upload: PendingUpload, error: string | null): FileUploadResponse {
return {
type: "file.upload.response",
payload: {
requestId: upload.requestId,
file: error
? null
: {
type: "uploaded_file",
id: upload.id,
fileName: upload.fileName,
mimeType: upload.mimeType,
size: upload.size,
path: upload.path,
},
error,
},
};
}
function sanitizeUploadId(value: string): string {
return value.replace(/[^a-zA-Z0-9._-]/g, "_") || "file";
}
function buildUploadId(requestId: string, attempt: number): string {
const baseId = `upload_${sanitizeUploadId(requestId)}`;
return attempt === 1 ? baseId : `${baseId}_${attempt}`;
}
function sanitizeFileName(value: string): string {
const name = basename(value)
.replace(/[^a-zA-Z0-9._ -]/g, "_")
.trim();
return name.length > 0 && name !== "." && name !== ".." ? name : "upload";
}
+2 -1
View File
@@ -236,6 +236,7 @@ interface SessionForTestOptions {
providerSnapshotManager?: ProviderSnapshotManager;
stt?: SessionOptions["stt"];
voice?: SessionOptions["voice"];
paseoHome?: string;
messages?: unknown[];
binaryMessages?: Uint8Array[];
}
@@ -272,7 +273,7 @@ function createSessionForTest(options: SessionForTestOptions = {}): Session {
logger,
downloadTokenStore: asDownloadTokenStore(),
pushTokenStore: asPushTokenStore(),
paseoHome: "/tmp/paseo-home",
paseoHome: options.paseoHome ?? "/tmp/paseo-home",
agentManager: asAgentManager({
listAgents: vi.fn(() => []),
subscribe: vi.fn(() => () => {}),
+26 -3
View File
@@ -16,6 +16,7 @@ import {
type SessionOutboundMessage,
type FileExplorerRequest,
type FileDownloadTokenRequest,
type FileUploadRequest,
type GitSetupOptions,
type CheckoutRenameBranchRequest,
type StartWorkspaceScriptRequest,
@@ -30,10 +31,12 @@ import {
import type { TerminalManager } from "../terminal/terminal-manager.js";
import { TerminalSessionController } from "../terminal/terminal-session-controller.js";
import {
type BinaryFrame,
encodeFileTransferFrame,
FileTransferOpcode,
type TerminalStreamFrame,
type FileTransferFrame,
} from "@getpaseo/protocol/binary-frames/index";
import { FileUploadStore } from "./file-upload/index.js";
import { CursorError } from "./pagination/cursor.js";
import { SortablePager, type SortSpec } from "./pagination/sortable-pager.js";
import { TTSManager } from "./agent/tts-manager.js";
@@ -809,6 +812,7 @@ export class Session {
private readonly workspaceSetupSnapshots: Map<string, WorkspaceSetupSnapshot>;
private readonly workspaceGitFetchSubscriptions = new Map<string, () => void>();
private readonly workspaceGitSubscriptions = new Map<string, () => void>();
private readonly fileUploads: FileUploadStore;
private readonly workspaceDirectory: WorkspaceDirectory;
private registerVoiceSpeakHandler?: (agentId: string, handler: VoiceSpeakHandler) => void;
private unregisterVoiceSpeakHandler?: (agentId: string) => void;
@@ -877,6 +881,7 @@ export class Session {
this.onLifecycleIntent = onLifecycleIntent ?? null;
this.downloadTokenStore = downloadTokenStore;
this.pushTokenStore = pushTokenStore;
this.fileUploads = new FileUploadStore({ paseoHome });
this.paseoHome = paseoHome;
this.worktreesRoot = worktreesRoot;
this.sessionLogger = logger.child({
@@ -2108,6 +2113,9 @@ export class Session {
return this.handleProjectIconRequest(msg);
case "file_download_token_request":
return this.handleFileDownloadTokenRequest(msg);
case "file.upload.request":
this.handleFileUploadRequest(msg);
return undefined;
default:
return undefined;
}
@@ -2212,8 +2220,12 @@ export class Session {
this.peakInflightRequests = this.inflightRequests;
}
public handleBinaryFrame(frame: TerminalStreamFrame): void {
this.terminalController.handleBinaryFrame(frame);
public async handleBinaryFrame(binaryFrame: BinaryFrame): Promise<void> {
if (binaryFrame.kind === "file_transfer") {
await this.handleFileTransferFrame(binaryFrame.frame);
return;
}
this.terminalController.handleBinaryFrame(binaryFrame.frame);
}
private async handleRestartServerRequest(requestId: string, reason?: string): Promise<void> {
@@ -5823,6 +5835,17 @@ export class Session {
}
}
private handleFileUploadRequest(request: FileUploadRequest): void {
this.fileUploads.beginUpload(request);
}
private async handleFileTransferFrame(frame: FileTransferFrame): Promise<void> {
const response = await this.fileUploads.receiveFrame(frame);
if (response) {
this.emit(response);
}
}
/**
* Handle project icon request for a given cwd
*/
@@ -105,6 +105,7 @@ import type { SpeechReadinessSnapshot } from "./speech/speech-runtime.js";
interface WebSocketServerInternals {
attachSocket(ws: unknown, req: unknown): Promise<void>;
socketMessageQueues: Map<unknown, Promise<void>>;
}
const TEST_DAEMON_VERSION = "1.2.3-test";
@@ -125,9 +126,12 @@ function parseSentEnvelope(data: unknown): z.infer<typeof WireEnvelopeSchema> {
}
const BinaryFrameSchema = z.object({
opcode: z.number(),
slot: z.number(),
payload: z.instanceof(Uint8Array),
kind: z.literal("terminal"),
frame: z.object({
opcode: z.number(),
slot: z.number(),
payload: z.instanceof(Uint8Array),
}),
});
class MockSocket {
@@ -369,7 +373,7 @@ async function attachRelayAndHello(params: {
}) {
await params.server.attachExternalSocket(params.socket, { transport: "relay" });
params.socket.emit("message", JSON.stringify(createHelloMessage(params.clientId)));
await Promise.resolve();
await waitForSocketMessages(params.server, params.socket);
expect(params.socket.sent.length).toBeGreaterThan(0);
const envelope = parseSentEnvelope(params.socket.sent[0]);
expect(envelope.type).toBe("session");
@@ -389,7 +393,7 @@ async function attachDirectAndHello(params: {
createDirectRequest(),
);
params.socket.emit("message", JSON.stringify(createHelloMessage(params.clientId)));
await Promise.resolve();
await waitForSocketMessages(params.server, params.socket);
expect(params.socket.sent.length).toBeGreaterThan(0);
const envelope = parseSentEnvelope(params.socket.sent[0]);
expect(envelope.type).toBe("session");
@@ -399,6 +403,13 @@ async function attachDirectAndHello(params: {
return serverInfo!;
}
async function waitForSocketMessages(
server: VoiceAssistantWebSocketServer,
socket: MockSocket,
): Promise<void> {
await asInternals<WebSocketServerInternals>(server).socketMessageQueues.get(socket);
}
describe("relay external socket reconnect behavior", () => {
beforeEach(() => {
sessionMock.instances.length = 0;
@@ -453,7 +464,7 @@ describe("relay external socket reconnect behavior", () => {
}),
),
);
await Promise.resolve();
await waitForSocketMessages(server, socket);
expect(sessionMock.instances).toHaveLength(1);
const session = sessionMock.instances[0];
@@ -550,7 +561,7 @@ describe("relay external socket reconnect behavior", () => {
},
}),
);
await Promise.resolve();
await waitForSocketMessages(server, socket);
expect(closeCode).toBe(4002);
expect(["Invalid hello", "Session message before hello"]).toContain(closeReason);
@@ -757,10 +768,10 @@ describe("relay external socket reconnect behavior", () => {
}),
),
);
await Promise.resolve();
await waitForSocketMessages(server, socket);
expect(session.handleBinaryFrame).toHaveBeenCalledTimes(1);
const frame = BinaryFrameSchema.parse(session.handleBinaryFrame.mock.calls[0]?.[0]);
const { frame } = BinaryFrameSchema.parse(session.handleBinaryFrame.mock.calls[0]?.[0]);
expect(frame.opcode).toBe(TerminalStreamOpcode.Input);
expect(frame.slot).toBe(9);
expect(new TextDecoder().decode(frame.payload)).toBe("ls\r");
+28 -8
View File
@@ -24,7 +24,7 @@ import {
type WSOutboundMessage,
wrapSessionMessage,
} from "./messages.js";
import { asUint8Array, decodeTerminalStreamFrame } from "@getpaseo/protocol/binary-frames/index";
import { asUint8Array, decodeBinaryFrame } from "@getpaseo/protocol/binary-frames/index";
import type { HostnamesConfig } from "./hostnames.js";
import { isHostnameAllowed } from "./hostnames.js";
import { Session, type SessionLifecycleIntent, type SessionRuntimeMetrics } from "./session.js";
@@ -324,6 +324,7 @@ export class VoiceAssistantWebSocketServer {
private readonly pendingConnections: Map<WebSocketLike, PendingConnection> = new Map();
private readonly sessions: Map<WebSocketLike, SessionConnection> = new Map();
private readonly externalSessionsByKey: Map<string, SessionConnection> = new Map();
private readonly socketMessageQueues: Map<WebSocketLike, Promise<void>> = new Map();
private readonly serverId: string;
private readonly daemonVersion: string;
private readonly daemonRuntimeConfig:
@@ -1096,7 +1097,7 @@ export class VoiceAssistantWebSocketServer {
private bindSocketHandlers(ws: WebSocketLike): void {
ws.on("message", (...args: unknown[]) => {
const data = args[0] as Buffer | ArrayBuffer | Buffer[] | string;
void this.handleRawMessage(ws, data);
this.enqueueRawMessage(ws, data);
});
ws.on("close", async (...args: unknown[]) => {
@@ -1119,6 +1120,25 @@ export class VoiceAssistantWebSocketServer {
});
}
private enqueueRawMessage(
ws: WebSocketLike,
data: Buffer | ArrayBuffer | Buffer[] | string,
): void {
const previous = this.socketMessageQueues.get(ws) ?? Promise.resolve();
const next = previous.then(
() => this.handleRawMessage(ws, data),
() => this.handleRawMessage(ws, data),
);
this.socketMessageQueues.set(ws, next);
void next
.catch(() => undefined)
.finally(() => {
if (this.socketMessageQueues.get(ws) === next) {
this.socketMessageQueues.delete(ws);
}
});
}
public resolveVoiceSpeakHandler(callerAgentId: string): VoiceSpeakHandler | null {
return this.voiceSpeakHandlers.get(callerAgentId) ?? null;
}
@@ -1298,19 +1318,19 @@ export class VoiceAssistantWebSocketServer {
);
}
private maybeHandleBinaryFrame(params: {
private async maybeHandleBinaryFrame(params: {
ws: WebSocketLike;
buffer: Buffer;
activeConnection: SessionConnection | undefined;
log: pino.Logger;
}): boolean {
}): Promise<boolean> {
const { ws, buffer, activeConnection, log } = params;
const asBytes = asUint8Array(buffer);
if (!asBytes) {
return false;
}
const frame = decodeTerminalStreamFrame(asBytes);
if (!frame) {
const decodedFrame = decodeBinaryFrame(asBytes);
if (!decodedFrame) {
return false;
}
if (!activeConnection) {
@@ -1324,7 +1344,7 @@ export class VoiceAssistantWebSocketServer {
}
return true;
}
activeConnection.session.handleBinaryFrame(frame);
await activeConnection.session.handleBinaryFrame(decodedFrame);
return true;
}
@@ -1369,7 +1389,7 @@ export class VoiceAssistantWebSocketServer {
try {
const buffer = bufferFromWsData(data);
const binaryHandled = this.maybeHandleBinaryFrame({
const binaryHandled = await this.maybeHandleBinaryFrame({
ws,
buffer,
activeConnection,