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:
@@ -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,
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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";
|
||||
}
|
||||
@@ -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(() => () => {}),
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user