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/client/core/src/client.ts b/packages/client/core/src/client.ts index b6fc026bb..a9c371ff3 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,13 @@ 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 { Sha256Hex } from './client/blob-cache'; import type { BrowserCommandExecutor } from './client/browser-host-channel'; import { BrowserHostChannel } from './client/browser-host-channel'; import type { @@ -111,6 +121,12 @@ 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 { Sha256Hex } from './client/blob-cache'; export type { BrowserCommandExecutor } from './client/browser-host-channel'; export type { ConversationReadClientOptions, @@ -146,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 { @@ -240,6 +258,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 +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, 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); @@ -326,6 +346,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 +632,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 +1378,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..f0d2c2f2a --- /dev/null +++ b/packages/client/core/src/client/attachment-channel.ts @@ -0,0 +1,193 @@ +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 type { Sha256Hex } from './blob-cache'; +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, + private readonly digest: Sha256Hex = sha256Hex, + ) {} + + 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 this.digest(input.bytes); + const begun = await this.beginUpload({ + declaredSha256, + declaredSize: input.bytes.byteLength, + name: input.name, + mimeType: input.mimeType, + attachmentKind: input.attachmentKind, + operationId: input.operationId, + }); + 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; + } + 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); + let page = first; + let offset = 0; + while (offset < first.sizeBytes) { + const slice = base64ToBytes(page.data); + // 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 }; + } + + 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; + 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 + 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..d27e28de9 --- /dev/null +++ b/packages/client/core/src/client/blob-cache.ts @@ -0,0 +1,94 @@ +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 { + 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 { + 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; +} + +/** 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 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++) { + 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..dee65cb4b --- /dev/null +++ b/packages/client/core/tests/integration/attachment-client.test.ts @@ -0,0 +1,324 @@ +import { + ATTACHMENT_STORE_WIRE_VERSION, + ATTACHMENT_UPLOAD_CHUNK_BYTES, + 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 { + AttachmentBlobCache, + 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(); + }); + + 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); + }); +}); diff --git a/packages/client/workbench/src/mock/dev-mock-host.ts b/packages/client/workbench/src/mock/dev-mock-host.ts index 83a4ef5b6..c750cf9cd 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, + AttachmentIdSchema, + blobIdFromSha256, + declaredMimeTypeMatches, 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,17 @@ 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(); + /** The daemon's `isReachable` roots: sessions whose prompt or resource names the attachment. */ + private readonly attachmentSessions = new Map>(); + private uploadSeq = 0; + private attachmentSeq = 0; constructor(private readonly transport: Transport) { this.terminals.set( @@ -453,6 +492,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({ @@ -914,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 }); @@ -1351,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 }); @@ -1857,6 +1931,224 @@ 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 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; + } + // 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) { + 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); + // 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); + } + + 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.attachmentSessions.get(payload.attachmentId)?.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 +2160,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 +2265,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); 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..93d2c86af --- /dev/null +++ b/packages/client/workbench/tests/integration/dev-mock-attachments.test.ts @@ -0,0 +1,141 @@ +import { LinkCodeClient } from '@linkcode/client-core'; +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'; + +async function connectedClient(): Promise { + const client = new LinkCodeClient(createDevMockTransport()); + await client.connect(); + 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(); + 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' }); + // 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('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( + 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', + ); + + 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(); + }); +}); 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/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/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/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..660d59998 --- /dev/null +++ b/packages/foundation/schema/tests/contract/wire/attachment.test.ts @@ -0,0 +1,190 @@ +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('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', + 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); + }); +}); 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..3a892e748 --- /dev/null +++ b/packages/host/engine/src/__tests__/attachment-upload.test.ts @@ -0,0 +1,388 @@ +import { createHash } from 'node:crypto'; +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 { + 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[] = []; +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[] = [], clock?: () => number) { + 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, 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 { + 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('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'); + 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__/blob-store.test.ts b/packages/host/engine/src/__tests__/blob-store.test.ts index d611649b5..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[] = []; @@ -94,30 +93,17 @@ describe('FsBlobStore', () => { expect(await store.list()).toEqual([]); await expect(store.delete(blobId)).resolves.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('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 }); - 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); + 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(); }); }); 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..272fa6359 --- /dev/null +++ b/packages/host/engine/src/__tests__/engine-attachments.test.ts @@ -0,0 +1,240 @@ +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' }), + ); + }); + }); + + 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/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/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..53ea5b1ed --- /dev/null +++ b/packages/host/engine/src/attachment/upload-service.ts @@ -0,0 +1,461 @@ +import { randomUUID } from 'node:crypto'; +import type { AttachmentId, BlobId, SessionId, UploadId, UploadLease } from '@linkcode/schema'; +import { + ATTACHMENT_UPLOAD_CHUNK_BYTES, + AttachmentIdSchema, + blobIdFromSha256, + declaredMimeTypeMatches, + 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'; + +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; + /** 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'; +} + +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); + const reapExpired = this.reapExpired.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'); + } + // 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', () => + 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, + gate: new AttachmentIoMutex(), + 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'); + } + 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 }; + }), + ); + }); + } + + commit(uploadId: UploadId): Effect.Effect { + 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* () { + 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)); + // 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); + yield* files('discard', () => discard(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; + }).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 { + 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'); + // `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, + 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); + // 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; + } + }); + if (!blobId) throw new Error('Attachment commit produced no blob id'); + return { attachmentId, blobId }; + }); + } + + 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); + } + } + + 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/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..4599f9079 100644 --- a/packages/host/engine/src/engine.ts +++ b/packages/host/engine/src/engine.ts @@ -16,10 +16,13 @@ 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'; import { AttachmentIoMutex } from './attachment/io-mutex'; +import { AttachmentRequestHandler } from './attachment/request-handler'; +import { AttachmentUploadService } from './attachment/upload-service'; import { InMemoryLoopStore, InMemoryScheduleStore, @@ -121,14 +124,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( @@ -141,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( @@ -245,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, @@ -334,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, @@ -477,6 +486,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; diff --git a/packages/host/engine/src/resource/service.ts b/packages/host/engine/src/resource/service.ts index 0f6726c6b..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'; @@ -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; } }); 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':