From 1cc3d137143f79d3f534c89bbcd7d54b88144c66 Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 17:44:54 +0800 Subject: [PATCH 01/11] feat(schema): add chunked attachment upload and read wire frames Additive kinds at wire 80: begin/chunk/commit/abort/read, 256 KiB chunks, and a feature-detect constant. The payload field is attachmentKind so it does not collide with the frame discriminator. --- .../foundation/schema/src/model/attachment.ts | 7 +- .../foundation/schema/src/wire/attachment.ts | 99 +++++++++++ packages/foundation/schema/src/wire/index.ts | 8 + .../foundation/schema/src/wire/payload.ts | 2 + .../tests/contract/wire/attachment.test.ts | 161 ++++++++++++++++++ 5 files changed, 275 insertions(+), 2 deletions(-) create mode 100644 packages/foundation/schema/src/wire/attachment.ts create mode 100644 packages/foundation/schema/tests/contract/wire/attachment.test.ts diff --git a/packages/foundation/schema/src/model/attachment.ts b/packages/foundation/schema/src/model/attachment.ts index 7f40e3bdb..bc328fe9e 100644 --- a/packages/foundation/schema/src/model/attachment.ts +++ b/packages/foundation/schema/src/model/attachment.ts @@ -23,8 +23,11 @@ export function blobIdFromSha256(hex: string): BlobId { return BlobIdSchema.parse(`sha256:${hex.toLowerCase()}`); } -/** Upload ID: daemon-minted identity of one in-flight upload lease. */ -export const UploadIdSchema = z.string().min(1).brand<'UploadId'>(); +const rUploadId = /^[\w-]{1,128}$/; + +/** Upload ID: daemon-minted identity of one in-flight upload lease. The charset is the staging + * filename — anything else is a path traversal. */ +export const UploadIdSchema = z.string().regex(rUploadId).brand<'UploadId'>(); export type UploadId = z.infer; export const BlobRecordSchema = z.object({ diff --git a/packages/foundation/schema/src/wire/attachment.ts b/packages/foundation/schema/src/wire/attachment.ts new file mode 100644 index 000000000..e2b73d85e --- /dev/null +++ b/packages/foundation/schema/src/wire/attachment.ts @@ -0,0 +1,99 @@ +import { z } from 'zod'; +import { + AttachmentKindSchema, + BlobIdSchema, + Sha256HexSchema, + UploadIdSchema, +} from '../model/attachment'; +import { MAX_ATTACHMENT_BYTES } from '../model/content'; +import { AttachmentIdSchema, OperationIdSchema, SessionIdSchema } from '../model/primitives'; +import { WireRequestIdSchema } from './request'; + +/** Raw bytes per upload/read chunk. One frame stays under the tunnel's 768 KiB chunk and + * workerd's 1 MiB cap after base64 (~341 KiB). */ +export const ATTACHMENT_UPLOAD_CHUNK_BYTES = 256 * 1024; + +/** Unacked chunks a client may have in flight; acks are cumulative. */ +export const ATTACHMENT_UPLOAD_WINDOW_CHUNKS = 2; + +/** `ATTACHMENT_UPLOAD_CHUNK_BYTES` as base64 — the per-frame `data` cap. */ +export const ATTACHMENT_UPLOAD_CHUNK_BASE64_MAX = 4 * Math.ceil(ATTACHMENT_UPLOAD_CHUNK_BYTES / 3); + +/** The wire version that introduced chunked attachment upload/read. Clients feature-detect on it. */ +export const ATTACHMENT_STORE_WIRE_VERSION = 80 as const; + +export const AttachmentUploadStateSchema = z.enum(['ready', 'exists']); +export type AttachmentUploadState = z.infer; + +/** Chunked attachment upload and read. Bytes travel as base64; `resource.source.upload` stays for + * older peers. Draft leases are not bound to a session — a later prompt or resource claims them. */ +export const attachmentWireVariants = [ + z.object({ + kind: z.literal('attachment.upload.begin'), + clientReqId: WireRequestIdSchema, + /** Replay key for a lost begin reply; a second begin with the same id returns the first. */ + operationId: OperationIdSchema.optional(), + declaredSha256: Sha256HexSchema, + declaredSize: z.number().int().nonnegative().max(MAX_ATTACHMENT_BYTES), + name: z.string().min(1), + mimeType: z.string().min(1).optional(), + /** AttachmentRecord.kind — not the frame discriminator. */ + attachmentKind: AttachmentKindSchema, + }), + z.object({ + kind: z.literal('attachment.upload.begun'), + replyTo: WireRequestIdSchema, + uploadId: UploadIdSchema, + chunkBytes: z.number().int().positive(), + state: AttachmentUploadStateSchema, + }), + z.object({ + kind: z.literal('attachment.upload.chunk'), + clientReqId: WireRequestIdSchema, + uploadId: UploadIdSchema, + offset: z.number().int().nonnegative(), + data: z.string().max(ATTACHMENT_UPLOAD_CHUNK_BASE64_MAX), + }), + z.object({ + kind: z.literal('attachment.upload.chunk.acked'), + replyTo: WireRequestIdSchema, + uploadId: UploadIdSchema, + /** Contiguous prefix received so far; the next chunk must start here. */ + receivedBytes: z.number().int().nonnegative(), + }), + z.object({ + kind: z.literal('attachment.upload.commit'), + clientReqId: WireRequestIdSchema, + uploadId: UploadIdSchema, + }), + z.object({ + kind: z.literal('attachment.upload.committed'), + replyTo: WireRequestIdSchema, + attachmentId: AttachmentIdSchema, + blobId: BlobIdSchema, + }), + z.object({ + kind: z.literal('attachment.upload.abort'), + clientReqId: WireRequestIdSchema, + uploadId: UploadIdSchema, + }), + z.object({ + kind: z.literal('attachment.read'), + clientReqId: WireRequestIdSchema, + sessionId: SessionIdSchema, + attachmentId: AttachmentIdSchema, + offset: z.number().int().nonnegative(), + length: z.number().int().positive().max(ATTACHMENT_UPLOAD_CHUNK_BYTES), + }), + z.object({ + kind: z.literal('attachment.read.result'), + replyTo: WireRequestIdSchema, + sessionId: SessionIdSchema, + attachmentId: AttachmentIdSchema, + blobId: BlobIdSchema, + offset: z.number().int().nonnegative(), + data: z.string().max(ATTACHMENT_UPLOAD_CHUNK_BASE64_MAX), + sizeBytes: z.number().int().nonnegative(), + eof: z.boolean(), + }), +] as const; diff --git a/packages/foundation/schema/src/wire/index.ts b/packages/foundation/schema/src/wire/index.ts index 784aaa437..1b4b41873 100644 --- a/packages/foundation/schema/src/wire/index.ts +++ b/packages/foundation/schema/src/wire/index.ts @@ -1,3 +1,11 @@ +export { + ATTACHMENT_STORE_WIRE_VERSION, + ATTACHMENT_UPLOAD_CHUNK_BASE64_MAX, + ATTACHMENT_UPLOAD_CHUNK_BYTES, + ATTACHMENT_UPLOAD_WINDOW_CHUNKS, + type AttachmentUploadState, + AttachmentUploadStateSchema, +} from './attachment'; export { CONVERSATION_GRAPH_WIRE_VERSION, type ConversationEvent, diff --git a/packages/foundation/schema/src/wire/payload.ts b/packages/foundation/schema/src/wire/payload.ts index 641872ec9..9c2350fb6 100644 --- a/packages/foundation/schema/src/wire/payload.ts +++ b/packages/foundation/schema/src/wire/payload.ts @@ -4,6 +4,7 @@ import { agentCatalogWireVariants } from './agent-catalog'; import { agentLoginWireVariants } from './agent-login'; import { agentRuntimeWireVariants } from './agent-runtime'; import { artifactWireVariants } from './artifact'; +import { attachmentWireVariants } from './attachment'; import { browserWireVariants } from './browser'; import { configWireVariants } from './config'; import { conversationWireVariants } from './conversation'; @@ -28,6 +29,7 @@ const wirePayloadVariants = [ ...conversationWireVariants, ...historyWireVariants, ...requestWireVariants, + ...attachmentWireVariants, ...resourceWireVariants, ...configWireVariants, ...agentRuntimeWireVariants, diff --git a/packages/foundation/schema/tests/contract/wire/attachment.test.ts b/packages/foundation/schema/tests/contract/wire/attachment.test.ts new file mode 100644 index 000000000..45a907f65 --- /dev/null +++ b/packages/foundation/schema/tests/contract/wire/attachment.test.ts @@ -0,0 +1,161 @@ +import { + ATTACHMENT_UPLOAD_CHUNK_BASE64_MAX, + ATTACHMENT_UPLOAD_CHUNK_BYTES, + MAX_ATTACHMENT_BYTES, + WIRE_PROTOCOL_VERSION, + WireMessageSchema, +} from '@linkcode/schema'; +import { describe, expect, it } from 'vitest'; + +function parses(payload: unknown): boolean { + return WireMessageSchema.safeParse({ + v: WIRE_PROTOCOL_VERSION, + id: 'message-1', + ts: 0, + payload, + }).success; +} + +const sha256 = 'a'.repeat(64); +const blobId = `sha256:${sha256}`; + +describe('attachment upload/read frames', () => { + it('round-trips begin, chunk, commit, abort, and read', () => { + expect( + parses({ + kind: 'attachment.upload.begin', + clientReqId: 'request-1', + declaredSha256: sha256, + declaredSize: 16, + name: 'shot.png', + mimeType: 'image/png', + attachmentKind: 'image', + }), + ).toBe(true); + expect( + parses({ + kind: 'attachment.upload.begun', + replyTo: 'request-1', + uploadId: 'upl-1', + chunkBytes: ATTACHMENT_UPLOAD_CHUNK_BYTES, + state: 'ready', + }), + ).toBe(true); + expect( + parses({ + kind: 'attachment.upload.chunk', + clientReqId: 'request-2', + uploadId: 'upl-1', + offset: 0, + data: 'aGVsbG8=', + }), + ).toBe(true); + expect( + parses({ + kind: 'attachment.upload.chunk.acked', + replyTo: 'request-2', + uploadId: 'upl-1', + receivedBytes: 5, + }), + ).toBe(true); + expect( + parses({ + kind: 'attachment.upload.commit', + clientReqId: 'request-3', + uploadId: 'upl-1', + }), + ).toBe(true); + expect( + parses({ + kind: 'attachment.upload.committed', + replyTo: 'request-3', + attachmentId: 'att-1', + blobId, + }), + ).toBe(true); + expect( + parses({ + kind: 'attachment.upload.abort', + clientReqId: 'request-4', + uploadId: 'upl-1', + }), + ).toBe(true); + expect( + parses({ + kind: 'attachment.read', + clientReqId: 'request-5', + sessionId: 'session-1', + attachmentId: 'att-1', + offset: 0, + length: 16, + }), + ).toBe(true); + expect( + parses({ + kind: 'attachment.read.result', + replyTo: 'request-5', + sessionId: 'session-1', + attachmentId: 'att-1', + blobId, + offset: 0, + data: 'aGVsbG8=', + sizeBytes: 5, + eof: true, + }), + ).toBe(true); + }); + + it('rejects a path-shaped upload id, an oversized claim, and an over-budget chunk', () => { + expect( + parses({ + kind: 'attachment.upload.chunk', + clientReqId: 'request-1', + uploadId: '../escape', + offset: 0, + data: 'YQ==', + }), + ).toBe(false); + expect( + parses({ + kind: 'attachment.upload.begin', + clientReqId: 'request-1', + declaredSha256: sha256, + declaredSize: MAX_ATTACHMENT_BYTES + 1, + name: 'too-big.bin', + attachmentKind: 'file', + }), + ).toBe(false); + expect( + parses({ + kind: 'attachment.upload.chunk', + clientReqId: 'request-1', + uploadId: 'upl-1', + offset: 0, + data: 'A'.repeat(ATTACHMENT_UPLOAD_CHUNK_BASE64_MAX + 1), + }), + ).toBe(false); + }); + + it('accepts a begin without operationId and an exists short-circuit', () => { + expect( + parses({ + kind: 'attachment.upload.begin', + clientReqId: 'request-1', + operationId: 'op-1', + declaredSha256: sha256, + declaredSize: 0, + name: 'empty.bin', + attachmentKind: 'file', + }), + ).toBe(true); + expect( + parses({ + kind: 'attachment.upload.begun', + replyTo: 'request-1', + uploadId: 'upl-1', + chunkBytes: ATTACHMENT_UPLOAD_CHUNK_BYTES, + state: 'exists', + }), + ).toBe(true); + }); +}); From c169b370d8ec22eb9e1388dc38d9a15d24c1836b Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 18:06:16 +0800 Subject: [PATCH 02/11] feat(engine,daemon): read blobs by offset and check attachment reachability Positional BlobStore.read, session-scoped isReachable (prompt refs or session resources), and a missing-file stat so a vanished CAS object is not treated as a dedupe hit. --- .../src/__tests__/attachment-store.test.ts | 71 +++++++++++++++++++ apps/daemon/src/attachment-store.ts | 36 +++++++++- .../engine/src/__tests__/blob-store.test.ts | 13 ++++ .../engine/src/attachment/attachment-store.ts | 20 +++++- .../host/engine/src/attachment/blob-store.ts | 23 ++++++ .../src/conversation/conversation-store.ts | 33 +++++++-- packages/host/engine/src/engine.ts | 41 ++++++++--- .../engine/src/resource/resource-store.ts | 10 +++ 8 files changed, 230 insertions(+), 17 deletions(-) diff --git a/apps/daemon/src/__tests__/attachment-store.test.ts b/apps/daemon/src/__tests__/attachment-store.test.ts index b5b876760..8ded5cccf 100644 --- a/apps/daemon/src/__tests__/attachment-store.test.ts +++ b/apps/daemon/src/__tests__/attachment-store.test.ts @@ -11,6 +11,7 @@ import { ConversationOperationSchema, ConversationTurnSchema, PromptRecordSchema, + SessionIdSchema, SessionRecordSchema, SessionResourceSchema, UploadIdSchema, @@ -209,4 +210,74 @@ describe('SQLite attachment store', () => { expect(await store.getBlob(leasedOnly.blobId)).toBeUndefined(); expect(await store.getBlob(viaResource.blobId)).toEqual(viaResource); }); + + it('reports reachability from a session prompt or resource, not another session', async () => { + const { database, path, store } = await fixture(); + await store.commitAttachment({ blob: blob('prompt'), attachment: attachment('att-prompt') }); + await store.commitAttachment({ + blob: blob('resource'), + attachment: attachment('att-resource'), + }); + await createConversationStore(database.client).persistTurnIntent({ + turn: ConversationTurnSchema.parse({ + turnId: 't-1', + sessionId: 's-1', + parentTurnId: null, + siblingOrdinal: 1, + input: { type: 'prompt', promptId: 'p-1' }, + runId: 'run-1', + state: 'preparing', + createdAt: 1, + }), + prompt: PromptRecordSchema.parse({ + promptId: 'p-1', + blocks: [{ type: 'attachment_ref', attachmentId: 'att-prompt' }], + contextAttachmentIds: [], + createdAt: 1, + }), + operation: ConversationOperationSchema.parse({ + operationId: 'op-reach', + sessionId: 's-1', + kind: 'turn.submit', + state: 'open', + createdAt: 1, + }), + }); + await createResourceStore(path).save( + SessionResourceSchema.parse({ + resourceId: 'resource-1', + sessionId: 's-1', + direction: 'source', + name: 'brief.txt', + kind: 'file', + status: 'ready', + locator: { type: 'managed-file', path: '/state/blobs/x' }, + attachmentId: 'att-resource', + createdAt: 1, + updatedAt: 1, + }), + ); + + expect( + await store.isReachable(SessionIdSchema.parse('s-1'), AttachmentIdSchema.parse('att-prompt')), + ).toBe(true); + expect( + await store.isReachable( + SessionIdSchema.parse('s-1'), + AttachmentIdSchema.parse('att-resource'), + ), + ).toBe(true); + expect( + await store.isReachable( + SessionIdSchema.parse('s-other'), + AttachmentIdSchema.parse('att-prompt'), + ), + ).toBe(false); + expect( + await store.isReachable( + SessionIdSchema.parse('s-1'), + AttachmentIdSchema.parse('att-missing'), + ), + ).toBe(false); + }); }); diff --git a/apps/daemon/src/attachment-store.ts b/apps/daemon/src/attachment-store.ts index 0282b8cf2..6382d652d 100644 --- a/apps/daemon/src/attachment-store.ts +++ b/apps/daemon/src/attachment-store.ts @@ -4,7 +4,14 @@ import type { AttachmentSweepWindow, StoredAttachment, } from '@linkcode/engine'; -import type { AttachmentId, BlobId, BlobRecord, UploadId, UploadLease } from '@linkcode/schema'; +import type { + AttachmentId, + BlobId, + BlobRecord, + SessionId, + UploadId, + UploadLease, +} from '@linkcode/schema'; import { AttachmentRecordSchema, BlobIdSchema, @@ -18,6 +25,7 @@ import { attachmentBlobs, attachments, blobs, + conversationTurns, promptAttachmentRefs, sessionResources, uploadLeases, @@ -90,6 +98,32 @@ export function createAttachmentStore(db: DaemonDatabaseClient): AttachmentStore return Promise.resolve(); }, + isReachable(sessionId: SessionId, attachmentId: AttachmentId): Promise { + const fromPrompt = db + .select({ id: promptAttachmentRefs.attachmentId }) + .from(promptAttachmentRefs) + .innerJoin(conversationTurns, eq(conversationTurns.promptId, promptAttachmentRefs.promptId)) + .where( + and( + eq(conversationTurns.sessionId, sessionId), + eq(promptAttachmentRefs.attachmentId, attachmentId), + ), + ) + .get(); + if (fromPrompt) return Promise.resolve(true); + const fromResource = db + .select({ id: sessionResources.attachmentId }) + .from(sessionResources) + .where( + and( + eq(sessionResources.sessionId, sessionId), + eq(sessionResources.attachmentId, attachmentId), + ), + ) + .get(); + return Promise.resolve(fromResource !== undefined); + }, + commitAttachment({ attachment, blob, uploadId }: AttachmentCommit): Promise { db.transaction((tx) => { tx.insert(blobs).values(blob).onConflictDoNothing().run(); diff --git a/packages/host/engine/src/__tests__/blob-store.test.ts b/packages/host/engine/src/__tests__/blob-store.test.ts index d611649b5..63e23bf1e 100644 --- a/packages/host/engine/src/__tests__/blob-store.test.ts +++ b/packages/host/engine/src/__tests__/blob-store.test.ts @@ -94,6 +94,19 @@ describe('FsBlobStore', () => { expect(await store.list()).toEqual([]); await expect(store.delete(blobId)).resolves.toBeUndefined(); }); + + it('reads a committed blob by offset and length', async () => { + const { store } = await storeInTempDir(); + const bytes = Buffer.from('abcdefghij'); + const stage = await store.stage('readable'); + await stage.write(0, bytes); + const blobId = await stage.commit({ sha256: sha256(bytes), sizeBytes: bytes.byteLength }); + + expect(Buffer.from((await store.read(blobId, 0, 4)) ?? [])).toEqual(Buffer.from('abcd')); + expect(Buffer.from((await store.read(blobId, 6, 16)) ?? [])).toEqual(Buffer.from('ghij')); + expect(Buffer.from((await store.read(blobId, 10, 4)) ?? [])).toEqual(Buffer.from('')); + expect(await store.read(blobIdFromSha256('b'.repeat(64)), 0, 4)).toBeUndefined(); + }); }); describe('mime sniff', () => { diff --git a/packages/host/engine/src/attachment/attachment-store.ts b/packages/host/engine/src/attachment/attachment-store.ts index d0ea54873..cf0822a1d 100644 --- a/packages/host/engine/src/attachment/attachment-store.ts +++ b/packages/host/engine/src/attachment/attachment-store.ts @@ -3,11 +3,13 @@ import type { AttachmentRecord, BlobId, BlobRecord, + SessionId, Timestamp, UploadId, UploadLease, } from '@linkcode/schema'; import { blobIdFromSha256 } from '@linkcode/schema'; +import { falseFn } from 'foxts/noop'; /** An attachment record joined with the blob holding its `original` bytes. */ export interface StoredAttachment extends AttachmentRecord { @@ -46,18 +48,28 @@ export interface AttachmentStore { /** Atomic: the blob row (if new), the attachment with its `original` variant, and the lease * pointed at the attachment. */ commitAttachment(commit: AttachmentCommit): Promise; + /** Whether a prompt of a turn in `sessionId`, or a session resource of that session, names the + * attachment. Integrity, not confidentiality — every peer of this store is one account. */ + isReachable(sessionId: SessionId, attachmentId: AttachmentId): Promise; /** Atomic reaper: expired leases first; then attachments, then blobs, that nothing roots and * that predate the grace window. Returns the blob ids whose bytes the caller must delete. */ sweep(window: AttachmentSweepWindow): Promise; } +/** Session-scoped root check for the in-memory store. */ +export type AttachmentReachability = (sessionId: SessionId, attachmentId: AttachmentId) => boolean; + export class InMemoryAttachmentStore implements AttachmentStore { private readonly blobs = new Map(); private readonly attachments = new Map(); private readonly leases = new Map(); - /** `roots` lists every attachment id a prompt or session resource currently references. */ - constructor(private readonly roots: () => Iterable = () => []) {} + /** `roots` lists every attachment id a prompt or session resource currently references. + * `reachable` is the same set sliced by session — used by `attachment.read`. */ + constructor( + private readonly roots: () => Iterable = () => [], + private readonly reachable: AttachmentReachability = falseFn, + ) {} getAttachment(attachmentId: AttachmentId): Promise { const attachment = this.attachments.get(attachmentId); @@ -98,6 +110,10 @@ export class InMemoryAttachmentStore implements AttachmentStore { return Promise.resolve(); } + isReachable(sessionId: SessionId, attachmentId: AttachmentId): Promise { + return Promise.resolve(this.reachable(sessionId, attachmentId)); + } + commitAttachment({ attachment, blob, uploadId }: AttachmentCommit): Promise { if (!this.blobs.has(blob.blobId)) this.blobs.set(blob.blobId, structuredClone(blob)); this.attachments.set(attachment.attachmentId, { diff --git a/packages/host/engine/src/attachment/blob-store.ts b/packages/host/engine/src/attachment/blob-store.ts index f751b4f93..a11e41116 100644 --- a/packages/host/engine/src/attachment/blob-store.ts +++ b/packages/host/engine/src/attachment/blob-store.ts @@ -28,6 +28,8 @@ export interface BlobStore { /** Absolute path for same-host consumers (materialization, hosted files); existence not implied. */ pathOf(blobId: BlobId): string; stat(blobId: BlobId): Promise<{ sizeBytes: number } | undefined>; + /** Positional read of an immutable blob. Missing file → `undefined`; a short read at EOF is ok. */ + read(blobId: BlobId, offset: number, length: number): Promise; /** Open staging for one upload; readers cannot observe the bytes until `commit`. */ stage(uploadId: string): Promise; delete(blobId: BlobId): Promise; @@ -108,6 +110,27 @@ export class FsBlobStore implements BlobStore { } } + async read(blobId: BlobId, offset: number, length: number): Promise { + if (offset < 0 || length < 0) throw new Error('Blob read offset and length must be >= 0'); + let handle: FileHandle; + try { + handle = await open(this.pathOf(blobId), 'r'); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') return undefined; + throw error; + } + try { + const sizeBytes = (await handle.stat()).size; + if (length === 0 || offset >= sizeBytes) return new Uint8Array(0); + const toRead = Math.min(length, sizeBytes - offset); + const buffer = Buffer.alloc(toRead); + const { bytesRead } = await handle.read(buffer, 0, toRead, offset); + return bytesRead === toRead ? buffer : buffer.subarray(0, bytesRead); + } finally { + await handle.close(); + } + } + async stage(uploadId: string): Promise { if (!rUploadId.test(uploadId)) throw new Error(`Invalid upload id: ${uploadId}`); const dir = join(this.root, 'tmp'); diff --git a/packages/host/engine/src/conversation/conversation-store.ts b/packages/host/engine/src/conversation/conversation-store.ts index d70cce86d..054310808 100644 --- a/packages/host/engine/src/conversation/conversation-store.ts +++ b/packages/host/engine/src/conversation/conversation-store.ts @@ -70,14 +70,25 @@ export class InMemoryConversationStore implements ConversationStore { referencedAttachmentIds(): AttachmentId[] { const ids: AttachmentId[] = []; for (const prompt of this.prompts.values()) { - for (let i = 0, len = prompt.contextAttachmentIds.length; i < len; i++) { - ids.push(prompt.contextAttachmentIds[i]); - } - for (let i = 0, len = prompt.blocks.length; i < len; i++) { - const block = prompt.blocks[i]; - if (block.type === 'attachment_ref') ids.push(block.attachmentId); + collectPromptAttachmentIds(prompt, ids); + } + return ids; + } + + /** Attachments a turn in this session references — the in-memory `attachment.read` reachability. */ + referencedAttachmentIdsForSession(sessionId: SessionId): AttachmentId[] { + const promptIds = new Set(); + for (const turn of this.turns.values()) { + if (turn.sessionId !== sessionId) continue; + if (turn.input.type === 'prompt' && turn.input.promptId !== null) { + promptIds.add(turn.input.promptId); } } + const ids: AttachmentId[] = []; + for (const promptId of promptIds) { + const prompt = this.prompts.get(promptId); + if (prompt) collectPromptAttachmentIds(prompt, ids); + } return ids; } @@ -195,3 +206,13 @@ export class InMemoryConversationStore implements ConversationStore { return Promise.resolve(); } } + +function collectPromptAttachmentIds(prompt: PromptRecord, ids: AttachmentId[]): void { + for (let i = 0, len = prompt.contextAttachmentIds.length; i < len; i++) { + ids.push(prompt.contextAttachmentIds[i]); + } + for (let i = 0, len = prompt.blocks.length; i < len; i++) { + const block = prompt.blocks[i]; + if (block.type === 'attachment_ref') ids.push(block.attachmentId); + } +} diff --git a/packages/host/engine/src/engine.ts b/packages/host/engine/src/engine.ts index 6dc387d9c..4daecb199 100644 --- a/packages/host/engine/src/engine.ts +++ b/packages/host/engine/src/engine.ts @@ -16,6 +16,7 @@ import { InMemoryProviderConfigStore } from './agent/provider-config'; import { AgentRequestHandler } from './agent/request-handler'; import { AgentRuntimeService } from './agent/runtime-service'; import { ManagedAssetService } from './asset/service'; +import type { AttachmentReachability } from './attachment/attachment-store'; import { InMemoryAttachmentStore } from './attachment/attachment-store'; import { FsBlobStore } from './attachment/blob-store'; import { AttachmentGc } from './attachment/gc'; @@ -121,14 +122,17 @@ export const createEngineRuntime = Effect.fn('Engine.create')(function* ( const blobStore = deps.blobStore ?? new FsBlobStore(join(stateDir, 'blobs')); const attachmentStore = deps.attachmentStore ?? - new InMemoryAttachmentStore(() => [ - ...(conversationStore instanceof InMemoryConversationStore - ? conversationStore.referencedAttachmentIds() - : []), - ...(resourceStore instanceof InMemoryResourceStore - ? resourceStore.referencedAttachmentIds() - : []), - ]); + new InMemoryAttachmentStore( + () => [ + ...(conversationStore instanceof InMemoryConversationStore + ? conversationStore.referencedAttachmentIds() + : []), + ...(resourceStore instanceof InMemoryResourceStore + ? resourceStore.referencedAttachmentIds() + : []), + ], + inMemoryAttachmentReachability(conversationStore, resourceStore), + ); const attachmentIo = new AttachmentIoMutex(); const attachmentGc = new AttachmentGc(attachmentStore, blobStore, Date.now, attachmentIo); const resources = new ResourceService( @@ -477,6 +481,27 @@ export const createEngineRuntime = Effect.fn('Engine.create')(function* ( }; }); +function inMemoryAttachmentReachability( + conversations: InMemoryConversationStore | object, + resources: InMemoryResourceStore | object, +): AttachmentReachability { + return (sessionId, attachmentId) => { + if (conversations instanceof InMemoryConversationStore) { + const ids = conversations.referencedAttachmentIdsForSession(sessionId); + for (let i = 0, len = ids.length; i < len; i++) { + if (ids[i] === attachmentId) return true; + } + } + if (resources instanceof InMemoryResourceStore) { + const ids = resources.referencedAttachmentIdsForSession(sessionId); + for (let i = 0, len = ids.length; i < len; i++) { + if (ids[i] === attachmentId) return true; + } + } + return false; + }; +} + function tryOperation( subsystem: OperationSubsystem, operation: string, diff --git a/packages/host/engine/src/resource/resource-store.ts b/packages/host/engine/src/resource/resource-store.ts index c05a41da1..387696aef 100644 --- a/packages/host/engine/src/resource/resource-store.ts +++ b/packages/host/engine/src/resource/resource-store.ts @@ -18,8 +18,18 @@ export class InMemoryResourceStore implements ResourceStore { /** GC roots for the in-memory attachment store: every attachment a resource is backed by. */ referencedAttachmentIds(): AttachmentId[] { + return this.collectReferencedIds(); + } + + /** Attachments this session's resources name — the in-memory `attachment.read` reachability. */ + referencedAttachmentIdsForSession(sessionId: SessionId): AttachmentId[] { + return this.collectReferencedIds(sessionId); + } + + private collectReferencedIds(sessionId?: SessionId): AttachmentId[] { const ids: AttachmentId[] = []; for (const resource of this.resources.values()) { + if (sessionId !== undefined && resource.sessionId !== sessionId) continue; if (resource.attachmentId !== undefined) ids.push(resource.attachmentId); } return ids; From d8b8b94f5aa58a948ff3877410895caa32a28210 Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 18:07:10 +0800 Subject: [PATCH 03/11] feat(engine): add the chunked attachment upload and read service Begin/chunk/commit/abort with a credit window and SHA-256 dedupe, MIME sniff at commit, and reachability-gated reads. Typed request errors for offset, hash, size, and missing ids. --- .../engine/src/attachment/request-handler.ts | 125 ++++++ .../engine/src/attachment/upload-service.ts | 424 ++++++++++++++++++ packages/host/engine/src/engine.ts | 5 + .../host/engine/src/wire/request-router.ts | 9 + 4 files changed, 563 insertions(+) create mode 100644 packages/host/engine/src/attachment/request-handler.ts create mode 100644 packages/host/engine/src/attachment/upload-service.ts diff --git a/packages/host/engine/src/attachment/request-handler.ts b/packages/host/engine/src/attachment/request-handler.ts new file mode 100644 index 000000000..e477b9f4e --- /dev/null +++ b/packages/host/engine/src/attachment/request-handler.ts @@ -0,0 +1,125 @@ +import type { WirePayload } from '@linkcode/schema'; +import type { Transport } from '@linkcode/transport'; +import { createWireMessage } from '@linkcode/transport'; +import { Effect } from 'effect'; +import type { WireResponder } from '../wire/responder'; +import type { AttachmentUploadService } from './upload-service'; + +type AttachmentRequest = Extract< + WirePayload, + { + kind: + | 'attachment.upload.begin' + | 'attachment.upload.chunk' + | 'attachment.upload.commit' + | 'attachment.upload.abort' + | 'attachment.read'; + } +>; + +export class AttachmentRequestHandler { + constructor( + private readonly transport: Transport, + private readonly uploads: AttachmentUploadService, + private readonly responder: WireResponder, + ) {} + + handle(payload: AttachmentRequest): Effect.Effect { + switch (payload.kind) { + case 'attachment.upload.begin': + return this.responder.reply( + payload.clientReqId, + this.uploads + .begin({ + operationId: payload.operationId, + declaredSha256: payload.declaredSha256, + declaredSize: payload.declaredSize, + name: payload.name, + mimeType: payload.mimeType, + attachmentKind: payload.attachmentKind, + }) + .pipe( + Effect.tap((result) => + Effect.sync(() => + this.transport.send( + createWireMessage({ + kind: 'attachment.upload.begun', + replyTo: payload.clientReqId, + ...result, + }), + ), + ), + ), + Effect.asVoid, + ), + ); + case 'attachment.upload.chunk': + return this.responder.reply( + payload.clientReqId, + this.uploads.chunk(payload.uploadId, payload.offset, payload.data).pipe( + Effect.tap((ack) => + Effect.sync(() => + this.transport.send( + createWireMessage({ + kind: 'attachment.upload.chunk.acked', + replyTo: payload.clientReqId, + ...ack, + }), + ), + ), + ), + Effect.asVoid, + ), + ); + case 'attachment.upload.commit': + return this.responder.reply( + payload.clientReqId, + this.uploads.commit(payload.uploadId).pipe( + Effect.tap((result) => + Effect.sync(() => + this.transport.send( + createWireMessage({ + kind: 'attachment.upload.committed', + replyTo: payload.clientReqId, + ...result, + }), + ), + ), + ), + Effect.asVoid, + ), + ); + case 'attachment.upload.abort': + return this.responder.reply( + payload.clientReqId, + this.uploads + .abort(payload.uploadId) + .pipe( + Effect.tap(() => Effect.sync(() => this.responder.sendSuccess(payload.clientReqId))), + ), + ); + case 'attachment.read': + return this.responder.reply( + payload.clientReqId, + this.uploads + .read(payload.sessionId, payload.attachmentId, payload.offset, payload.length) + .pipe( + Effect.tap((result) => + Effect.sync(() => + this.transport.send( + createWireMessage({ + kind: 'attachment.read.result', + replyTo: payload.clientReqId, + ...result, + }), + ), + ), + ), + Effect.asVoid, + ), + ); + default: + return Effect.void; + } + } +} diff --git a/packages/host/engine/src/attachment/upload-service.ts b/packages/host/engine/src/attachment/upload-service.ts new file mode 100644 index 000000000..9f2b0d60f --- /dev/null +++ b/packages/host/engine/src/attachment/upload-service.ts @@ -0,0 +1,424 @@ +import { randomUUID } from 'node:crypto'; +import type { AttachmentId, BlobId, SessionId, UploadId, UploadLease } from '@linkcode/schema'; +import { + ATTACHMENT_UPLOAD_CHUNK_BYTES, + ATTACHMENT_UPLOAD_WINDOW_CHUNKS, + AttachmentIdSchema, + blobIdFromSha256, + MAX_ATTACHMENT_BYTES, + UploadIdSchema, +} from '@linkcode/schema'; +import { Effect } from 'effect'; +import { noop } from 'foxts/noop'; +import { OperationError, RequestError } from '../failure'; +import type { AttachmentStore } from './attachment-store'; +import type { BlobStage, BlobStore } from './blob-store'; +import { BlobIntegrityError } from './blob-store'; +import { UPLOAD_LEASE_TTL_MS } from './gc'; +import { AttachmentIoMutex } from './io-mutex'; +import { declaredMimeTypeMatches } from './mime-sniff'; + +const HEAD_BYTES = 16; +const rUploadId = /^[\w-]{1,128}$/; + +export interface AttachmentBeginInput { + readonly operationId?: string; + readonly declaredSha256: string; + readonly declaredSize: number; + readonly name: string; + readonly mimeType?: string; + readonly attachmentKind: string; +} + +export interface AttachmentBeginResult { + readonly uploadId: UploadId; + readonly chunkBytes: number; + readonly state: 'ready' | 'exists'; +} + +export interface AttachmentChunkAck { + readonly uploadId: UploadId; + readonly receivedBytes: number; +} + +export interface AttachmentCommitResult { + readonly attachmentId: AttachmentId; + readonly blobId: BlobId; +} + +export interface AttachmentReadResult { + readonly sessionId: SessionId; + readonly attachmentId: AttachmentId; + readonly blobId: BlobId; + readonly offset: number; + readonly data: string; + readonly sizeBytes: number; + readonly eof: boolean; +} + +interface LiveUpload { + readonly lease: UploadLease; + readonly stage: BlobStage | undefined; + receivedBytes: number; + head: Uint8Array; + readonly state: 'ready' | 'exists'; +} + +export class AttachmentUploadService { + private readonly live = new Map(); + private readonly begunByOperation = new Map(); + + constructor( + private readonly blobs: BlobStore, + private readonly attachments: AttachmentStore, + private readonly io: AttachmentIoMutex = new AttachmentIoMutex(), + private readonly clock: () => number = Date.now, + ) {} + + begin( + input: AttachmentBeginInput, + ): Effect.Effect { + const store = this.store.bind(this); + const files = this.files.bind(this); + return Effect.gen({ self: this }, function* () { + if (input.operationId !== undefined) { + const replayed = this.begunByOperation.get(input.operationId); + if (replayed) return replayed; + } + if (input.declaredSize > MAX_ATTACHMENT_BYTES) { + return yield* invalid('limit_exceeded', 'Attachment exceeds the 8 MiB limit'); + } + const uploadId = UploadIdSchema.parse(`upl-${randomUUID()}`); + const now = this.clock(); + const lease = yield* store('begin', () => + this.attachments.beginUpload({ + uploadId, + declaredSha256: input.declaredSha256, + declaredSize: input.declaredSize, + name: input.name, + mimeType: input.mimeType, + kind: input.attachmentKind, + expiresAt: now + UPLOAD_LEASE_TTL_MS, + createdAt: now, + }), + ); + const pinnedBlobId = lease.blobId; + let state: 'ready' | 'exists' = 'ready'; + if (pinnedBlobId !== undefined) { + const info = yield* files('stat', () => this.blobs.stat(pinnedBlobId)); + if (info?.sizeBytes === input.declaredSize) state = 'exists'; + } + const stage = + state === 'ready' + ? yield* files('stage', () => this.blobs.stage(uploadId)).pipe( + Effect.tapError(() => + store('begin-cleanup', () => this.attachments.deleteLease(uploadId)).pipe( + Effect.ignore, + ), + ), + ) + : undefined; + this.live.set(uploadId, { + lease, + stage, + receivedBytes: state === 'exists' ? input.declaredSize : 0, + head: new Uint8Array(0), + state, + }); + const result: AttachmentBeginResult = { + uploadId, + chunkBytes: ATTACHMENT_UPLOAD_CHUNK_BYTES, + state, + }; + if (input.operationId !== undefined) this.begunByOperation.set(input.operationId, result); + return result; + }); + } + + chunk( + uploadId: UploadId, + offset: number, + data: string, + ): Effect.Effect { + const store = this.store.bind(this); + const files = this.files.bind(this); + return Effect.gen({ self: this }, function* () { + if (!rUploadId.test(uploadId)) return yield* invalid('not_found', 'Upload not found'); + const live = this.live.get(uploadId); + const stage = live?.stage; + if (!live || stage === undefined) { + const lease = yield* store('getLease', () => this.attachments.getLease(uploadId)); + if (!lease) return yield* invalid('not_found', 'Upload not found'); + if (live?.state === 'exists' || lease.blobId !== undefined) { + return yield* invalid('invalid_request', 'Blob already stored; commit without chunks'); + } + return yield* invalid('conflict', 'Upload is not accepting chunks; retry from begin'); + } + if (offset !== live.receivedBytes) { + return yield* invalid( + 'invalid_request', + `Expected offset ${live.receivedBytes}, got ${offset}`, + ); + } + const windowEnd = + live.receivedBytes + ATTACHMENT_UPLOAD_WINDOW_CHUNKS * ATTACHMENT_UPLOAD_CHUNK_BYTES; + if (offset >= windowEnd && live.receivedBytes < live.lease.declaredSize) { + return yield* invalid('invalid_request', 'Chunk is outside the credit window'); + } + const bytes = yield* Effect.try({ + try: () => decodeChunk(data), + catch: (cause) => mapCause(cause, 'store', 'chunk'), + }); + if (live.receivedBytes + bytes.byteLength > live.lease.declaredSize) { + return yield* invalid('invalid_request', 'Chunk exceeds the declared size'); + } + yield* files('write', () => stage.write(offset, bytes)); + if (offset === 0) live.head = bytes.subarray(0, Math.min(HEAD_BYTES, bytes.byteLength)); + live.receivedBytes += bytes.byteLength; + return { uploadId, receivedBytes: live.receivedBytes }; + }); + } + + commit(uploadId: UploadId): Effect.Effect { + const store = this.store.bind(this); + const files = this.files.bind(this); + const forget = this.forget.bind(this); + const publishExists = this.publishExists.bind(this); + const publishReady = this.publishReady.bind(this); + return Effect.gen({ self: this }, function* () { + if (!rUploadId.test(uploadId)) return yield* invalid('not_found', 'Upload not found'); + const live = this.live.get(uploadId); + const lease = + live?.lease ?? (yield* store('getLease', () => this.attachments.getLease(uploadId))); + if (!lease) return yield* invalid('not_found', 'Upload not found'); + if (lease.attachmentId !== undefined && lease.blobId !== undefined) { + forget(uploadId); + return { attachmentId: lease.attachmentId, blobId: lease.blobId }; + } + const pinnedBlobId = lease.blobId; + const existsFile = + pinnedBlobId === undefined + ? undefined + : yield* files('stat', () => this.blobs.stat(pinnedBlobId)); + const exists = live?.state === 'exists' || existsFile !== undefined; + if (exists) { + const blobId = pinnedBlobId ?? blobIdFromSha256(lease.declaredSha256); + const head = + (yield* files('read', () => this.blobs.read(blobId, 0, HEAD_BYTES))) ?? new Uint8Array(0); + yield* assertMime(lease.mimeType, head); + const committed = yield* publishExists(lease, blobId); + forget(uploadId); + return committed; + } + if (!live?.stage) { + return yield* invalid('conflict', 'Upload is not accepting commit; retry from begin'); + } + if (live.receivedBytes !== lease.declaredSize) { + return yield* invalid( + 'invalid_request', + `Uploaded ${live.receivedBytes} bytes, declared ${lease.declaredSize}`, + ); + } + yield* assertMime(lease.mimeType, live.head); + const committed = yield* publishReady(live); + forget(uploadId); + return committed; + }); + } + + abort(uploadId: UploadId): Effect.Effect { + const store = this.store.bind(this); + const files = this.files.bind(this); + const forget = this.forget.bind(this); + return Effect.gen({ self: this }, function* () { + if (!rUploadId.test(uploadId)) return yield* invalid('not_found', 'Upload not found'); + const live = this.live.get(uploadId); + const stage = live?.stage; + if (stage) { + yield* files('abort', () => stage.abort()); + } else { + const lease = yield* store('getLease', () => this.attachments.getLease(uploadId)); + if (!lease && !live) return yield* invalid('not_found', 'Upload not found'); + } + yield* store('deleteLease', () => this.attachments.deleteLease(uploadId)); + forget(uploadId); + }); + } + + read( + sessionId: SessionId, + attachmentId: AttachmentId, + offset: number, + length: number, + ): Effect.Effect { + const store = this.store.bind(this); + const files = this.files.bind(this); + return Effect.gen({ self: this }, function* () { + const reachable = yield* store('isReachable', () => + this.attachments.isReachable(sessionId, attachmentId), + ); + if (!reachable) return yield* invalid('not_found', 'Attachment not found'); + const attachment = yield* store('getAttachment', () => + this.attachments.getAttachment(attachmentId), + ); + if (!attachment) return yield* invalid('not_found', 'Attachment not found'); + if (offset > attachment.sizeBytes) { + return yield* invalid('invalid_request', 'Read offset is past the end of the attachment'); + } + const bytes = + (yield* files('read', () => this.blobs.read(attachment.blobId, offset, length))) ?? + undefined; + if (bytes === undefined) return yield* invalid('not_found', 'Attachment bytes are missing'); + const end = offset + bytes.byteLength; + return { + sessionId, + attachmentId, + blobId: attachment.blobId, + offset, + data: Buffer.from(bytes).toString('base64'), + sizeBytes: attachment.sizeBytes, + eof: end >= attachment.sizeBytes, + }; + }); + } + + private publishExists( + lease: UploadLease, + blobId: BlobId, + ): Effect.Effect { + const { attachments, clock, io } = this; + const attachmentId = AttachmentIdSchema.parse(`att-${randomUUID()}`); + const now = clock(); + return this.files('commit', async () => { + await io.run(async () => { + await attachments.commitAttachment({ + blob: { blobId, sizeBytes: lease.declaredSize, createdAt: now }, + attachment: recordFromLease(lease, attachmentId, now), + uploadId: lease.uploadId, + }); + }); + return { attachmentId, blobId }; + }); + } + + private publishReady( + live: LiveUpload, + ): Effect.Effect { + const { attachments, blobs, clock, io } = this; + const { lease, stage } = live; + if (!stage) { + return invalid('conflict', 'Upload is not accepting commit; retry from begin'); + } + const attachmentId = AttachmentIdSchema.parse(`att-${randomUUID()}`); + const now = clock(); + const expected = { sha256: lease.declaredSha256, sizeBytes: lease.declaredSize }; + return this.files('commit', async () => { + let blobId: BlobId | undefined; + await io.run(async () => { + try { + blobId = await stage.commit(expected); + await attachments.commitAttachment({ + blob: { blobId, sizeBytes: lease.declaredSize, createdAt: now }, + attachment: recordFromLease(lease, attachmentId, now), + uploadId: lease.uploadId, + }); + } catch (error) { + await stage.abort().catch(noop); + if (blobId) await blobs.delete(blobId); + throw error; + } + }); + if (!blobId) throw new Error('Attachment commit produced no blob id'); + return { attachmentId, blobId }; + }); + } + + private forget(uploadId: UploadId): void { + this.live.delete(uploadId); + for (const [operationId, begun] of this.begunByOperation) { + if (begun.uploadId === uploadId) this.begunByOperation.delete(operationId); + } + } + + private store( + operation: string, + work: () => Promise, + ): Effect.Effect { + return Effect.tryPromise({ + try: async () => work(), + catch: (cause) => mapCause(cause, 'store', operation), + }); + } + + private files( + operation: string, + work: () => Promise, + ): Effect.Effect { + return Effect.tryPromise({ + try: async () => work(), + catch: (cause) => mapCause(cause, 'filesystem', operation), + }); + } +} + +function recordFromLease(lease: UploadLease, attachmentId: AttachmentId, now: number) { + return { + attachmentId, + kind: lease.kind, + name: lease.name, + mimeType: lease.mimeType ?? 'application/octet-stream', + sizeBytes: lease.declaredSize, + metadata: {}, + createdAt: now, + }; +} + +function decodeChunk(data: string): Uint8Array { + if (data.length === 0) return new Uint8Array(0); + const bytes = Buffer.from(data, 'base64'); + if (bytes.byteLength === 0) { + throw new RequestError({ code: 'invalid_request', message: 'Chunk data is not valid base64' }); + } + if (bytes.byteLength > ATTACHMENT_UPLOAD_CHUNK_BYTES) { + throw new RequestError({ + code: 'invalid_request', + message: 'Chunk exceeds the negotiated size', + }); + } + return bytes; +} + +function assertMime( + mimeType: string | undefined, + head: Uint8Array, +): Effect.Effect { + const declared = mimeType ?? 'application/octet-stream'; + if (!declaredMimeTypeMatches(declared, head)) { + return invalid('invalid_request', `File contents are not ${declared}`); + } + return Effect.void; +} + +function invalid( + code: 'invalid_request' | 'limit_exceeded' | 'not_found' | 'conflict', + message: string, +): Effect.Effect { + return Effect.fail(new RequestError({ code, message })); +} + +function mapCause( + cause: unknown, + subsystem: 'store' | 'filesystem', + operation: string, +): RequestError | OperationError { + if (cause instanceof RequestError) return cause; + if (cause instanceof BlobIntegrityError) { + return new RequestError({ code: 'invalid_request', message: cause.message }); + } + return new OperationError({ + subsystem, + operation: `attachments.${operation}`, + publicMessage: 'Attachment operation failed', + cause, + }); +} diff --git a/packages/host/engine/src/engine.ts b/packages/host/engine/src/engine.ts index 4daecb199..4599f9079 100644 --- a/packages/host/engine/src/engine.ts +++ b/packages/host/engine/src/engine.ts @@ -21,6 +21,8 @@ import { InMemoryAttachmentStore } from './attachment/attachment-store'; import { FsBlobStore } from './attachment/blob-store'; import { AttachmentGc } from './attachment/gc'; import { AttachmentIoMutex } from './attachment/io-mutex'; +import { AttachmentRequestHandler } from './attachment/request-handler'; +import { AttachmentUploadService } from './attachment/upload-service'; import { InMemoryLoopStore, InMemoryScheduleStore, @@ -145,6 +147,7 @@ export const createEngineRuntime = Effect.fn('Engine.create')(function* ( attachmentStore, attachmentIo, ); + const uploads = new AttachmentUploadService(blobStore, attachmentStore, attachmentIo); const plugins = new PluginService(deps.pluginFactory ?? createPluginProviderAdapter); const translator = deps.translator; const startOptions = new SessionStartOptionsResolver( @@ -249,6 +252,7 @@ export const createEngineRuntime = Effect.fn('Engine.create')(function* ( const artifacts = new ArtifactHostService(routes); const artifactRequests = new ArtifactRequestHandler(transport, artifacts, responder); const resourceRequests = new ResourceRequestHandler(transport, resources, responder); + const attachmentRequests = new AttachmentRequestHandler(transport, uploads, responder); const conversationCheckpoints = new ConversationCheckpointService( conversationTurns, records, @@ -338,6 +342,7 @@ export const createEngineRuntime = Effect.fn('Engine.create')(function* ( script: scriptRequests, artifact: artifactRequests, resource: resourceRequests, + attachment: attachmentRequests, automation: automationRequests, terminal: terminalRequests, simulator: simulatorRequests, diff --git a/packages/host/engine/src/wire/request-router.ts b/packages/host/engine/src/wire/request-router.ts index 95fb06821..cea683856 100644 --- a/packages/host/engine/src/wire/request-router.ts +++ b/packages/host/engine/src/wire/request-router.ts @@ -4,6 +4,7 @@ import { createWireMessage, pong } from '@linkcode/transport'; import { Effect } from 'effect'; import type { AgentRequestHandler } from '../agent/request-handler'; import type { ManagedAssetService } from '../asset/service'; +import type { AttachmentRequestHandler } from '../attachment/request-handler'; import type { AutomationRequestHandler } from '../automation/request-handler'; import type { BrowserRequestHandler } from '../browser/request-handler'; import type { ConversationRequestHandler } from '../conversation/request-handler'; @@ -33,6 +34,7 @@ interface RequestHandlers { readonly script: ScriptRequestHandler; readonly artifact: ArtifactRequestHandler; readonly resource: ResourceRequestHandler; + readonly attachment: AttachmentRequestHandler; readonly automation: AutomationRequestHandler; readonly terminal: TerminalRequestHandler; readonly simulator: SimulatorRequestHandler; @@ -136,6 +138,13 @@ export class WireRequestRouter { case 'resource.host': { return this.handlers.resource.handle(p); } + case 'attachment.upload.begin': + case 'attachment.upload.chunk': + case 'attachment.upload.commit': + case 'attachment.upload.abort': + case 'attachment.read': { + return this.handlers.attachment.handle(p); + } case 'schedule.create': case 'schedule.update': case 'schedule.delete': From 849ae3416d21d089c679840639f556b301e5a29c Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 18:07:32 +0800 Subject: [PATCH 04/11] test(engine): cover attachment upload reject paths and wire reads Wrong offset/hash/size, missing blob files, path-shaped ids, and cross-session reads fail typed. Dedupe exists and resource-rooted reads go through the engine. --- .../src/__tests__/attachment-upload.test.ts | 221 ++++++++++++++++++ .../src/__tests__/engine-attachments.test.ts | 168 +++++++++++++ 2 files changed, 389 insertions(+) create mode 100644 packages/host/engine/src/__tests__/attachment-upload.test.ts create mode 100644 packages/host/engine/src/__tests__/engine-attachments.test.ts diff --git a/packages/host/engine/src/__tests__/attachment-upload.test.ts b/packages/host/engine/src/__tests__/attachment-upload.test.ts new file mode 100644 index 000000000..9d3640c1b --- /dev/null +++ b/packages/host/engine/src/__tests__/attachment-upload.test.ts @@ -0,0 +1,221 @@ +import { createHash } from 'node:crypto'; +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import type { AttachmentId, UploadId } from '@linkcode/schema'; +import { AttachmentIdSchema, blobIdFromSha256, SessionIdSchema } from '@linkcode/schema'; +import { Effect } from 'effect'; +import { afterEach, describe, expect, it } from 'vitest'; +import { InMemoryAttachmentStore } from '../attachment/attachment-store'; +import { FsBlobStore } from '../attachment/blob-store'; +import { AttachmentUploadService } from '../attachment/upload-service'; + +const temporaryDirectories: string[] = []; +const sessionId = SessionIdSchema.parse('session-1'); +const otherSession = SessionIdSchema.parse('session-2'); + +afterEach(async () => { + await Promise.all( + temporaryDirectories.splice(0).map((path) => rm(path, { recursive: true, force: true })), + ); +}); + +function sha256(bytes: Uint8Array): string { + return createHash('sha256').update(bytes).digest('hex'); +} + +async function makeService(reachableIds: AttachmentId[] = []) { + const root = await mkdtemp(join(tmpdir(), 'linkcode-upload-')); + temporaryDirectories.push(root); + const reachable = new Set(reachableIds); + const blobs = new FsBlobStore(join(root, 'blobs')); + const attachments = new InMemoryAttachmentStore( + () => reachable, + (sid, id) => sid === sessionId && reachable.has(id), + ); + const uploads = new AttachmentUploadService(blobs, attachments); + return { attachments, blobs, reachable, uploads }; +} + +async function run(effect: Effect.Effect): Promise { + return Effect.runPromise(effect); +} + +describe('AttachmentUploadService', () => { + it('uploads bytes, then reads them only from a session that references the attachment', async () => { + const { reachable, uploads } = await makeService(); + const bytes = Buffer.from('hello attachment'); + const begun = await run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'note.txt', + mimeType: 'text/plain', + attachmentKind: 'file', + }), + ); + expect(begun.state).toBe('ready'); + const ack = await run(uploads.chunk(begun.uploadId, 0, bytes.toString('base64'))); + expect(ack.receivedBytes).toBe(bytes.byteLength); + const committed = await run(uploads.commit(begun.uploadId)); + reachable.add(committed.attachmentId); + + const page = await run(uploads.read(sessionId, committed.attachmentId, 0, bytes.byteLength)); + expect(Buffer.from(page.data, 'base64').toString()).toBe('hello attachment'); + expect(page.eof).toBe(true); + expect(page.blobId).toBe(blobIdFromSha256(sha256(bytes))); + + await expect( + run(uploads.read(otherSession, committed.attachmentId, 0, 4)), + ).rejects.toMatchObject({ _tag: 'RequestError', code: 'not_found' }); + }); + + it('short-circuits begin when the blob already exists on disk', async () => { + const { uploads } = await makeService(); + const bytes = Buffer.from('same screenshot twice'); + const input = { + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'shot.png', + mimeType: 'text/plain', + attachmentKind: 'file', + }; + const first = await run(uploads.begin(input)); + await run(uploads.chunk(first.uploadId, 0, bytes.toString('base64'))); + const committed = await run(uploads.commit(first.uploadId)); + + const second = await run(uploads.begin({ ...input, name: 'copy.png' })); + expect(second.state).toBe('exists'); + const again = await run(uploads.commit(second.uploadId)); + expect(again.blobId).toBe(committed.blobId); + expect(again.attachmentId).not.toBe(committed.attachmentId); + }); + + it('treats a missing blob file as ready even when the row is pinned', async () => { + const { blobs, uploads } = await makeService(); + const bytes = Buffer.from('will vanish'); + const input = { + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'gone.bin', + attachmentKind: 'file', + }; + const first = await run(uploads.begin(input)); + await run(uploads.chunk(first.uploadId, 0, bytes.toString('base64'))); + const committed = await run(uploads.commit(first.uploadId)); + await blobs.delete(committed.blobId); + + const again = await run(uploads.begin(input)); + expect(again.state).toBe('ready'); + await run(uploads.chunk(again.uploadId, 0, bytes.toString('base64'))); + const restored = await run(uploads.commit(again.uploadId)); + expect(restored.blobId).toBe(committed.blobId); + expect(await blobs.stat(restored.blobId)).toEqual({ sizeBytes: bytes.byteLength }); + }); + + it('rejects a wrong offset, a hash mismatch, an oversize claim, and a path-shaped id', async () => { + const { uploads } = await makeService(); + const bytes = Buffer.from('payload'); + const begun = await run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'payload.bin', + attachmentKind: 'file', + }), + ); + await expect( + run(uploads.chunk(begun.uploadId, 1, bytes.toString('base64'))), + ).rejects.toMatchObject({ _tag: 'RequestError', code: 'invalid_request' }); + + const wrongHash = await run( + uploads.begin({ + declaredSha256: sha256(Buffer.from('other')), + declaredSize: bytes.byteLength, + name: 'wrong.bin', + attachmentKind: 'file', + }), + ); + await run(uploads.chunk(wrongHash.uploadId, 0, bytes.toString('base64'))); + await expect(run(uploads.commit(wrongHash.uploadId))).rejects.toMatchObject({ + _tag: 'RequestError', + code: 'invalid_request', + }); + + await expect( + run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: 9 * 1024 * 1024, + name: 'huge.bin', + attachmentKind: 'file', + }), + ), + ).rejects.toMatchObject({ _tag: 'RequestError', code: 'limit_exceeded' }); + + await expect(run(uploads.chunk('upl-../escape' as UploadId, 0, 'YQ=='))).rejects.toMatchObject({ + _tag: 'RequestError', + code: 'not_found', + }); + }); + + it('replays begin by operationId and aborts a live stage', async () => { + const { attachments, uploads } = await makeService(); + const bytes = Buffer.from('draft'); + const first = await run( + uploads.begin({ + operationId: 'op-1', + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'draft.bin', + attachmentKind: 'file', + }), + ); + const replayed = await run( + uploads.begin({ + operationId: 'op-1', + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'draft.bin', + attachmentKind: 'file', + }), + ); + expect(replayed.uploadId).toBe(first.uploadId); + await run(uploads.abort(first.uploadId)); + expect(await attachments.getLease(first.uploadId)).toBeUndefined(); + await expect( + run(uploads.chunk(first.uploadId, 0, bytes.toString('base64'))), + ).rejects.toMatchObject({ _tag: 'RequestError', code: 'not_found' }); + }); + + it('rejects a sniffed image whose bytes are not that type', async () => { + const { uploads } = await makeService(); + const bytes = Buffer.from('not a png'); + const begun = await run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'fake.png', + mimeType: 'image/png', + attachmentKind: 'image', + }), + ); + await run(uploads.chunk(begun.uploadId, 0, bytes.toString('base64'))); + await expect(run(uploads.commit(begun.uploadId))).rejects.toMatchObject({ + _tag: 'RequestError', + code: 'invalid_request', + }); + }); +}); + +describe('in-memory attachment reachability', () => { + it('is false until the session-scoped callback says otherwise', async () => { + const id = AttachmentIdSchema.parse('att-1'); + const store = new InMemoryAttachmentStore( + () => [id], + (sid) => sid === sessionId, + ); + expect(await store.isReachable(sessionId, id)).toBe(true); + expect(await store.isReachable(otherSession, id)).toBe(false); + }); +}); diff --git a/packages/host/engine/src/__tests__/engine-attachments.test.ts b/packages/host/engine/src/__tests__/engine-attachments.test.ts new file mode 100644 index 000000000..4475db9e1 --- /dev/null +++ b/packages/host/engine/src/__tests__/engine-attachments.test.ts @@ -0,0 +1,168 @@ +import { createHash } from 'node:crypto'; +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import type { WirePayload } from '@linkcode/schema'; +import { ATTACHMENT_UPLOAD_CHUNK_BYTES, SessionIdSchema } from '@linkcode/schema'; +import { nullthrow } from 'foxts/guard'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { createSessionHarness, startedSessionId } from './fixtures/session-harness'; + +const temporaryDirectories: string[] = []; + +afterEach(async () => { + await Promise.all( + temporaryDirectories.splice(0).map((path) => rm(path, { recursive: true, force: true })), + ); +}); + +function sha256(bytes: Uint8Array): string { + return createHash('sha256').update(bytes).digest('hex'); +} + +async function tempDirectory(): Promise { + const path = await mkdtemp(join(tmpdir(), 'linkcode-attachments-')); + temporaryDirectories.push(path); + return path; +} + +function replyOf( + sent: WirePayload[], + kind: K, + replyTo: string, +): Extract { + const reply = sent.find( + (payload) => payload.kind === kind && 'replyTo' in payload && payload.replyTo === replyTo, + ); + if (reply?.kind !== kind) throw new Error(`no ${kind} for ${replyTo}`); + return reply as Extract; +} + +describe('engine attachment upload/read', () => { + it('uploads through the wire, then reads after a session resource roots the attachment', async () => { + const stateDir = await tempDirectory(); + const h = createSessionHarness( + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + { stateDir }, + ); + await h.engine.start(); + await h.inject({ + kind: 'session.start', + clientReqId: 'start', + opts: { kind: 'claude-code', cwd: stateDir }, + }); + const sessionId = startedSessionId(h.sent, 'start'); + const bytes = Buffer.from('wire upload'); + await h.inject({ + kind: 'attachment.upload.begin', + clientReqId: 'begin', + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'note.txt', + mimeType: 'text/plain', + attachmentKind: 'file', + }); + await vi.waitFor(() => { + expect(h.sent).toContainEqual( + expect.objectContaining({ kind: 'attachment.upload.begun', replyTo: 'begin' }), + ); + }); + const begun = replyOf(h.sent, 'attachment.upload.begun', 'begin'); + expect(begun.state).toBe('ready'); + expect(begun.chunkBytes).toBe(ATTACHMENT_UPLOAD_CHUNK_BYTES); + + await h.inject({ + kind: 'attachment.upload.chunk', + clientReqId: 'bad-offset', + uploadId: begun.uploadId, + offset: 99, + data: 'YQ==', + }); + await vi.waitFor(() => { + expect(h.sent).toContainEqual( + expect.objectContaining({ + kind: 'request.failed', + replyTo: 'bad-offset', + code: 'invalid_request', + }), + ); + }); + + await h.inject({ + kind: 'attachment.upload.chunk', + clientReqId: 'chunk', + uploadId: begun.uploadId, + offset: 0, + data: bytes.toString('base64'), + }); + await vi.waitFor(() => { + expect(h.sent).toContainEqual( + expect.objectContaining({ kind: 'attachment.upload.chunk.acked', replyTo: 'chunk' }), + ); + }); + + await h.inject({ + kind: 'attachment.upload.commit', + clientReqId: 'commit', + uploadId: begun.uploadId, + }); + await vi.waitFor(() => { + expect(h.sent).toContainEqual( + expect.objectContaining({ kind: 'attachment.upload.committed', replyTo: 'commit' }), + ); + }); + const committed = replyOf(h.sent, 'attachment.upload.committed', 'commit'); + + await h.inject({ + kind: 'resource.source.upload', + clientReqId: 'resource', + sessionId, + name: 'note.txt', + mimeType: 'text/plain', + data: bytes.toString('base64'), + }); + await vi.waitFor(() => { + expect(h.sent).toContainEqual( + expect.objectContaining({ kind: 'resource.uploaded', replyTo: 'resource' }), + ); + }); + const resource = replyOf(h.sent, 'resource.uploaded', 'resource').resource; + const attachmentId = nullthrow(resource.attachmentId, 'resource missing attachmentId'); + + await h.inject({ + kind: 'attachment.read', + clientReqId: 'read', + sessionId, + attachmentId, + offset: 0, + length: bytes.byteLength, + }); + await vi.waitFor(() => { + expect(h.sent).toContainEqual( + expect.objectContaining({ kind: 'attachment.read.result', replyTo: 'read' }), + ); + }); + const page = replyOf(h.sent, 'attachment.read.result', 'read'); + expect(Buffer.from(page.data, 'base64').toString()).toBe('wire upload'); + expect(page.blobId).toBe(committed.blobId); + + await h.inject({ + kind: 'attachment.read', + clientReqId: 'cross', + sessionId: SessionIdSchema.parse('session-other'), + attachmentId, + offset: 0, + length: 4, + }); + await vi.waitFor(() => { + expect(h.sent).toContainEqual( + expect.objectContaining({ kind: 'request.failed', replyTo: 'cross', code: 'not_found' }), + ); + }); + }); +}); From 82298d3b28b07fc6958926861ba08e5658d66a25 Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 18:08:07 +0800 Subject: [PATCH 05/11] feat(client-core): add chunked attachment upload, read, and blob cache Feature-detects wire 80, pipelines chunks in a two-chunk credit window, and caches decoded bytes by blobId. --- packages/client/core/src/client.ts | 84 +++++++++ .../core/src/client/attachment-channel.ts | 171 ++++++++++++++++++ packages/client/core/src/client/blob-cache.ts | 59 ++++++ .../core/src/client/pending-registry.ts | 28 +++ .../integration/attachment-client.test.ts | 114 ++++++++++++ 5 files changed, 456 insertions(+) create mode 100644 packages/client/core/src/client/attachment-channel.ts create mode 100644 packages/client/core/src/client/blob-cache.ts create mode 100644 packages/client/core/tests/integration/attachment-client.test.ts diff --git a/packages/client/core/src/client.ts b/packages/client/core/src/client.ts index b6fc026bb..423e09b36 100644 --- a/packages/client/core/src/client.ts +++ b/packages/client/core/src/client.ts @@ -10,6 +10,7 @@ import type { AgentKind, AgentRuntimes, AgentStartCatalog, + AttachmentId, ContentBlock, CustomMcpServerPatchOp, CustomMcpServerPublic, @@ -67,6 +68,7 @@ import type { StartOptions, TerminalMetadata, TerminalReplayEvent, + UploadId, WireMessage, WorkspaceFile, WorkspaceId, @@ -75,6 +77,7 @@ import type { WorkspaceScript, } from '@linkcode/schema'; import { + ATTACHMENT_STORE_WIRE_VERSION, CONVERSATION_GRAPH_WIRE_VERSION, MIN_COMPATIBLE_WIRE_VERSION, WIRE_PROTOCOL_VERSION, @@ -85,6 +88,12 @@ import { extractErrorMessage, isErrorLikeObject } from 'foxts/extract-error-mess import { noop } from 'foxts/noop'; import type { AgentLoginHandlers } from './client/agent-login-channel'; import { AgentLoginChannel } from './client/agent-login-channel'; +import type { + AttachmentBeginInput, + AttachmentPutInput, + AttachmentReadBytes, +} from './client/attachment-channel'; +import { AttachmentChannel } from './client/attachment-channel'; import type { BrowserCommandExecutor } from './client/browser-host-channel'; import { BrowserHostChannel } from './client/browser-host-channel'; import type { @@ -111,6 +120,11 @@ import { PendingRegistry, resolveRandomUUID } from './client/pending-registry'; import { TerminalChannel } from './client/terminal-channel'; export type { AgentLoginHandlers, AgentLoginSettled } from './client/agent-login-channel'; +export type { + AttachmentBeginInput, + AttachmentPutInput, + AttachmentReadBytes, +} from './client/attachment-channel'; export type { BrowserCommandExecutor } from './client/browser-host-channel'; export type { ConversationReadClientOptions, @@ -240,6 +254,7 @@ export function isRequestFailureReportedInConversation(error: unknown): boolean export class LinkCodeClient { private readonly pending: PendingRegistry; private readonly control: ControlChannel; + private readonly attachments: AttachmentChannel; private readonly events = new EventBuffer(); private readonly graphChanges = new ConversationGraphChanges(); private readonly terminals: TerminalChannel; @@ -277,6 +292,7 @@ export class LinkCodeClient { const randomUUID = resolveRandomUUID(options.randomUUID); this.pending = new PendingRegistry(randomUUID); this.control = new ControlChannel(transport, this.pending); + this.attachments = new AttachmentChannel(transport, this.pending); this.terminals = new TerminalChannel(transport, this.pending, randomUUID); this.browserHost = new BrowserHostChannel(transport, this.pending, randomUUID); this.agentLogin = new AgentLoginChannel(transport, this.pending); @@ -326,6 +342,11 @@ export class LinkCodeClient { return this.peerWire !== null && this.peerWire.version >= CONVERSATION_GRAPH_WIRE_VERSION; } + /** Whether the host serves chunked `attachment.upload.*` / `attachment.read`. */ + get supportsAttachmentStore(): boolean { + return this.peerWire !== null && this.peerWire.version >= ATTACHMENT_STORE_WIRE_VERSION; + } + private async handshake(): Promise { let settled = false; let cancelTimer: () => void = noop; @@ -607,6 +628,36 @@ export class LinkCodeClient { case 'resource.hosted': this.pending.resolve('resourceHost', p.replyTo, p.hosted); break; + case 'attachment.upload.begun': + this.pending.resolve('attachmentBegin', p.replyTo, { + uploadId: p.uploadId, + chunkBytes: p.chunkBytes, + state: p.state, + }); + break; + case 'attachment.upload.chunk.acked': + this.pending.resolve('attachmentChunk', p.replyTo, { + uploadId: p.uploadId, + receivedBytes: p.receivedBytes, + }); + break; + case 'attachment.upload.committed': + this.pending.resolve('attachmentCommit', p.replyTo, { + attachmentId: p.attachmentId, + blobId: p.blobId, + }); + break; + case 'attachment.read.result': + this.pending.resolve('attachmentRead', p.replyTo, { + sessionId: p.sessionId, + attachmentId: p.attachmentId, + blobId: p.blobId, + offset: p.offset, + data: p.data, + sizeBytes: p.sizeBytes, + eof: p.eof, + }); + break; case 'resource.changed': for (const cb of this.resourceEventSubs) cb({ type: 'changed', resource: p.resource }); break; @@ -1323,6 +1374,39 @@ export class LinkCodeClient { hostResource(resourceId: SessionResourceId): Promise { return this.control.hostResource(resourceId); } + + beginAttachmentUpload(input: AttachmentBeginInput) { + return this.attachments.beginUpload(input); + } + + sendAttachmentChunk(uploadId: UploadId, offset: number, data: string) { + return this.attachments.sendChunk(uploadId, offset, data); + } + + commitAttachmentUpload(uploadId: UploadId) { + return this.attachments.commit(uploadId); + } + + abortAttachmentUpload(uploadId: UploadId) { + return this.attachments.abort(uploadId); + } + + readAttachment(sessionId: SessionId, attachmentId: AttachmentId, offset: number, length: number) { + return this.attachments.read(sessionId, attachmentId, offset, length); + } + + /** Hash + windowed chunked upload. Identical bytes commit with no transfer. */ + putAttachment(input: AttachmentPutInput) { + return this.attachments.put(input); + } + + /** Read every byte of an attachment, cached by `blobId`. */ + getAttachmentBytes( + sessionId: SessionId, + attachmentId: AttachmentId, + ): Promise { + return this.attachments.get(sessionId, attachmentId); + } subscribeResources(cb: ResourceEventCb): Unsubscribe { this.resourceEventSubs.add(cb); return () => this.resourceEventSubs.delete(cb); diff --git a/packages/client/core/src/client/attachment-channel.ts b/packages/client/core/src/client/attachment-channel.ts new file mode 100644 index 000000000..abdb212da --- /dev/null +++ b/packages/client/core/src/client/attachment-channel.ts @@ -0,0 +1,171 @@ +import type { AttachmentId, OperationId, SessionId, UploadId } from '@linkcode/schema'; +import { ATTACHMENT_UPLOAD_CHUNK_BYTES, ATTACHMENT_UPLOAD_WINDOW_CHUNKS } from '@linkcode/schema'; +import type { Transport } from '@linkcode/transport'; +import { noop } from 'foxts/noop'; +import { AttachmentBlobCache, base64ToBytes, bytesToBase64, sha256Hex } from './blob-cache'; +import type { + AttachmentChunkAck, + AttachmentCommitResult, + AttachmentReadResult, + AttachmentUploadBegun, + PendingRegistry, +} from './pending-registry'; +import { sendCorrelated } from './pending-registry'; + +export interface AttachmentBeginInput { + readonly declaredSha256: string; + readonly declaredSize: number; + readonly name: string; + readonly mimeType?: string; + readonly attachmentKind: string; + readonly operationId?: OperationId; +} + +export interface AttachmentPutInput { + readonly bytes: Uint8Array; + readonly name: string; + readonly mimeType?: string; + readonly attachmentKind: string; + readonly operationId?: OperationId; +} + +/** + * Chunked attachment upload/read. Credit-windowed puts retry from zero; completed re-sends are + * free via the daemon's SHA-256 short-circuit. The cache is keyed by `blobId`. + */ +export class AttachmentChannel { + readonly cache = new AttachmentBlobCache(); + + constructor( + private readonly transport: Transport, + private readonly pending: PendingRegistry, + ) {} + + beginUpload(input: AttachmentBeginInput): Promise { + return sendCorrelated(this.transport, this.pending, 'attachmentBegin', (clientReqId) => ({ + kind: 'attachment.upload.begin', + clientReqId, + declaredSha256: input.declaredSha256, + declaredSize: input.declaredSize, + name: input.name, + mimeType: input.mimeType, + attachmentKind: input.attachmentKind, + operationId: input.operationId, + })); + } + + sendChunk(uploadId: UploadId, offset: number, data: string): Promise { + return sendCorrelated(this.transport, this.pending, 'attachmentChunk', (clientReqId) => ({ + kind: 'attachment.upload.chunk', + clientReqId, + uploadId, + offset, + data, + })); + } + + commit(uploadId: UploadId): Promise { + return sendCorrelated(this.transport, this.pending, 'attachmentCommit', (clientReqId) => ({ + kind: 'attachment.upload.commit', + clientReqId, + uploadId, + })); + } + + abort(uploadId: UploadId): Promise<{ ok: true }> { + return sendCorrelated(this.transport, this.pending, 'ack', (clientReqId) => ({ + kind: 'attachment.upload.abort', + clientReqId, + uploadId, + })); + } + + read( + sessionId: SessionId, + attachmentId: AttachmentId, + offset: number, + length: number, + ): Promise { + return sendCorrelated(this.transport, this.pending, 'attachmentRead', (clientReqId) => ({ + kind: 'attachment.read', + clientReqId, + sessionId, + attachmentId, + offset, + length, + })); + } + + /** Hash, begin, windowed chunks, commit. Identical bytes short-circuit to `exists`. */ + async put(input: AttachmentPutInput): Promise { + const declaredSha256 = await sha256Hex(input.bytes); + const begun = await this.beginUpload({ + declaredSha256, + declaredSize: input.bytes.byteLength, + name: input.name, + mimeType: input.mimeType, + attachmentKind: input.attachmentKind, + operationId: input.operationId, + }); + if (begun.state === 'ready') { + await this.sendWindowed(begun.uploadId, begun.chunkBytes, input.bytes); + } + const committed = await this.commit(begun.uploadId); + this.cache.set(committed.blobId, input.bytes); + return committed; + } + + /** Assemble the attachment through `attachment.read`, using the cache when the blob is known. */ + async get(sessionId: SessionId, attachmentId: AttachmentId): Promise { + const first = await this.read(sessionId, attachmentId, 0, ATTACHMENT_UPLOAD_CHUNK_BYTES); + const cached = this.cache.get(first.blobId); + if (cached?.byteLength === first.sizeBytes) { + return { blobId: first.blobId, bytes: cached, sizeBytes: first.sizeBytes }; + } + const bytes = new Uint8Array(first.sizeBytes); + const firstSlice = base64ToBytes(first.data); + bytes.set(firstSlice, first.offset); + let offset = first.offset + firstSlice.byteLength; + while (offset < first.sizeBytes) { + // eslint-disable-next-line no-await-in-loop -- sequential pages of one attachment + const page = await this.read(sessionId, attachmentId, offset, ATTACHMENT_UPLOAD_CHUNK_BYTES); + const slice = base64ToBytes(page.data); + bytes.set(slice, page.offset); + offset = page.offset + slice.byteLength; + if (page.eof) break; + } + this.cache.set(first.blobId, bytes); + return { blobId: first.blobId, bytes, sizeBytes: first.sizeBytes }; + } + + private async sendWindowed( + uploadId: UploadId, + chunkBytes: number, + bytes: Uint8Array, + ): Promise { + const acks: Array> = []; + let offset = 0; + while (offset < bytes.byteLength) { + const inflight = acks.length; + if (inflight >= ATTACHMENT_UPLOAD_WINDOW_CHUNKS) { + // eslint-disable-next-line no-await-in-loop -- credit window drains the oldest ack + await acks[inflight - ATTACHMENT_UPLOAD_WINDOW_CHUNKS]; + } + const end = Math.min(offset + chunkBytes, bytes.byteLength); + const at = offset; + const data = bytesToBase64(bytes.subarray(at, end)); + offset = end; + acks.push(this.sendChunk(uploadId, at, data).then(noop)); + } + for (let i = 0, len = acks.length; i < len; i++) { + // eslint-disable-next-line no-await-in-loop -- drain remaining acks + await acks[i]; + } + } +} + +export interface AttachmentReadBytes { + readonly blobId: string; + readonly bytes: Uint8Array; + readonly sizeBytes: number; +} diff --git a/packages/client/core/src/client/blob-cache.ts b/packages/client/core/src/client/blob-cache.ts new file mode 100644 index 000000000..70c845cb3 --- /dev/null +++ b/packages/client/core/src/client/blob-cache.ts @@ -0,0 +1,59 @@ +/** Decoded attachment bytes keyed by content-addressed `blobId`. Blobs are immutable. */ +export class AttachmentBlobCache { + private readonly blobs = new Map(); + + get(blobId: string): Uint8Array | undefined { + return this.blobs.get(blobId); + } + + set(blobId: string, bytes: Uint8Array): void { + this.blobs.set(blobId, bytes); + } + + has(blobId: string): boolean { + return this.blobs.has(blobId); + } +} + +export function bytesToBase64(bytes: Uint8Array): string { + const nodeBuffer = ( + globalThis as { Buffer?: { from(data: Uint8Array): { toString(encoding: string): string } } } + ).Buffer; + if (nodeBuffer) return nodeBuffer.from(bytes).toString('base64'); + let binary = ''; + for (let i = 0, len = bytes.byteLength; i < len; i++) { + binary += String.fromCharCode(bytes[i]); + } + return btoa(binary); +} + +export function base64ToBytes(data: string): Uint8Array { + const nodeBuffer = ( + globalThis as { + Buffer?: { from(data: string, encoding: string): Uint8Array }; + } + ).Buffer; + if (nodeBuffer) return new Uint8Array(nodeBuffer.from(data, 'base64')); + const binary = atob(data); + const bytes = new Uint8Array(binary.length); + for (let i = 0, len = binary.length; i < len; i++) { + bytes[i] = binary.charCodeAt(i); + } + return bytes; +} + +export async function sha256Hex(bytes: Uint8Array): Promise { + const digest = await crypto.subtle.digest('SHA-256', arrayBufferOf(bytes)); + const view = new Uint8Array(digest); + let hex = ''; + for (let i = 0, len = view.byteLength; i < len; i++) { + hex += view[i].toString(16).padStart(2, '0'); + } + return hex; +} + +function arrayBufferOf(bytes: Uint8Array): ArrayBuffer { + const buffer = new ArrayBuffer(bytes.byteLength); + new Uint8Array(buffer).set(bytes); + return buffer; +} diff --git a/packages/client/core/src/client/pending-registry.ts b/packages/client/core/src/client/pending-registry.ts index 7b88e0cee..5abfe21ff 100644 --- a/packages/client/core/src/client/pending-registry.ts +++ b/packages/client/core/src/client/pending-registry.ts @@ -88,6 +88,26 @@ export type ConversationReadPage = Omit< 'kind' | 'replyTo' >; +export type AttachmentUploadBegun = Omit< + Extract, + 'kind' | 'replyTo' +>; + +export type AttachmentChunkAck = Omit< + Extract, + 'kind' | 'replyTo' +>; + +export type AttachmentCommitResult = Omit< + Extract, + 'kind' | 'replyTo' +>; + +export type AttachmentReadResult = Omit< + Extract, + 'kind' | 'replyTo' +>; + export type RandomUUID = () => string; export function resolveRandomUUID(provider?: RandomUUID): RandomUUID { @@ -138,6 +158,10 @@ export interface PendingValueMap { resourceList: SessionResource[]; resourceUpload: SessionResource; resourceHost: HostedSessionResource; + attachmentBegin: AttachmentUploadBegun; + attachmentChunk: AttachmentChunkAck; + attachmentCommit: AttachmentCommitResult; + attachmentRead: AttachmentReadResult; workspaceList: WorkspaceRecord[]; workspaceRegister: WorkspaceRecord; scheduleCreate: Schedule; @@ -202,6 +226,10 @@ export class PendingRegistry { resourceList: new Map(), resourceUpload: new Map(), resourceHost: new Map(), + attachmentBegin: new Map(), + attachmentChunk: new Map(), + attachmentCommit: new Map(), + attachmentRead: new Map(), workspaceList: new Map(), workspaceRegister: new Map(), scheduleCreate: new Map(), diff --git a/packages/client/core/tests/integration/attachment-client.test.ts b/packages/client/core/tests/integration/attachment-client.test.ts new file mode 100644 index 000000000..7678eed8f --- /dev/null +++ b/packages/client/core/tests/integration/attachment-client.test.ts @@ -0,0 +1,114 @@ +import { + ATTACHMENT_STORE_WIRE_VERSION, + AttachmentIdSchema, + BlobIdSchema, + SessionIdSchema, + UploadIdSchema, +} from '@linkcode/schema'; +import { createLocalTransportPair, createWireMessage } from '@linkcode/transport'; +import { describe, expect, it } from 'vitest'; +import { LinkCodeClient } from '../../src/client'; +import { base64ToBytes, bytesToBase64, sha256Hex } from '../../src/client/blob-cache'; +import { createConnectedLocalClient } from '../support/local-client'; + +describe('LinkCodeClient attachment store API', () => { + it('advertises the store only for hosts at or above its wire version', async () => { + const current = await createConnectedLocalClient(); + expect(current.client.supportsAttachmentStore).toBe(true); + current.client.dispose(); + current.serverTransport.close(); + + const [clientTransport, serverTransport] = createLocalTransportPair(); + await serverTransport.connect(); + serverTransport.onMessage((message) => { + if (message.payload.kind === 'ping') { + serverTransport.send( + createWireMessage({ + kind: 'pong', + version: ATTACHMENT_STORE_WIRE_VERSION - 1, + minCompatible: ATTACHMENT_STORE_WIRE_VERSION - 4, + }), + ); + } + }); + const older = new LinkCodeClient(clientTransport); + await older.connect(); + expect(older.supportsAttachmentStore).toBe(false); + older.dispose(); + serverTransport.close(); + }); + + it('uploads and reads through correlated frames and caches by blobId', async () => { + const { client, serverTransport } = await createConnectedLocalClient(); + const bytes = new TextEncoder().encode('cached blob'); + const digest = await sha256Hex(bytes); + const blobId = BlobIdSchema.parse(`sha256:${digest}`); + const attachmentId = AttachmentIdSchema.parse('att-1'); + const sessionId = SessionIdSchema.parse('session-1'); + + serverTransport.onMessage((message) => { + const p = message.payload; + if (p.kind === 'attachment.upload.begin') { + serverTransport.send( + createWireMessage({ + kind: 'attachment.upload.begun', + replyTo: p.clientReqId, + uploadId: UploadIdSchema.parse('upl-1'), + chunkBytes: 256 * 1024, + state: 'ready', + }), + ); + } + if (p.kind === 'attachment.upload.chunk') { + serverTransport.send( + createWireMessage({ + kind: 'attachment.upload.chunk.acked', + replyTo: p.clientReqId, + uploadId: p.uploadId, + receivedBytes: p.offset + base64ToBytes(p.data).byteLength, + }), + ); + } + if (p.kind === 'attachment.upload.commit') { + serverTransport.send( + createWireMessage({ + kind: 'attachment.upload.committed', + replyTo: p.clientReqId, + attachmentId, + blobId, + }), + ); + } + if (p.kind === 'attachment.read') { + serverTransport.send( + createWireMessage({ + kind: 'attachment.read.result', + replyTo: p.clientReqId, + sessionId: p.sessionId, + attachmentId: p.attachmentId, + blobId, + offset: p.offset, + data: bytesToBase64(bytes.subarray(p.offset, p.offset + p.length)), + sizeBytes: bytes.byteLength, + eof: p.offset + p.length >= bytes.byteLength, + }), + ); + } + }); + + const put = await client.putAttachment({ + bytes, + name: 'note.txt', + mimeType: 'text/plain', + attachmentKind: 'file', + }); + expect(put).toEqual({ attachmentId, blobId }); + + const got = await client.getAttachmentBytes(sessionId, attachmentId); + expect(new TextDecoder().decode(got.bytes)).toBe('cached blob'); + expect(got.blobId).toBe(blobId); + + client.dispose(); + serverTransport.close(); + }); +}); From 535f39e877f72b080c0b047507fdb42ffba18385 Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 18:08:27 +0800 Subject: [PATCH 06/11] feat(workbench): implement chunked attachment frames on the dev mock Begin/chunk/commit/abort/read run in memory so showcase and mock E2E exercise the new protocol instead of silently dropping the kinds. --- .../workbench/src/mock/dev-mock-host.ts | 277 +++++++++++++++++- 1 file changed, 276 insertions(+), 1 deletion(-) diff --git a/packages/client/workbench/src/mock/dev-mock-host.ts b/packages/client/workbench/src/mock/dev-mock-host.ts index 83a4ef5b6..632b0ef3b 100644 --- a/packages/client/workbench/src/mock/dev-mock-host.ts +++ b/packages/client/workbench/src/mock/dev-mock-host.ts @@ -6,6 +6,8 @@ import type { AgentInput, AgentKind, AgentRuntimes, + AttachmentId, + BlobId, ContentBlock, ConversationGraphTurn, ConversationReadItem, @@ -34,6 +36,7 @@ import type { ToolCall, TurnId, TurnSubmitInput, + UploadId, WireMessage, WirePayload, WorkspaceId, @@ -42,6 +45,10 @@ import type { } from '@linkcode/schema'; import { AGENT_INPUT_CAPABILITIES, + ATTACHMENT_UPLOAD_CHUNK_BYTES, + ATTACHMENT_UPLOAD_WINDOW_CHUNKS, + AttachmentIdSchema, + blobIdFromSha256, managedAgentAssetId, managedAssetIdEquals, managedAssetKey, @@ -49,6 +56,7 @@ import { normalizeCwdKey, SessionResourceIdSchema, textBlock, + UploadIdSchema, userRowMessageId, } from '@linkcode/schema'; import type { Transport } from '@linkcode/transport'; @@ -175,6 +183,26 @@ interface MockTerminal { attachments: Map; } +interface MockAttachmentUpload { + uploadId: UploadId; + declaredSha256: string; + declaredSize: number; + name: string; + mimeType?: string; + attachmentKind: string; + received: number; + bytes: Uint8Array; + state: 'ready' | 'exists'; + attachmentId?: AttachmentId; + blobId?: BlobId; +} + +interface MockAttachmentBegin { + uploadId: UploadId; + chunkBytes: number; + state: 'ready' | 'exists'; +} + function createMockTerminal( terminalId: string, opts: { @@ -237,6 +265,15 @@ export class DevMockHost { private readonly installedAssets = new Set(); private readonly cleanGitWorkspaces = new Set(); private readonly createdGitBranches = new Map>(); + private readonly attachmentUploads = new Map(); + private readonly attachmentBlobs = new Map(); + private readonly attachmentRecords = new Map< + string, + { blobId: BlobId; sizeBytes: number; name: string } + >(); + private readonly attachmentBegins = new Map(); + private uploadSeq = 0; + private attachmentSeq = 0; constructor(private readonly transport: Transport) { this.terminals.set( @@ -453,6 +490,25 @@ export class DevMockHost { await wait(CONTROL_LATENCY_MS); this.hostResource(p.clientReqId, p.resourceId); break; + case 'attachment.upload.begin': + await wait(CONTROL_LATENCY_MS); + this.beginAttachmentUpload(p); + break; + case 'attachment.upload.chunk': + this.chunkAttachmentUpload(p); + break; + case 'attachment.upload.commit': + await wait(CONTROL_LATENCY_MS); + await this.commitAttachmentUpload(p); + break; + case 'attachment.upload.abort': + await wait(CONTROL_LATENCY_MS); + this.abortAttachmentUpload(p); + break; + case 'attachment.read': + await wait(CONTROL_LATENCY_MS); + this.readAttachment(p); + break; case 'config.get': await wait(CONTROL_LATENCY_MS); this.send({ @@ -1857,6 +1913,196 @@ export class DevMockHost { }); } + private beginAttachmentUpload( + payload: Extract, + ): void { + if (payload.operationId !== undefined) { + const replayed = this.attachmentBegins.get(payload.operationId); + if (replayed) { + this.send({ + kind: 'attachment.upload.begun', + replyTo: payload.clientReqId, + ...replayed, + }); + return; + } + } + const existing = this.attachmentBlobs.get(payload.declaredSha256); + const state = existing?.byteLength === payload.declaredSize ? 'exists' : 'ready'; + this.uploadSeq += 1; + const uploadId = UploadIdSchema.parse(`upl-mock-${this.uploadSeq}`); + const bytes = + existing !== undefined && state === 'exists' + ? existing + : new Uint8Array(payload.declaredSize); + this.attachmentUploads.set(uploadId, { + uploadId, + declaredSha256: payload.declaredSha256, + declaredSize: payload.declaredSize, + name: payload.name, + mimeType: payload.mimeType, + attachmentKind: payload.attachmentKind, + received: state === 'exists' ? payload.declaredSize : 0, + bytes, + state, + blobId: state === 'exists' ? blobIdFromSha256(payload.declaredSha256) : undefined, + }); + const begun: MockAttachmentBegin = { + uploadId, + chunkBytes: ATTACHMENT_UPLOAD_CHUNK_BYTES, + state, + }; + if (payload.operationId !== undefined) this.attachmentBegins.set(payload.operationId, begun); + this.send({ kind: 'attachment.upload.begun', replyTo: payload.clientReqId, ...begun }); + } + + private chunkAttachmentUpload( + payload: Extract, + ): void { + const upload = this.attachmentUploads.get(payload.uploadId); + if (!upload) { + this.sendFailure(payload.clientReqId, 'Upload not found', { code: 'not_found' }); + return; + } + if (upload.state === 'exists') { + this.sendFailure(payload.clientReqId, 'Blob already stored; commit without chunks', { + code: 'invalid_request', + }); + return; + } + if (payload.offset !== upload.received) { + this.sendFailure( + payload.clientReqId, + `Expected offset ${upload.received}, got ${payload.offset}`, + { code: 'invalid_request' }, + ); + return; + } + const windowEnd = + upload.received + ATTACHMENT_UPLOAD_WINDOW_CHUNKS * ATTACHMENT_UPLOAD_CHUNK_BYTES; + if (payload.offset >= windowEnd && upload.received < upload.declaredSize) { + this.sendFailure(payload.clientReqId, 'Chunk is outside the credit window', { + code: 'invalid_request', + }); + return; + } + const chunk = mockBase64ToBytes(payload.data); + if (upload.received + chunk.byteLength > upload.declaredSize) { + this.sendFailure(payload.clientReqId, 'Chunk exceeds the declared size', { + code: 'invalid_request', + }); + return; + } + upload.bytes.set(chunk, payload.offset); + upload.received += chunk.byteLength; + this.send({ + kind: 'attachment.upload.chunk.acked', + replyTo: payload.clientReqId, + uploadId: payload.uploadId, + receivedBytes: upload.received, + }); + } + + private async commitAttachmentUpload( + payload: Extract, + ): Promise { + const upload = this.attachmentUploads.get(payload.uploadId); + if (!upload) { + this.sendFailure(payload.clientReqId, 'Upload not found', { code: 'not_found' }); + return; + } + if (upload.attachmentId !== undefined && upload.blobId !== undefined) { + this.send({ + kind: 'attachment.upload.committed', + replyTo: payload.clientReqId, + attachmentId: upload.attachmentId, + blobId: upload.blobId, + }); + return; + } + if (upload.state === 'ready' && upload.received !== upload.declaredSize) { + this.sendFailure( + payload.clientReqId, + `Uploaded ${upload.received} bytes, declared ${upload.declaredSize}`, + { code: 'invalid_request' }, + ); + return; + } + if (upload.state === 'ready') { + const digest = await mockSha256Hex(upload.bytes); + if (digest !== upload.declaredSha256) { + this.sendFailure(payload.clientReqId, 'Uploaded bytes do not match the declared SHA-256', { + code: 'invalid_request', + }); + return; + } + this.attachmentBlobs.set(upload.declaredSha256, upload.bytes); + } + this.attachmentSeq += 1; + const attachmentId = AttachmentIdSchema.parse(`att-mock-${this.attachmentSeq}`); + const blobId = blobIdFromSha256(upload.declaredSha256); + upload.attachmentId = attachmentId; + upload.blobId = blobId; + this.attachmentRecords.set(attachmentId, { + blobId, + sizeBytes: upload.declaredSize, + name: upload.name, + }); + this.send({ + kind: 'attachment.upload.committed', + replyTo: payload.clientReqId, + attachmentId, + blobId, + }); + } + + private abortAttachmentUpload( + payload: Extract, + ): void { + if (!this.attachmentUploads.has(payload.uploadId)) { + this.sendFailure(payload.clientReqId, 'Upload not found', { code: 'not_found' }); + return; + } + this.attachmentUploads.delete(payload.uploadId); + this.sendSuccess(payload.clientReqId); + } + + private readAttachment(payload: Extract): void { + if (!this.sessions.has(payload.sessionId)) { + this.sendFailure(payload.clientReqId, 'Attachment not found', { code: 'not_found' }); + return; + } + const record = this.attachmentRecords.get(payload.attachmentId); + if (!record) { + this.sendFailure(payload.clientReqId, 'Attachment not found', { code: 'not_found' }); + return; + } + const hex = record.blobId.slice('sha256:'.length); + const bytes = this.attachmentBlobs.get(hex); + if (!bytes) { + this.sendFailure(payload.clientReqId, 'Attachment bytes are missing', { code: 'not_found' }); + return; + } + if (payload.offset > bytes.byteLength) { + this.sendFailure(payload.clientReqId, 'Read offset is past the end of the attachment', { + code: 'invalid_request', + }); + return; + } + const slice = bytes.subarray(payload.offset, payload.offset + payload.length); + this.send({ + kind: 'attachment.read.result', + replyTo: payload.clientReqId, + sessionId: payload.sessionId, + attachmentId: payload.attachmentId, + blobId: record.blobId, + offset: payload.offset, + data: mockBytesToBase64(slice), + sizeBytes: bytes.byteLength, + eof: payload.offset + slice.byteLength >= bytes.byteLength, + }); + } + private send(payload: WirePayload): void { this.transport.send(createWireMessage(payload)); } @@ -1868,7 +2114,7 @@ export class DevMockHost { private sendFailure( replyTo: string, message: string, - reporting: { reportedInConversation?: true } = {}, + reporting: { reportedInConversation?: true; code?: string } = {}, ): void { this.send({ kind: 'request.failed', replyTo, message, ...reporting }); } @@ -1973,6 +2219,35 @@ function isRunningTurn(session: MockSession, epoch: number): boolean { return session.epoch === epoch && session.status === 'running'; } +function mockBase64ToBytes(data: string): Uint8Array { + const binary = atob(data); + const bytes = new Uint8Array(binary.length); + for (let i = 0, len = binary.length; i < len; i++) { + bytes[i] = binary.charCodeAt(i); + } + return bytes; +} + +function mockBytesToBase64(bytes: Uint8Array): string { + let binary = ''; + for (let i = 0, len = bytes.byteLength; i < len; i++) { + binary += String.fromCharCode(bytes[i]); + } + return btoa(binary); +} + +async function mockSha256Hex(bytes: Uint8Array): Promise { + const buffer = new ArrayBuffer(bytes.byteLength); + new Uint8Array(buffer).set(bytes); + const digest = await crypto.subtle.digest('SHA-256', buffer); + const view = new Uint8Array(digest); + let hex = ''; + for (let i = 0, len = view.byteLength; i < len; i++) { + hex += view[i].toString(16).padStart(2, '0'); + } + return hex; +} + async function waitForShowcaseStep(session: MockSession, epoch: number): Promise { await wait(SHOWCASE_SCRIPT_STEP_LATENCY_MS); return isRunningTurn(session, epoch); From a55f6eeda46b9090c9ad3165fb23c7e9aad0c2a0 Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 19:05:02 +0800 Subject: [PATCH 07/11] fix(engine): accept pipelined chunks, verify dedupe hits, and release abandoned stages --- .../src/__tests__/attachment-upload.test.ts | 177 +++++++++++++++++- .../src/__tests__/engine-attachments.test.ts | 72 +++++++ .../engine/src/attachment/upload-service.ts | 93 ++++++--- packages/host/engine/src/resource/service.ts | 4 +- 4 files changed, 312 insertions(+), 34 deletions(-) diff --git a/packages/host/engine/src/__tests__/attachment-upload.test.ts b/packages/host/engine/src/__tests__/attachment-upload.test.ts index 9d3640c1b..3a892e748 100644 --- a/packages/host/engine/src/__tests__/attachment-upload.test.ts +++ b/packages/host/engine/src/__tests__/attachment-upload.test.ts @@ -1,13 +1,19 @@ import { createHash } from 'node:crypto'; -import { mkdtemp, rm } from 'node:fs/promises'; +import { mkdtemp, readdir, rm, writeFile } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; import type { AttachmentId, UploadId } from '@linkcode/schema'; -import { AttachmentIdSchema, blobIdFromSha256, SessionIdSchema } from '@linkcode/schema'; +import { + ATTACHMENT_UPLOAD_CHUNK_BYTES, + AttachmentIdSchema, + blobIdFromSha256, + SessionIdSchema, +} from '@linkcode/schema'; import { Effect } from 'effect'; import { afterEach, describe, expect, it } from 'vitest'; import { InMemoryAttachmentStore } from '../attachment/attachment-store'; import { FsBlobStore } from '../attachment/blob-store'; +import { UPLOAD_LEASE_TTL_MS } from '../attachment/gc'; import { AttachmentUploadService } from '../attachment/upload-service'; const temporaryDirectories: string[] = []; @@ -24,7 +30,7 @@ function sha256(bytes: Uint8Array): string { return createHash('sha256').update(bytes).digest('hex'); } -async function makeService(reachableIds: AttachmentId[] = []) { +async function makeService(reachableIds: AttachmentId[] = [], clock?: () => number) { const root = await mkdtemp(join(tmpdir(), 'linkcode-upload-')); temporaryDirectories.push(root); const reachable = new Set(reachableIds); @@ -33,8 +39,21 @@ async function makeService(reachableIds: AttachmentId[] = []) { () => reachable, (sid, id) => sid === sessionId && reachable.has(id), ); - const uploads = new AttachmentUploadService(blobs, attachments); - return { attachments, blobs, reachable, uploads }; + const uploads = new AttachmentUploadService(blobs, attachments, undefined, clock); + return { attachments, blobs, reachable, root, uploads }; +} + +function stagingEntries(root: string): Promise { + return readdir(join(root, 'blobs', 'tmp')); +} + +function chunksOf(bytes: Buffer): Array<{ offset: number; data: string }> { + const chunks: Array<{ offset: number; data: string }> = []; + for (let offset = 0; offset < bytes.byteLength; offset += ATTACHMENT_UPLOAD_CHUNK_BYTES) { + const end = Math.min(offset + ATTACHMENT_UPLOAD_CHUNK_BYTES, bytes.byteLength); + chunks.push({ offset, data: bytes.subarray(offset, end).toString('base64') }); + } + return chunks; } async function run(effect: Effect.Effect): Promise { @@ -188,6 +207,154 @@ describe('AttachmentUploadService', () => { ).rejects.toMatchObject({ _tag: 'RequestError', code: 'not_found' }); }); + it('accepts the whole credit window in flight, the way the client sends it', async () => { + const { uploads } = await makeService(); + // Three chunks so the client's two-chunk window overlaps a write on both sides. + const bytes = Buffer.alloc(ATTACHMENT_UPLOAD_CHUNK_BYTES * 2 + 11, 7); + const begun = await run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'window.bin', + attachmentKind: 'file', + }), + ); + const acks = await Promise.all( + chunksOf(bytes).map((chunk) => run(uploads.chunk(begun.uploadId, chunk.offset, chunk.data))), + ); + expect(acks.at(-1)?.receivedBytes).toBe(bytes.byteLength); + const committed = await run(uploads.commit(begun.uploadId)); + expect(committed.blobId).toBe(blobIdFromSha256(sha256(bytes))); + }); + + it('refuses a dedupe hit whose stored bytes are not the declared size', async () => { + const { uploads } = await makeService(); + const bytes = Buffer.from('eleven byte'); + const stored = await run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'honest.bin', + attachmentKind: 'file', + }), + ); + await run(uploads.chunk(stored.uploadId, 0, bytes.toString('base64'))); + await run(uploads.commit(stored.uploadId)); + + // Same hash, a size that never matched the bytes, and no chunks sent at all. + const lying = await run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: 4096, + name: 'liar.bin', + attachmentKind: 'file', + }), + ); + expect(lying.state).toBe('ready'); + await expect(run(uploads.commit(lying.uploadId))).rejects.toMatchObject({ + _tag: 'RequestError', + code: 'invalid_request', + }); + }); + + it('keeps a rowed blob on disk when a same-hash upload fails to publish its record', async () => { + const { attachments, blobs, reachable, uploads } = await makeService(); + const bytes = Buffer.from('the same screenshot twice'); + const input = { + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + attachmentKind: 'file', + }; + // Both begins land before either commit, so neither sees a blob row and both stage bytes. + const a = await run(uploads.begin({ ...input, name: 'a.png' })); + const b = await run(uploads.begin({ ...input, name: 'b.png' })); + await run(uploads.chunk(a.uploadId, 0, bytes.toString('base64'))); + await run(uploads.chunk(b.uploadId, 0, bytes.toString('base64'))); + const first = await run(uploads.commit(a.uploadId)); + reachable.add(first.attachmentId); + + const commitAttachment = attachments.commitAttachment.bind(attachments); + attachments.commitAttachment = () => Promise.reject(new Error('row insert failed')); + await expect(run(uploads.commit(b.uploadId))).rejects.toBeDefined(); + attachments.commitAttachment = commitAttachment; + + expect(await blobs.stat(first.blobId)).toEqual({ sizeBytes: bytes.byteLength }); + const page = await run(uploads.read(sessionId, first.attachmentId, 0, bytes.byteLength)); + expect(Buffer.from(page.data, 'base64').toString()).toBe('the same screenshot twice'); + }); + + it('reads no further than the bytes on disk and says so typed', async () => { + const { blobs, reachable, uploads } = await makeService(); + const bytes = Buffer.from('full length payload'); + const begun = await run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'truncated.bin', + attachmentKind: 'file', + }), + ); + await run(uploads.chunk(begun.uploadId, 0, bytes.toString('base64'))); + const committed = await run(uploads.commit(begun.uploadId)); + reachable.add(committed.attachmentId); + // Bitrot or a half-finished GC unlink: the row outlives some of its bytes. + const path = blobs.pathOf(committed.blobId); + await rm(path, { force: true }); + await writeFile(path, bytes.subarray(0, 4)); + + await expect( + run(uploads.read(sessionId, committed.attachmentId, 4, bytes.byteLength)), + ).rejects.toMatchObject({ _tag: 'RequestError', code: 'not_found' }); + }); + + it('releases the staging handle of an expired lease and of a rejected commit', async () => { + let now = 1_000; + const { blobs, root, uploads } = await makeService([], () => now); + const bytes = Buffer.from('abandoned'); + const abandoned = await run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'abandoned.bin', + attachmentKind: 'file', + }), + ); + expect(await stagingEntries(root)).toEqual([abandoned.uploadId]); + + // A commit that cannot succeed ends the upload: no resume frame exists to continue it. + const doomed = await run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'doomed.bin', + attachmentKind: 'file', + }), + ); + await expect(run(uploads.commit(doomed.uploadId))).rejects.toMatchObject({ + _tag: 'RequestError', + code: 'invalid_request', + }); + await expect( + run(uploads.chunk(doomed.uploadId, 0, bytes.toString('base64'))), + ).rejects.toMatchObject({ _tag: 'RequestError', code: 'conflict' }); + + now += UPLOAD_LEASE_TTL_MS + 1; + const fresh = await run( + uploads.begin({ + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'fresh.bin', + attachmentKind: 'file', + }), + ); + expect(await stagingEntries(root)).toEqual([fresh.uploadId]); + // The lease row outlives the handle until the GC's own sweep; the id is known but dead. + await expect( + run(uploads.chunk(abandoned.uploadId, 0, bytes.toString('base64'))), + ).rejects.toMatchObject({ _tag: 'RequestError', code: 'conflict' }); + expect(await blobs.stat(blobIdFromSha256(sha256(bytes)))).toBeUndefined(); + }); + it('rejects a sniffed image whose bytes are not that type', async () => { const { uploads } = await makeService(); const bytes = Buffer.from('not a png'); diff --git a/packages/host/engine/src/__tests__/engine-attachments.test.ts b/packages/host/engine/src/__tests__/engine-attachments.test.ts index 4475db9e1..272fa6359 100644 --- a/packages/host/engine/src/__tests__/engine-attachments.test.ts +++ b/packages/host/engine/src/__tests__/engine-attachments.test.ts @@ -165,4 +165,76 @@ describe('engine attachment upload/read', () => { ); }); }); + + it('accepts two chunk frames delivered before the first write lands', async () => { + const stateDir = await tempDirectory(); + const h = createSessionHarness( + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + { + stateDir, + }, + ); + await h.engine.start(); + const bytes = Buffer.alloc(ATTACHMENT_UPLOAD_CHUNK_BYTES + 11, 9); + await h.inject({ + kind: 'attachment.upload.begin', + clientReqId: 'begin', + declaredSha256: sha256(bytes), + declaredSize: bytes.byteLength, + name: 'window.bin', + attachmentKind: 'file', + }); + await vi.waitFor(() => { + expect(h.sent).toContainEqual( + expect.objectContaining({ kind: 'attachment.upload.begun', replyTo: 'begin' }), + ); + }); + const begun = replyOf(h.sent, 'attachment.upload.begun', 'begin'); + + // Each frame is handled in its own fiber, so both are in flight before either write resolves. + await Promise.all([ + h.inject({ + kind: 'attachment.upload.chunk', + clientReqId: 'chunk-0', + uploadId: begun.uploadId, + offset: 0, + data: bytes.subarray(0, ATTACHMENT_UPLOAD_CHUNK_BYTES).toString('base64'), + }), + h.inject({ + kind: 'attachment.upload.chunk', + clientReqId: 'chunk-1', + uploadId: begun.uploadId, + offset: ATTACHMENT_UPLOAD_CHUNK_BYTES, + data: bytes.subarray(ATTACHMENT_UPLOAD_CHUNK_BYTES).toString('base64'), + }), + ]); + await vi.waitFor(() => { + expect(h.sent).toContainEqual( + expect.objectContaining({ + kind: 'attachment.upload.chunk.acked', + replyTo: 'chunk-1', + receivedBytes: bytes.byteLength, + }), + ); + }); + + await h.inject({ + kind: 'attachment.upload.commit', + clientReqId: 'commit', + uploadId: begun.uploadId, + }); + await vi.waitFor(() => { + expect(h.sent).toContainEqual( + expect.objectContaining({ kind: 'attachment.upload.committed', replyTo: 'commit' }), + ); + }); + expect(replyOf(h.sent, 'attachment.upload.committed', 'commit').blobId).toBe( + `sha256:${sha256(bytes)}`, + ); + }); }); diff --git a/packages/host/engine/src/attachment/upload-service.ts b/packages/host/engine/src/attachment/upload-service.ts index 9f2b0d60f..aaecf2853 100644 --- a/packages/host/engine/src/attachment/upload-service.ts +++ b/packages/host/engine/src/attachment/upload-service.ts @@ -2,7 +2,6 @@ import { randomUUID } from 'node:crypto'; import type { AttachmentId, BlobId, SessionId, UploadId, UploadLease } from '@linkcode/schema'; import { ATTACHMENT_UPLOAD_CHUNK_BYTES, - ATTACHMENT_UPLOAD_WINDOW_CHUNKS, AttachmentIdSchema, blobIdFromSha256, MAX_ATTACHMENT_BYTES, @@ -59,6 +58,9 @@ export interface AttachmentReadResult { interface LiveUpload { readonly lease: UploadLease; readonly stage: BlobStage | undefined; + /** Chunk frames arrive in order but each is handled in its own fiber; the contiguity check and + * the write that advances `receivedBytes` must not interleave across the client's credit window. */ + readonly gate: AttachmentIoMutex; receivedBytes: number; head: Uint8Array; readonly state: 'ready' | 'exists'; @@ -80,6 +82,7 @@ export class AttachmentUploadService { ): Effect.Effect { const store = this.store.bind(this); const files = this.files.bind(this); + const reapExpired = this.reapExpired.bind(this); return Effect.gen({ self: this }, function* () { if (input.operationId !== undefined) { const replayed = this.begunByOperation.get(input.operationId); @@ -88,6 +91,9 @@ export class AttachmentUploadService { if (input.declaredSize > MAX_ATTACHMENT_BYTES) { return yield* invalid('limit_exceeded', 'Attachment exceeds the 8 MiB limit'); } + // A staging handle lives in `live` until commit or abort; a client that vanishes mid-upload + // never sends either, so expired leases release their descriptors here. + yield* files('reap', reapExpired); const uploadId = UploadIdSchema.parse(`upl-${randomUUID()}`); const now = this.clock(); const lease = yield* store('begin', () => @@ -121,6 +127,7 @@ export class AttachmentUploadService { this.live.set(uploadId, { lease, stage, + gate: new AttachmentIoMutex(), receivedBytes: state === 'exists' ? input.declaredSize : 0, head: new Uint8Array(0), state, @@ -154,28 +161,27 @@ export class AttachmentUploadService { } return yield* invalid('conflict', 'Upload is not accepting chunks; retry from begin'); } - if (offset !== live.receivedBytes) { - return yield* invalid( - 'invalid_request', - `Expected offset ${live.receivedBytes}, got ${offset}`, - ); - } - const windowEnd = - live.receivedBytes + ATTACHMENT_UPLOAD_WINDOW_CHUNKS * ATTACHMENT_UPLOAD_CHUNK_BYTES; - if (offset >= windowEnd && live.receivedBytes < live.lease.declaredSize) { - return yield* invalid('invalid_request', 'Chunk is outside the credit window'); - } - const bytes = yield* Effect.try({ - try: () => decodeChunk(data), - catch: (cause) => mapCause(cause, 'store', 'chunk'), - }); - if (live.receivedBytes + bytes.byteLength > live.lease.declaredSize) { - return yield* invalid('invalid_request', 'Chunk exceeds the declared size'); - } - yield* files('write', () => stage.write(offset, bytes)); - if (offset === 0) live.head = bytes.subarray(0, Math.min(HEAD_BYTES, bytes.byteLength)); - live.receivedBytes += bytes.byteLength; - return { uploadId, receivedBytes: live.receivedBytes }; + return yield* files('write', () => + live.gate.run(async () => { + if (offset !== live.receivedBytes) { + throw new RequestError({ + code: 'invalid_request', + message: `Expected offset ${live.receivedBytes}, got ${offset}`, + }); + } + const bytes = decodeChunk(data); + if (live.receivedBytes + bytes.byteLength > live.lease.declaredSize) { + throw new RequestError({ + code: 'invalid_request', + message: 'Chunk exceeds the declared size', + }); + } + await stage.write(offset, bytes); + if (offset === 0) live.head = bytes.subarray(0, Math.min(HEAD_BYTES, bytes.byteLength)); + live.receivedBytes += bytes.byteLength; + return { uploadId, receivedBytes: live.receivedBytes }; + }), + ); }); } @@ -183,6 +189,7 @@ export class AttachmentUploadService { const store = this.store.bind(this); const files = this.files.bind(this); const forget = this.forget.bind(this); + const discard = this.discard.bind(this); const publishExists = this.publishExists.bind(this); const publishReady = this.publishReady.bind(this); return Effect.gen({ self: this }, function* () { @@ -200,14 +207,16 @@ export class AttachmentUploadService { pinnedBlobId === undefined ? undefined : yield* files('stat', () => this.blobs.stat(pinnedBlobId)); - const exists = live?.state === 'exists' || existsFile !== undefined; + // A dedupe hit only counts when the stored bytes are the size the client declared; otherwise + // the declared size is a lie and the exists path would skip every coverage check. + const exists = live?.state === 'exists' || existsFile?.sizeBytes === lease.declaredSize; if (exists) { const blobId = pinnedBlobId ?? blobIdFromSha256(lease.declaredSha256); const head = (yield* files('read', () => this.blobs.read(blobId, 0, HEAD_BYTES))) ?? new Uint8Array(0); yield* assertMime(lease.mimeType, head); const committed = yield* publishExists(lease, blobId); - forget(uploadId); + yield* files('discard', () => discard(uploadId)); return committed; } if (!live?.stage) { @@ -223,7 +232,11 @@ export class AttachmentUploadService { const committed = yield* publishReady(live); forget(uploadId); return committed; - }); + }).pipe( + // v1 has no resume frame, so a rejected commit ends the upload: release its staging handle + // instead of leaving a dead entry a later chunk would fail on. + Effect.tapError(() => files('discard', () => discard(uploadId)).pipe(Effect.ignore)), + ); } abort(uploadId: UploadId): Effect.Effect { @@ -269,6 +282,11 @@ export class AttachmentUploadService { (yield* files('read', () => this.blobs.read(attachment.blobId, offset, length))) ?? undefined; if (bytes === undefined) return yield* invalid('not_found', 'Attachment bytes are missing'); + // `length` is positive on the wire, so an empty read below the recorded size means the blob + // is shorter than its row. Saying `eof` there would stall the caller's walk forever. + if (bytes.byteLength === 0 && offset < attachment.sizeBytes) { + return yield* invalid('not_found', 'Attachment bytes are truncated'); + } const end = offset + bytes.byteLength; return { sessionId, @@ -324,7 +342,9 @@ export class AttachmentUploadService { }); } catch (error) { await stage.abort().catch(noop); - if (blobId) await blobs.delete(blobId); + // Content addressing means a concurrent upload of the same bytes may already own a row + // for this blob; unlinking then would strand its attachment. + if (blobId && !(await attachments.getBlob(blobId))) await blobs.delete(blobId); throw error; } }); @@ -333,7 +353,24 @@ export class AttachmentUploadService { }); } - private forget(uploadId: UploadId): void { + private async discard(uploadId: string): Promise { + const live = this.live.get(uploadId); + this.forget(uploadId); + await live?.stage?.abort().catch(noop); + } + + private async reapExpired(): Promise { + const now = this.clock(); + const dead: BlobStage[] = []; + for (const [uploadId, live] of this.live) { + if (live.lease.expiresAt > now) continue; + if (live.stage) dead.push(live.stage); + this.forget(uploadId); + } + await Promise.all(dead.map((stage) => stage.abort().catch(noop))); + } + + private forget(uploadId: string): void { this.live.delete(uploadId); for (const [operationId, begun] of this.begunByOperation) { if (begun.uploadId === uploadId) this.begunByOperation.delete(operationId); diff --git a/packages/host/engine/src/resource/service.ts b/packages/host/engine/src/resource/service.ts index 0f6726c6b..4219a0322 100644 --- a/packages/host/engine/src/resource/service.ts +++ b/packages/host/engine/src/resource/service.ts @@ -127,7 +127,9 @@ export class ResourceService { }); } catch (error) { await stage.abort().catch(noop); - await blobs.delete(blobId); + // Content addressing means another attachment may already own a row for this blob; + // unlinking then would strand its bytes. + if (!(await attachments.getBlob(blobId))) await blobs.delete(blobId); throw error; } }); From 36b22ce1eb322d85ea8cf2a067e5ad9ca36c1c82 Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 19:06:29 +0800 Subject: [PATCH 08/11] fix(client-core): terminate the attachment read walk, abort a failed upload, bound the blob cache --- packages/client/core/src/client.ts | 6 +- .../core/src/client/attachment-channel.ts | 48 ++-- packages/client/core/src/client/blob-cache.ts | 41 +++- .../integration/attachment-client.test.ts | 212 +++++++++++++++++- 4 files changed, 289 insertions(+), 18 deletions(-) diff --git a/packages/client/core/src/client.ts b/packages/client/core/src/client.ts index 423e09b36..a9c371ff3 100644 --- a/packages/client/core/src/client.ts +++ b/packages/client/core/src/client.ts @@ -94,6 +94,7 @@ import type { AttachmentReadBytes, } from './client/attachment-channel'; import { AttachmentChannel } from './client/attachment-channel'; +import type { Sha256Hex } from './client/blob-cache'; import type { BrowserCommandExecutor } from './client/browser-host-channel'; import { BrowserHostChannel } from './client/browser-host-channel'; import type { @@ -125,6 +126,7 @@ export type { AttachmentPutInput, AttachmentReadBytes, } from './client/attachment-channel'; +export type { Sha256Hex } from './client/blob-cache'; export type { BrowserCommandExecutor } from './client/browser-host-channel'; export type { ConversationReadClientOptions, @@ -160,6 +162,8 @@ type TerminalReplayTruncatedCb = (truncated: boolean) => void; export interface LinkCodeClientOptions { randomUUID?: RandomUUID; + /** Attachment upload hashes its bytes; hosts without `crypto.subtle` must supply the digest. */ + sha256Hex?: Sha256Hex; } export interface TerminalAttachResult { @@ -292,7 +296,7 @@ export class LinkCodeClient { const randomUUID = resolveRandomUUID(options.randomUUID); this.pending = new PendingRegistry(randomUUID); this.control = new ControlChannel(transport, this.pending); - this.attachments = new AttachmentChannel(transport, this.pending); + this.attachments = new AttachmentChannel(transport, this.pending, options.sha256Hex); this.terminals = new TerminalChannel(transport, this.pending, randomUUID); this.browserHost = new BrowserHostChannel(transport, this.pending, randomUUID); this.agentLogin = new AgentLoginChannel(transport, this.pending); diff --git a/packages/client/core/src/client/attachment-channel.ts b/packages/client/core/src/client/attachment-channel.ts index abdb212da..f0d2c2f2a 100644 --- a/packages/client/core/src/client/attachment-channel.ts +++ b/packages/client/core/src/client/attachment-channel.ts @@ -2,6 +2,7 @@ import type { AttachmentId, OperationId, SessionId, UploadId } from '@linkcode/s import { ATTACHMENT_UPLOAD_CHUNK_BYTES, ATTACHMENT_UPLOAD_WINDOW_CHUNKS } from '@linkcode/schema'; import type { Transport } from '@linkcode/transport'; import { noop } from 'foxts/noop'; +import type { Sha256Hex } from './blob-cache'; import { AttachmentBlobCache, base64ToBytes, bytesToBase64, sha256Hex } from './blob-cache'; import type { AttachmentChunkAck, @@ -39,6 +40,7 @@ export class AttachmentChannel { constructor( private readonly transport: Transport, private readonly pending: PendingRegistry, + private readonly digest: Sha256Hex = sha256Hex, ) {} beginUpload(input: AttachmentBeginInput): Promise { @@ -98,7 +100,7 @@ export class AttachmentChannel { /** Hash, begin, windowed chunks, commit. Identical bytes short-circuit to `exists`. */ async put(input: AttachmentPutInput): Promise { - const declaredSha256 = await sha256Hex(input.bytes); + const declaredSha256 = await this.digest(input.bytes); const begun = await this.beginUpload({ declaredSha256, declaredSize: input.bytes.byteLength, @@ -107,10 +109,17 @@ export class AttachmentChannel { attachmentKind: input.attachmentKind, operationId: input.operationId, }); - if (begun.state === 'ready') { - await this.sendWindowed(begun.uploadId, begun.chunkBytes, input.bytes); + let committed: AttachmentCommitResult; + try { + if (begun.state === 'ready') { + await this.sendWindowed(begun.uploadId, begun.chunkBytes, input.bytes); + } + committed = await this.commit(begun.uploadId); + } catch (error) { + // Release the lease and the daemon's staging file now; the TTL reaper is a 24h backstop. + await this.abort(begun.uploadId).catch(noop); + throw error; } - const committed = await this.commit(begun.uploadId); this.cache.set(committed.blobId, input.bytes); return committed; } @@ -123,16 +132,26 @@ export class AttachmentChannel { return { blobId: first.blobId, bytes: cached, sizeBytes: first.sizeBytes }; } const bytes = new Uint8Array(first.sizeBytes); - const firstSlice = base64ToBytes(first.data); - bytes.set(firstSlice, first.offset); - let offset = first.offset + firstSlice.byteLength; + let page = first; + let offset = 0; while (offset < first.sizeBytes) { - // eslint-disable-next-line no-await-in-loop -- sequential pages of one attachment - const page = await this.read(sessionId, attachmentId, offset, ATTACHMENT_UPLOAD_CHUNK_BYTES); const slice = base64ToBytes(page.data); - bytes.set(slice, page.offset); - offset = page.offset + slice.byteLength; - if (page.eof) break; + // A page that repeats an offset, returns nothing, or overruns the recorded size cannot be + // assembled — without this the walk never advances and zero-fills what it could not read. + if ( + page.offset !== offset || + slice.byteLength === 0 || + offset + slice.byteLength > bytes.byteLength + ) { + throw new Error( + `Attachment ${attachmentId} returned ${slice.byteLength} bytes at ${page.offset}, expected more at ${offset} of ${first.sizeBytes}`, + ); + } + bytes.set(slice, offset); + offset += slice.byteLength; + if (offset >= first.sizeBytes) break; + // eslint-disable-next-line no-await-in-loop -- sequential pages of one attachment + page = await this.read(sessionId, attachmentId, offset, ATTACHMENT_UPLOAD_CHUNK_BYTES); } this.cache.set(first.blobId, bytes); return { blobId: first.blobId, bytes, sizeBytes: first.sizeBytes }; @@ -155,7 +174,10 @@ export class AttachmentChannel { const at = offset; const data = bytesToBase64(bytes.subarray(at, end)); offset = end; - acks.push(this.sendChunk(uploadId, at, data).then(noop)); + const ack = this.sendChunk(uploadId, at, data).then(noop); + // The awaits below surface the first failure; the rest must not become unhandled rejections. + ack.catch(noop); + acks.push(ack); } for (let i = 0, len = acks.length; i < len; i++) { // eslint-disable-next-line no-await-in-loop -- drain remaining acks diff --git a/packages/client/core/src/client/blob-cache.ts b/packages/client/core/src/client/blob-cache.ts index 70c845cb3..d27e28de9 100644 --- a/packages/client/core/src/client/blob-cache.ts +++ b/packages/client/core/src/client/blob-cache.ts @@ -1,13 +1,39 @@ -/** Decoded attachment bytes keyed by content-addressed `blobId`. Blobs are immutable. */ +import { nullthrow } from 'foxts/guard'; + +/** Decoded bytes this cache will hold before evicting the least recently read blob. */ +export const ATTACHMENT_CACHE_MAX_BYTES = 64 * 1024 * 1024; + +/** + * Decoded attachment bytes keyed by content-addressed `blobId`. Blobs are immutable, so a hit is + * always correct; the byte budget exists because this cache outlives every session it serves. + */ export class AttachmentBlobCache { private readonly blobs = new Map(); + private totalBytes = 0; + + constructor(private readonly maxBytes: number = ATTACHMENT_CACHE_MAX_BYTES) {} get(blobId: string): Uint8Array | undefined { - return this.blobs.get(blobId); + const bytes = this.blobs.get(blobId); + // Map iterates in insertion order, so re-inserting a hit makes eviction least-recently-used. + if (bytes) { + this.blobs.delete(blobId); + this.blobs.set(blobId, bytes); + } + return bytes; } set(blobId: string, bytes: Uint8Array): void { + const previous = this.blobs.get(blobId); + if (previous) this.totalBytes -= previous.byteLength; + this.blobs.delete(blobId); this.blobs.set(blobId, bytes); + this.totalBytes += bytes.byteLength; + for (const [oldest, stale] of this.blobs) { + if (oldest === blobId || this.totalBytes <= this.maxBytes) break; + this.blobs.delete(oldest); + this.totalBytes -= stale.byteLength; + } } has(blobId: string): boolean { @@ -42,8 +68,17 @@ export function base64ToBytes(data: string): Uint8Array { return bytes; } +/** A host without `crypto.subtle` (React Native) injects its own digest, as it already does for + * `randomUUID`. */ +export type Sha256Hex = (bytes: Uint8Array) => Promise; + export async function sha256Hex(bytes: Uint8Array): Promise { - const digest = await crypto.subtle.digest('SHA-256', arrayBufferOf(bytes)); + const subtle = (Reflect.get(globalThis, 'crypto') as { subtle?: SubtleCrypto } | undefined) + ?.subtle; + const digest = await nullthrow( + subtle, + 'LinkCodeClient: no crypto.subtle — pass options.sha256Hex', + ).digest('SHA-256', arrayBufferOf(bytes)); const view = new Uint8Array(digest); let hex = ''; for (let i = 0, len = view.byteLength; i < len; i++) { diff --git a/packages/client/core/tests/integration/attachment-client.test.ts b/packages/client/core/tests/integration/attachment-client.test.ts index 7678eed8f..dee65cb4b 100644 --- a/packages/client/core/tests/integration/attachment-client.test.ts +++ b/packages/client/core/tests/integration/attachment-client.test.ts @@ -1,5 +1,6 @@ import { ATTACHMENT_STORE_WIRE_VERSION, + ATTACHMENT_UPLOAD_CHUNK_BYTES, AttachmentIdSchema, BlobIdSchema, SessionIdSchema, @@ -8,7 +9,12 @@ import { import { createLocalTransportPair, createWireMessage } from '@linkcode/transport'; import { describe, expect, it } from 'vitest'; import { LinkCodeClient } from '../../src/client'; -import { base64ToBytes, bytesToBase64, sha256Hex } from '../../src/client/blob-cache'; +import { + AttachmentBlobCache, + base64ToBytes, + bytesToBase64, + sha256Hex, +} from '../../src/client/blob-cache'; import { createConnectedLocalClient } from '../support/local-client'; describe('LinkCodeClient attachment store API', () => { @@ -111,4 +117,208 @@ describe('LinkCodeClient attachment store API', () => { client.dispose(); serverTransport.close(); }); + + it('pipelines a multi-chunk upload past a receiver that only accepts the next offset', async () => { + const { client, serverTransport } = await createConnectedLocalClient(); + const bytes = new Uint8Array(ATTACHMENT_UPLOAD_CHUNK_BYTES * 2 + 9).fill(7); + const blobId = BlobIdSchema.parse(`sha256:${'c'.repeat(64)}`); + const offsets: number[] = []; + let received = 0; + + serverTransport.onMessage((message) => { + const p = message.payload; + if (p.kind === 'attachment.upload.begin') { + serverTransport.send( + createWireMessage({ + kind: 'attachment.upload.begun', + replyTo: p.clientReqId, + uploadId: UploadIdSchema.parse('upl-2'), + chunkBytes: ATTACHMENT_UPLOAD_CHUNK_BYTES, + state: 'ready', + }), + ); + } + // The daemon's contract: strictly the next contiguous offset, acked cumulatively. + if (p.kind === 'attachment.upload.chunk') { + offsets.push(p.offset); + if (p.offset !== received) { + serverTransport.send( + createWireMessage({ + kind: 'request.failed', + replyTo: p.clientReqId, + message: `Expected offset ${received}, got ${p.offset}`, + code: 'invalid_request', + }), + ); + return; + } + received += base64ToBytes(p.data).byteLength; + serverTransport.send( + createWireMessage({ + kind: 'attachment.upload.chunk.acked', + replyTo: p.clientReqId, + uploadId: p.uploadId, + receivedBytes: received, + }), + ); + } + if (p.kind === 'attachment.upload.commit') { + serverTransport.send( + createWireMessage({ + kind: 'attachment.upload.committed', + replyTo: p.clientReqId, + attachmentId: AttachmentIdSchema.parse('att-2'), + blobId, + }), + ); + } + }); + + await client.putAttachment({ bytes, name: 'big.bin', attachmentKind: 'file' }); + expect(offsets).toEqual([0, ATTACHMENT_UPLOAD_CHUNK_BYTES, ATTACHMENT_UPLOAD_CHUNK_BYTES * 2]); + expect(received).toBe(bytes.byteLength); + + client.dispose(); + serverTransport.close(); + }); + + it('assembles a multi-page read once and serves the rest from the cache', async () => { + const { client, serverTransport } = await createConnectedLocalClient(); + const bytes = new Uint8Array(ATTACHMENT_UPLOAD_CHUNK_BYTES + 31).fill(3); + const blobId = BlobIdSchema.parse(`sha256:${'d'.repeat(64)}`); + const attachmentId = AttachmentIdSchema.parse('att-3'); + const sessionId = SessionIdSchema.parse('session-1'); + let reads = 0; + + serverTransport.onMessage((message) => { + const p = message.payload; + if (p.kind !== 'attachment.read') return; + reads += 1; + const slice = bytes.subarray(p.offset, p.offset + p.length); + serverTransport.send( + createWireMessage({ + kind: 'attachment.read.result', + replyTo: p.clientReqId, + sessionId: p.sessionId, + attachmentId: p.attachmentId, + blobId, + offset: p.offset, + data: bytesToBase64(slice), + sizeBytes: bytes.byteLength, + eof: p.offset + slice.byteLength >= bytes.byteLength, + }), + ); + }); + + const first = await client.getAttachmentBytes(sessionId, attachmentId); + expect(first.bytes).toEqual(bytes); + expect(reads).toBe(2); + const again = await client.getAttachmentBytes(sessionId, attachmentId); + expect(again.bytes).toEqual(bytes); + // Only the one probe page: the rest came from the blobId cache. + expect(reads).toBe(3); + + client.dispose(); + serverTransport.close(); + }); + + it('fails the read walk instead of spinning when the bytes run out early', async () => { + const { client, serverTransport } = await createConnectedLocalClient(); + const bytes = new Uint8Array(16).fill(1); + const blobId = BlobIdSchema.parse(`sha256:${'e'.repeat(64)}`); + const attachmentId = AttachmentIdSchema.parse('att-4'); + const sessionId = SessionIdSchema.parse('session-1'); + let reads = 0; + + serverTransport.onMessage((message) => { + const p = message.payload; + if (p.kind !== 'attachment.read') return; + reads += 1; + const slice = bytes.subarray(p.offset, p.offset + p.length); + serverTransport.send( + createWireMessage({ + kind: 'attachment.read.result', + replyTo: p.clientReqId, + sessionId: p.sessionId, + attachmentId: p.attachmentId, + blobId, + offset: p.offset, + data: bytesToBase64(slice), + // A row that outlived some of its bytes: the size never becomes reachable. + sizeBytes: bytes.byteLength * 4, + eof: false, + }), + ); + }); + + await expect(client.getAttachmentBytes(sessionId, attachmentId)).rejects.toThrow('att-4'); + expect(reads).toBe(2); + + client.dispose(); + serverTransport.close(); + }); + + it('aborts the upload when a chunk is rejected', async () => { + const { client, serverTransport } = await createConnectedLocalClient(); + const bytes = new Uint8Array(ATTACHMENT_UPLOAD_CHUNK_BYTES * 2).fill(5); + const uploadId = UploadIdSchema.parse('upl-3'); + let aborted: string | undefined; + + serverTransport.onMessage((message) => { + const p = message.payload; + if (p.kind === 'attachment.upload.begin') { + serverTransport.send( + createWireMessage({ + kind: 'attachment.upload.begun', + replyTo: p.clientReqId, + uploadId, + chunkBytes: ATTACHMENT_UPLOAD_CHUNK_BYTES, + state: 'ready', + }), + ); + } + if (p.kind === 'attachment.upload.chunk') { + serverTransport.send( + createWireMessage({ + kind: 'request.failed', + replyTo: p.clientReqId, + message: 'disk is full', + code: 'invalid_request', + }), + ); + } + if (p.kind === 'attachment.upload.abort') { + aborted = p.uploadId; + serverTransport.send( + createWireMessage({ kind: 'request.succeeded', replyTo: p.clientReqId }), + ); + } + }); + + await expect( + client.putAttachment({ bytes, name: 'doomed.bin', attachmentKind: 'file' }), + ).rejects.toThrow('disk is full'); + expect(aborted).toBe(uploadId); + + client.dispose(); + serverTransport.close(); + }); +}); + +describe('AttachmentBlobCache', () => { + it('evicts the least recently read blob once the byte budget is spent', () => { + const cache = new AttachmentBlobCache(10); + cache.set('a', new Uint8Array(4)); + cache.set('b', new Uint8Array(4)); + expect(cache.get('a')).toBeDefined(); + cache.set('c', new Uint8Array(4)); + // 'b' was the least recently read of the two that fit alongside 'c'. + expect(cache.has('b')).toBe(false); + expect(cache.has('a')).toBe(true); + expect(cache.has('c')).toBe(true); + + cache.set('huge', new Uint8Array(40)); + expect(cache.has('huge')).toBe(true); + expect(cache.has('a')).toBe(false); + }); }); From 91b1c7c36f5376fa3ddd647c6ec5ce70e8bd1449 Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 19:06:43 +0800 Subject: [PATCH 09/11] test(workbench,schema): cover the dev mock's attachment frames and pin the chunk budget --- .../workbench/src/mock/dev-mock-host.ts | 13 ++--- .../integration/dev-mock-attachments.test.ts | 49 +++++++++++++++++++ .../tests/contract/wire/attachment.test.ts | 31 +++++++++++- 3 files changed, 83 insertions(+), 10 deletions(-) create mode 100644 packages/client/workbench/tests/integration/dev-mock-attachments.test.ts diff --git a/packages/client/workbench/src/mock/dev-mock-host.ts b/packages/client/workbench/src/mock/dev-mock-host.ts index 632b0ef3b..c4d02cc7b 100644 --- a/packages/client/workbench/src/mock/dev-mock-host.ts +++ b/packages/client/workbench/src/mock/dev-mock-host.ts @@ -46,7 +46,6 @@ import type { import { AGENT_INPUT_CAPABILITIES, ATTACHMENT_UPLOAD_CHUNK_BYTES, - ATTACHMENT_UPLOAD_WINDOW_CHUNKS, AttachmentIdSchema, blobIdFromSha256, managedAgentAssetId, @@ -1978,14 +1977,6 @@ export class DevMockHost { ); return; } - const windowEnd = - upload.received + ATTACHMENT_UPLOAD_WINDOW_CHUNKS * ATTACHMENT_UPLOAD_CHUNK_BYTES; - if (payload.offset >= windowEnd && upload.received < upload.declaredSize) { - this.sendFailure(payload.clientReqId, 'Chunk is outside the credit window', { - code: 'invalid_request', - }); - return; - } const chunk = mockBase64ToBytes(payload.data); if (upload.received + chunk.byteLength > upload.declaredSize) { this.sendFailure(payload.clientReqId, 'Chunk exceeds the declared size', { @@ -2064,6 +2055,10 @@ export class DevMockHost { return; } this.attachmentUploads.delete(payload.uploadId); + // The replay must die with the upload it names, or a retried operationId resolves to a dead id. + for (const [operationId, begun] of this.attachmentBegins) { + if (begun.uploadId === payload.uploadId) this.attachmentBegins.delete(operationId); + } this.sendSuccess(payload.clientReqId); } diff --git a/packages/client/workbench/tests/integration/dev-mock-attachments.test.ts b/packages/client/workbench/tests/integration/dev-mock-attachments.test.ts new file mode 100644 index 000000000..89d528311 --- /dev/null +++ b/packages/client/workbench/tests/integration/dev-mock-attachments.test.ts @@ -0,0 +1,49 @@ +import { LinkCodeClient } from '@linkcode/client-core'; +import { ATTACHMENT_UPLOAD_CHUNK_BYTES, AttachmentIdSchema } from '@linkcode/schema'; +import { describe, expect, it } from 'vitest'; +import { createDevMockTransport } from '../../src/mock/dev-mock-transport'; + +async function connectedClient(): Promise { + const client = new LinkCodeClient(createDevMockTransport()); + await client.connect(); + return client; +} + +describe('dev mock attachment store', () => { + it('round-trips a multi-chunk upload and dedupes the second copy', async () => { + const client = await connectedClient(); + const sessionId = await client.startSession({ kind: 'codex', cwd: '/mock/repo' }); + const bytes = new Uint8Array(ATTACHMENT_UPLOAD_CHUNK_BYTES + 17); + for (let i = 0; i < bytes.byteLength; i++) bytes[i] = i % 251; + + const first = await client.putAttachment({ bytes, name: 'shot.bin', attachmentKind: 'file' }); + const read = await client.getAttachmentBytes(sessionId, first.attachmentId); + expect(read.bytes).toEqual(bytes); + expect(read.blobId).toBe(first.blobId); + + // Same bytes, new record: the mock must answer `exists` and transfer nothing. + const second = await client.putAttachment({ bytes, name: 'copy.bin', attachmentKind: 'file' }); + expect(second.blobId).toBe(first.blobId); + expect(second.attachmentId).not.toBe(first.attachmentId); + client.dispose(); + }); + + it('rejects an unknown attachment and a chunk at the wrong offset', async () => { + const client = await connectedClient(); + const sessionId = await client.startSession({ kind: 'codex', cwd: '/mock/repo' }); + await expect( + client.readAttachment(sessionId, AttachmentIdSchema.parse('att-nope'), 0, 16), + ).rejects.toThrow('Attachment not found'); + + const begun = await client.beginAttachmentUpload({ + declaredSha256: 'b'.repeat(64), + declaredSize: 32, + name: 'offset.bin', + attachmentKind: 'file', + }); + await expect(client.sendAttachmentChunk(begun.uploadId, 8, 'YQ==')).rejects.toThrow( + 'Expected offset 0', + ); + client.dispose(); + }); +}); diff --git a/packages/foundation/schema/tests/contract/wire/attachment.test.ts b/packages/foundation/schema/tests/contract/wire/attachment.test.ts index 45a907f65..660d59998 100644 --- a/packages/foundation/schema/tests/contract/wire/attachment.test.ts +++ b/packages/foundation/schema/tests/contract/wire/attachment.test.ts @@ -136,12 +136,41 @@ describe('attachment upload/read frames', () => { ).toBe(false); }); + it('keeps one chunk frame inside the tunnel budget', () => { + // 341 KiB of base64 for 256 KiB raw: one tunnel chunk frame, far under workerd's 1 MiB cap. + expect(ATTACHMENT_UPLOAD_CHUNK_BYTES).toBe(262_144); + expect(ATTACHMENT_UPLOAD_CHUNK_BASE64_MAX).toBe(349_528); + expect(ATTACHMENT_UPLOAD_CHUNK_BASE64_MAX).toBeLessThan(768 * 1024); + }); + + it('rejects a reply that drops a required field', () => { + expect( + parses({ + kind: 'attachment.read.result', + replyTo: 'request-1', + sessionId: 'session-1', + attachmentId: 'att-1', + blobId, + offset: 0, + data: 'aGVsbG8=', + eof: true, + }), + ).toBe(false); + expect( + parses({ + kind: 'attachment.upload.begun', + replyTo: 'request-1', + uploadId: 'upl-1', + state: 'ready', + }), + ).toBe(false); + }); + it('accepts a begin without operationId and an exists short-circuit', () => { expect( parses({ kind: 'attachment.upload.begin', clientReqId: 'request-1', - operationId: 'op-1', declaredSha256: sha256, declaredSize: 0, name: 'empty.bin', From b7b79e6c82e3985c8ef13943320893f35d1ddc48 Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 20:09:43 +0800 Subject: [PATCH 10/11] refactor(schema): host the attachment mime sniff so clients can validate too --- .../src/model/__tests__/mime-sniff.test.ts | 32 +++++++++++++++++++ packages/foundation/schema/src/model/index.ts | 1 + .../schema/src/model}/mime-sniff.ts | 12 ++----- .../engine/src/__tests__/blob-store.test.ts | 27 ---------------- .../engine/src/attachment/upload-service.ts | 2 +- packages/host/engine/src/resource/service.ts | 2 +- 6 files changed, 38 insertions(+), 38 deletions(-) create mode 100644 packages/foundation/schema/src/model/__tests__/mime-sniff.test.ts rename packages/{host/engine/src/attachment => foundation/schema/src/model}/mime-sniff.ts (85%) diff --git a/packages/foundation/schema/src/model/__tests__/mime-sniff.test.ts b/packages/foundation/schema/src/model/__tests__/mime-sniff.test.ts new file mode 100644 index 000000000..f9e99f9e0 --- /dev/null +++ b/packages/foundation/schema/src/model/__tests__/mime-sniff.test.ts @@ -0,0 +1,32 @@ +import { describe, expect, it } from 'vitest'; +import { declaredMimeTypeMatches, sniffImageMimeType } from '../mime-sniff'; + +const bytes = (text: string): Uint8Array => new TextEncoder().encode(text); + +describe('mime sniff', () => { + const png = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0, 0]); + const webp = new Uint8Array(16); + webp.set(bytes('RIFF'), 0); + webp.set(bytes('WEBPVP8 '), 8); + + it('recognizes the supported image types and nothing else', () => { + expect(sniffImageMimeType(png)).toBe('image/png'); + expect(sniffImageMimeType(new Uint8Array([0xff, 0xd8, 0xff, 0xe0]))).toBe('image/jpeg'); + expect(sniffImageMimeType(bytes('GIF89a......'))).toBe('image/gif'); + expect(sniffImageMimeType(webp)).toBe('image/webp'); + expect(sniffImageMimeType(bytes('RIFF....WAVE'))).toBeUndefined(); + expect(sniffImageMimeType(bytes('%PDF-1.7'))).toBeUndefined(); + expect(sniffImageMimeType(new Uint8Array(0))).toBeUndefined(); + }); + + it('holds sniffable image declarations to their bytes and trusts the rest', () => { + expect(declaredMimeTypeMatches('image/png', png)).toBe(true); + expect(declaredMimeTypeMatches('image/jpeg', png)).toBe(false); + expect(declaredMimeTypeMatches('image/jpeg', bytes(''))).toBe(false); + expect(declaredMimeTypeMatches('image/svg+xml', bytes(''))).toBe(true); + expect(declaredMimeTypeMatches('image/svg+xml', png)).toBe(false); + expect(declaredMimeTypeMatches('image/heic', bytes('ftypheic'))).toBe(true); + expect(declaredMimeTypeMatches('application/pdf', bytes('%PDF-1.7'))).toBe(true); + expect(declaredMimeTypeMatches('text/plain', png)).toBe(true); + }); +}); diff --git a/packages/foundation/schema/src/model/index.ts b/packages/foundation/schema/src/model/index.ts index 53f15a07e..15aa9b99e 100644 --- a/packages/foundation/schema/src/model/index.ts +++ b/packages/foundation/schema/src/model/index.ts @@ -16,6 +16,7 @@ export * from './linkcode-marketplace'; export * from './linkcode-plugin'; export * from './loop'; export * from './managed-asset'; +export * from './mime-sniff'; export * from './permission'; export * from './plan'; export * from './plugin'; diff --git a/packages/host/engine/src/attachment/mime-sniff.ts b/packages/foundation/schema/src/model/mime-sniff.ts similarity index 85% rename from packages/host/engine/src/attachment/mime-sniff.ts rename to packages/foundation/schema/src/model/mime-sniff.ts index baca02d0c..269d16922 100644 --- a/packages/host/engine/src/attachment/mime-sniff.ts +++ b/packages/foundation/schema/src/model/mime-sniff.ts @@ -1,4 +1,5 @@ -import type { SupportedAttachmentImageMimeType } from '@linkcode/schema'; +import type { SupportedAttachmentImageMimeType } from './content'; +import { isSupportedAttachmentImageMimeType } from './content'; const PNG_MAGIC = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]; @@ -23,13 +24,6 @@ export function sniffImageMimeType(head: Uint8Array): SupportedAttachmentImageMi return undefined; } -const SNIFFABLE_IMAGE_TYPES = new Set([ - 'image/jpeg', - 'image/png', - 'image/gif', - 'image/webp', -]); - /** A declared sniffable `image/*` type must match its bytes — model APIs refuse the mismatch * later and less legibly. Other declarations (svg, heic, pdf, …) have no reliable sniff here * and are trusted. */ @@ -37,5 +31,5 @@ export function declaredMimeTypeMatches(declared: string, head: Uint8Array): boo if (!declared.startsWith('image/')) return true; const sniffed = sniffImageMimeType(head); if (sniffed !== undefined) return sniffed === declared; - return !SNIFFABLE_IMAGE_TYPES.has(declared); + return !isSupportedAttachmentImageMimeType(declared); } diff --git a/packages/host/engine/src/__tests__/blob-store.test.ts b/packages/host/engine/src/__tests__/blob-store.test.ts index 63e23bf1e..2ac1eab8c 100644 --- a/packages/host/engine/src/__tests__/blob-store.test.ts +++ b/packages/host/engine/src/__tests__/blob-store.test.ts @@ -5,7 +5,6 @@ import { join } from 'node:path'; import { blobIdFromSha256 } from '@linkcode/schema'; import { afterEach, describe, expect, it } from 'vitest'; import { BlobIntegrityError, FsBlobStore } from '../attachment/blob-store'; -import { declaredMimeTypeMatches, sniffImageMimeType } from '../attachment/mime-sniff'; const temporaryDirectories: string[] = []; @@ -108,29 +107,3 @@ describe('FsBlobStore', () => { expect(await store.read(blobIdFromSha256('b'.repeat(64)), 0, 4)).toBeUndefined(); }); }); - -describe('mime sniff', () => { - const png = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0, 0]); - const webp = Buffer.concat([Buffer.from('RIFF'), Buffer.alloc(4), Buffer.from('WEBPVP8 ')]); - - it('recognizes the supported image types and nothing else', () => { - expect(sniffImageMimeType(png)).toBe('image/png'); - expect(sniffImageMimeType(new Uint8Array([0xff, 0xd8, 0xff, 0xe0]))).toBe('image/jpeg'); - expect(sniffImageMimeType(Buffer.from('GIF89a......'))).toBe('image/gif'); - expect(sniffImageMimeType(webp)).toBe('image/webp'); - expect(sniffImageMimeType(Buffer.from('RIFF....WAVE'))).toBeUndefined(); - expect(sniffImageMimeType(Buffer.from('%PDF-1.7'))).toBeUndefined(); - expect(sniffImageMimeType(new Uint8Array(0))).toBeUndefined(); - }); - - it('holds sniffable image declarations to their bytes and trusts the rest', () => { - expect(declaredMimeTypeMatches('image/png', png)).toBe(true); - expect(declaredMimeTypeMatches('image/jpeg', png)).toBe(false); - expect(declaredMimeTypeMatches('image/jpeg', Buffer.from(''))).toBe(false); - expect(declaredMimeTypeMatches('image/svg+xml', Buffer.from(''))).toBe(true); - expect(declaredMimeTypeMatches('image/svg+xml', png)).toBe(false); - expect(declaredMimeTypeMatches('image/heic', Buffer.from('ftypheic'))).toBe(true); - expect(declaredMimeTypeMatches('application/pdf', Buffer.from('%PDF-1.7'))).toBe(true); - expect(declaredMimeTypeMatches('text/plain', png)).toBe(true); - }); -}); diff --git a/packages/host/engine/src/attachment/upload-service.ts b/packages/host/engine/src/attachment/upload-service.ts index aaecf2853..53ea5b1ed 100644 --- a/packages/host/engine/src/attachment/upload-service.ts +++ b/packages/host/engine/src/attachment/upload-service.ts @@ -4,6 +4,7 @@ import { ATTACHMENT_UPLOAD_CHUNK_BYTES, AttachmentIdSchema, blobIdFromSha256, + declaredMimeTypeMatches, MAX_ATTACHMENT_BYTES, UploadIdSchema, } from '@linkcode/schema'; @@ -15,7 +16,6 @@ import type { BlobStage, BlobStore } from './blob-store'; import { BlobIntegrityError } from './blob-store'; import { UPLOAD_LEASE_TTL_MS } from './gc'; import { AttachmentIoMutex } from './io-mutex'; -import { declaredMimeTypeMatches } from './mime-sniff'; const HEAD_BYTES = 16; const rUploadId = /^[\w-]{1,128}$/; diff --git a/packages/host/engine/src/resource/service.ts b/packages/host/engine/src/resource/service.ts index 4219a0322..1f8769edd 100644 --- a/packages/host/engine/src/resource/service.ts +++ b/packages/host/engine/src/resource/service.ts @@ -7,6 +7,7 @@ import type { SessionId, SessionResource, SessionResourceId } from '@linkcode/sc import { AttachmentIdSchema, blobIdFromSha256, + declaredMimeTypeMatches, MAX_ATTACHMENT_BYTES, SessionResourceIdSchema, } from '@linkcode/schema'; @@ -17,7 +18,6 @@ import { noop } from 'foxts/noop'; import type { AttachmentStore } from '../attachment/attachment-store'; import type { BlobStore } from '../attachment/blob-store'; import { AttachmentIoMutex } from '../attachment/io-mutex'; -import { declaredMimeTypeMatches } from '../attachment/mime-sniff'; import { OperationError, RequestError } from '../failure'; import type { FileHostService } from '../preview/file-host-service'; import type { SessionRecordRegistry } from '../session/session-record-registry'; From d3351d3deebff5615dda69f422ff76b937f04edb Mon Sep 17 00:00:00 2001 From: Zerlight Wu Date: Thu, 3 Sep 2026 20:11:12 +0800 Subject: [PATCH 11/11] fix(workbench): gate dev-mock attachment reads on a session root and sniff its commits --- .../workbench/src/mock/dev-mock-host.ts | 55 ++++++++- .../integration/dev-mock-attachments.test.ts | 106 ++++++++++++++++-- 2 files changed, 152 insertions(+), 9 deletions(-) diff --git a/packages/client/workbench/src/mock/dev-mock-host.ts b/packages/client/workbench/src/mock/dev-mock-host.ts index c4d02cc7b..c750cf9cd 100644 --- a/packages/client/workbench/src/mock/dev-mock-host.ts +++ b/packages/client/workbench/src/mock/dev-mock-host.ts @@ -48,6 +48,7 @@ import { ATTACHMENT_UPLOAD_CHUNK_BYTES, AttachmentIdSchema, blobIdFromSha256, + declaredMimeTypeMatches, managedAgentAssetId, managedAssetIdEquals, managedAssetKey, @@ -271,6 +272,8 @@ export class DevMockHost { { blobId: BlobId; sizeBytes: number; name: string } >(); private readonly attachmentBegins = new Map(); + /** The daemon's `isReachable` roots: sessions whose prompt or resource names the attachment. */ + private readonly attachmentSessions = new Map>(); private uploadSeq = 0; private attachmentSeq = 0; @@ -969,7 +972,16 @@ export class DevMockHost { this.resources.set(resourceId, processing); this.send({ kind: 'resource.changed', resource: processing }); await wait(CONTROL_LATENCY_MS); - const ready: SessionResource = { ...processing, status: 'ready', updatedAt: Date.now() }; + // Resource bytes land in the attachment store on the daemon, which is what roots them for + // `attachment.read`; a resource with no attachment id would be unreadable through the wire. + const attachmentId = await this.publishResourceAttachment(payload); + this.rootAttachment(payload.sessionId, attachmentId); + const ready: SessionResource = { + ...processing, + status: 'ready', + attachmentId, + updatedAt: Date.now(), + }; this.resources.set(resourceId, ready); this.send({ kind: 'resource.changed', resource: ready }); this.send({ kind: 'resource.uploaded', replyTo: payload.clientReqId, resource: ready }); @@ -1406,6 +1418,13 @@ export class DevMockHost { this.sendFailure(p.clientReqId, 'Dev mock host does not support explicit-parent submits.'); return; } + if (p.input.type === 'prompt') { + const blocks = p.input.blocks; + for (let i = 0, len = blocks.length; i < len; i++) { + const block = blocks[i]; + if (block.type === 'attachment_ref') this.rootAttachment(p.sessionId, block.attachmentId); + } + } const content = turnSubmitContent(p.input); const turn = this.beginTurn(session, content, p.input.type === 'prompt' ? undefined : p.input); this.send({ kind: 'turn.submitted', replyTo: p.clientReqId, turnId: turn.graph.turnId }); @@ -2019,6 +2038,15 @@ export class DevMockHost { ); return; } + // Engine order: size, then declared MIME vs bytes, then SHA-256 — a rejected commit must + // leave no blob behind. + const declaredMime = upload.mimeType ?? 'application/octet-stream'; + if (!declaredMimeTypeMatches(declaredMime, upload.bytes.subarray(0, 16))) { + this.sendFailure(payload.clientReqId, `File contents are not ${declaredMime}`, { + code: 'invalid_request', + }); + return; + } if (upload.state === 'ready') { const digest = await mockSha256Hex(upload.bytes); if (digest !== upload.declaredSha256) { @@ -2062,8 +2090,31 @@ export class DevMockHost { this.sendSuccess(payload.clientReqId); } + private async publishResourceAttachment( + payload: Extract, + ): Promise { + const bytes = mockBase64ToBytes(payload.data); + const digest = await mockSha256Hex(bytes); + this.attachmentBlobs.set(digest, bytes); + this.attachmentSeq += 1; + const attachmentId = AttachmentIdSchema.parse(`att-mock-${this.attachmentSeq}`); + this.attachmentRecords.set(attachmentId, { + blobId: blobIdFromSha256(digest), + sizeBytes: bytes.byteLength, + name: payload.name, + }); + return attachmentId; + } + + /** Root an attachment in a session, the way persisting a prompt or a resource does on the daemon. */ + private rootAttachment(sessionId: SessionId, attachmentId: AttachmentId): void { + const rooted = this.attachmentSessions.get(attachmentId) ?? new Set(); + rooted.add(sessionId); + this.attachmentSessions.set(attachmentId, rooted); + } + private readAttachment(payload: Extract): void { - if (!this.sessions.has(payload.sessionId)) { + if (!this.attachmentSessions.get(payload.attachmentId)?.has(payload.sessionId)) { this.sendFailure(payload.clientReqId, 'Attachment not found', { code: 'not_found' }); return; } diff --git a/packages/client/workbench/tests/integration/dev-mock-attachments.test.ts b/packages/client/workbench/tests/integration/dev-mock-attachments.test.ts index 89d528311..93d2c86af 100644 --- a/packages/client/workbench/tests/integration/dev-mock-attachments.test.ts +++ b/packages/client/workbench/tests/integration/dev-mock-attachments.test.ts @@ -1,5 +1,13 @@ import { LinkCodeClient } from '@linkcode/client-core'; -import { ATTACHMENT_UPLOAD_CHUNK_BYTES, AttachmentIdSchema } from '@linkcode/schema'; +import type { AttachmentId, SessionId } from '@linkcode/schema'; +import { + ATTACHMENT_UPLOAD_CHUNK_BYTES, + AttachmentIdSchema, + OperationIdSchema, +} from '@linkcode/schema'; +import type { Transport } from '@linkcode/transport'; +import { createWireMessage } from '@linkcode/transport'; +import { nullthrow } from 'foxts/guard'; import { describe, expect, it } from 'vitest'; import { createDevMockTransport } from '../../src/mock/dev-mock-transport'; @@ -9,18 +17,47 @@ async function connectedClient(): Promise { return client; } +/** `turn.submit` has no client-core method yet (CODE-638), so the prompt-ref root is driven raw. */ +function submitPromptRef( + transport: Transport, + sessionId: SessionId, + attachmentId: AttachmentId, +): Promise { + return new Promise((resolve, reject) => { + const clientReqId = 'creq-attachment-ref'; + const unsubscribe = transport.onMessage((message) => { + const p = message.payload; + if (!('replyTo' in p) || p.replyTo !== clientReqId) return; + unsubscribe(); + if (p.kind === 'turn.submitted') resolve(); + else reject(new Error(p.kind === 'request.failed' ? p.message : `unexpected ${p.kind}`)); + }); + transport.send( + createWireMessage({ + kind: 'turn.submit', + clientReqId, + sessionId, + operationId: OperationIdSchema.parse('op-attachment-ref'), + input: { + type: 'prompt', + blocks: [ + { type: 'text', text: 'look at this' }, + { type: 'attachment_ref', attachmentId }, + ], + }, + }), + ); + }); +} + describe('dev mock attachment store', () => { it('round-trips a multi-chunk upload and dedupes the second copy', async () => { const client = await connectedClient(); - const sessionId = await client.startSession({ kind: 'codex', cwd: '/mock/repo' }); + await client.startSession({ kind: 'codex', cwd: '/mock/repo' }); const bytes = new Uint8Array(ATTACHMENT_UPLOAD_CHUNK_BYTES + 17); for (let i = 0; i < bytes.byteLength; i++) bytes[i] = i % 251; const first = await client.putAttachment({ bytes, name: 'shot.bin', attachmentKind: 'file' }); - const read = await client.getAttachmentBytes(sessionId, first.attachmentId); - expect(read.bytes).toEqual(bytes); - expect(read.blobId).toBe(first.blobId); - // Same bytes, new record: the mock must answer `exists` and transfer nothing. const second = await client.putAttachment({ bytes, name: 'copy.bin', attachmentKind: 'file' }); expect(second.blobId).toBe(first.blobId); @@ -28,7 +65,53 @@ describe('dev mock attachment store', () => { client.dispose(); }); - it('rejects an unknown attachment and a chunk at the wrong offset', async () => { + it('reads an attachment only from a session that roots it', async () => { + const client = await connectedClient(); + const sessionId = await client.startSession({ kind: 'codex', cwd: '/mock/repo' }); + const bytes = new TextEncoder().encode('resource bytes'); + + // An upload alone is a draft lease — the daemon's isReachable roots nothing yet. + const draft = await client.putAttachment({ bytes, name: 'draft.txt', attachmentKind: 'file' }); + await expect(client.getAttachmentBytes(sessionId, draft.attachmentId)).rejects.toThrow( + 'Attachment not found', + ); + + // A session resource is a root, and carries the attachment its bytes landed in. + const resource = await client.uploadSource( + sessionId, + 'brief.txt', + btoa('resource bytes'), + 'text/plain', + ); + const attachmentId = nullthrow(resource.attachmentId, 'resource missing attachmentId'); + const read = await client.getAttachmentBytes(sessionId, attachmentId); + expect(read.bytes).toEqual(bytes); + + const otherSession = await client.startSession({ kind: 'codex', cwd: '/mock/other' }); + await expect(client.getAttachmentBytes(otherSession, attachmentId)).rejects.toThrow( + 'Attachment not found', + ); + client.dispose(); + }); + + it('roots a draft attachment once a prompt of that session references it', async () => { + const transport = createDevMockTransport(); + const client = new LinkCodeClient(transport); + await client.connect(); + const sessionId = await client.startSession({ kind: 'codex', cwd: '/mock/repo' }); + const bytes = new TextEncoder().encode('attached by prompt'); + const draft = await client.putAttachment({ bytes, name: 'note.txt', attachmentKind: 'file' }); + + await expect(client.getAttachmentBytes(sessionId, draft.attachmentId)).rejects.toThrow( + 'Attachment not found', + ); + await submitPromptRef(transport, sessionId, draft.attachmentId); + const read = await client.getAttachmentBytes(sessionId, draft.attachmentId); + expect(read.bytes).toEqual(bytes); + client.dispose(); + }); + + it('rejects an unknown attachment, a wrong offset, and bytes that are not the declared image', async () => { const client = await connectedClient(); const sessionId = await client.startSession({ kind: 'codex', cwd: '/mock/repo' }); await expect( @@ -44,6 +127,15 @@ describe('dev mock attachment store', () => { await expect(client.sendAttachmentChunk(begun.uploadId, 8, 'YQ==')).rejects.toThrow( 'Expected offset 0', ); + + await expect( + client.putAttachment({ + bytes: new TextEncoder().encode('not a png'), + name: 'fake.png', + mimeType: 'image/png', + attachmentKind: 'image', + }), + ).rejects.toThrow('File contents are not image/png'); client.dispose(); }); });