fix(workbench): gate dev-mock attachment reads on a session root and sniff its commits

Author: Zerlight

packages/client/workbench/src/mock/dev-mock-host.ts

@@ -48,6 +48,7 @@ import {
   ATTACHMENT_UPLOAD_CHUNK_BYTES,
   AttachmentIdSchema,
   blobIdFromSha256,
+  declaredMimeTypeMatches,
   managedAgentAssetId,
   managedAssetIdEquals,
   managedAssetKey,
@@ -271,6 +272,8 @@ export class DevMockHost {
     { blobId: BlobId; sizeBytes: number; name: string }
   >();
   private readonly attachmentBegins = new Map<string, MockAttachmentBegin>();
+  /** The daemon's `isReachable` roots: sessions whose prompt or resource names the attachment. */
+  private readonly attachmentSessions = new Map<AttachmentId, Set<SessionId>>();
   private uploadSeq = 0;
   private attachmentSeq = 0;
 
@@ -969,7 +972,16 @@ export class DevMockHost {
     this.resources.set(resourceId, processing);
     this.send({ kind: 'resource.changed', resource: processing });
     await wait(CONTROL_LATENCY_MS);
-    const ready: SessionResource = { ...processing, status: 'ready', updatedAt: Date.now() };
+    // Resource bytes land in the attachment store on the daemon, which is what roots them for
+    // `attachment.read`; a resource with no attachment id would be unreadable through the wire.
+    const attachmentId = await this.publishResourceAttachment(payload);
+    this.rootAttachment(payload.sessionId, attachmentId);
+    const ready: SessionResource = {
+      ...processing,
+      status: 'ready',
+      attachmentId,
+      updatedAt: Date.now(),
+    };
     this.resources.set(resourceId, ready);
     this.send({ kind: 'resource.changed', resource: ready });
     this.send({ kind: 'resource.uploaded', replyTo: payload.clientReqId, resource: ready });
@@ -1406,6 +1418,13 @@ export class DevMockHost {
       this.sendFailure(p.clientReqId, 'Dev mock host does not support explicit-parent submits.');
       return;
     }
+    if (p.input.type === 'prompt') {
+      const blocks = p.input.blocks;
+      for (let i = 0, len = blocks.length; i < len; i++) {
+        const block = blocks[i];
+        if (block.type === 'attachment_ref') this.rootAttachment(p.sessionId, block.attachmentId);
+      }
+    }
     const content = turnSubmitContent(p.input);
     const turn = this.beginTurn(session, content, p.input.type === 'prompt' ? undefined : p.input);
     this.send({ kind: 'turn.submitted', replyTo: p.clientReqId, turnId: turn.graph.turnId });
@@ -2019,6 +2038,15 @@ export class DevMockHost {
       );
       return;
     }
+    // Engine order: size, then declared MIME vs bytes, then SHA-256 — a rejected commit must
+    // leave no blob behind.
+    const declaredMime = upload.mimeType ?? 'application/octet-stream';
+    if (!declaredMimeTypeMatches(declaredMime, upload.bytes.subarray(0, 16))) {
+      this.sendFailure(payload.clientReqId, `File contents are not ${declaredMime}`, {
+        code: 'invalid_request',
+      });
+      return;
+    }
     if (upload.state === 'ready') {
       const digest = await mockSha256Hex(upload.bytes);
       if (digest !== upload.declaredSha256) {
@@ -2062,8 +2090,31 @@ export class DevMockHost {
     this.sendSuccess(payload.clientReqId);
   }
 
+  private async publishResourceAttachment(
+    payload: Extract<WirePayload, { kind: 'resource.source.upload' }>,
+  ): Promise<AttachmentId> {
+    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<SessionId>();
+    rooted.add(sessionId);
+    this.attachmentSessions.set(attachmentId, rooted);
+  }
+
   private readAttachment(payload: Extract<WirePayload, { kind: 'attachment.read' }>): void {
-    if (!this.sessions.has(payload.sessionId)) {
+    if (!this.attachmentSessions.get(payload.attachmentId)?.has(payload.sessionId)) {
       this.sendFailure(payload.clientReqId, 'Attachment not found', { code: 'not_found' });
       return;
     }

packages/client/workbench/tests/integration/dev-mock-attachments.test.ts

@@ -1,5 +1,13 @@
 import { LinkCodeClient } from '@linkcode/client-core';
-import { ATTACHMENT_UPLOAD_CHUNK_BYTES, AttachmentIdSchema } from '@linkcode/schema';
+import type { AttachmentId, SessionId } from '@linkcode/schema';
+import {
+  ATTACHMENT_UPLOAD_CHUNK_BYTES,
+  AttachmentIdSchema,
+  OperationIdSchema,
+} from '@linkcode/schema';
+import type { Transport } from '@linkcode/transport';
+import { createWireMessage } from '@linkcode/transport';
+import { nullthrow } from 'foxts/guard';
 import { describe, expect, it } from 'vitest';
 import { createDevMockTransport } from '../../src/mock/dev-mock-transport';
 
@@ -9,26 +17,101 @@ async function connectedClient(): Promise<LinkCodeClient> {
   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<void> {
+  return new Promise((resolve, reject) => {
+    const clientReqId = 'creq-attachment-ref';
+    const unsubscribe = transport.onMessage((message) => {
+      const p = message.payload;
+      if (!('replyTo' in p) || p.replyTo !== clientReqId) return;
+      unsubscribe();
+      if (p.kind === 'turn.submitted') resolve();
+      else reject(new Error(p.kind === 'request.failed' ? p.message : `unexpected ${p.kind}`));
+    });
+    transport.send(
+      createWireMessage({
+        kind: 'turn.submit',
+        clientReqId,
+        sessionId,
+        operationId: OperationIdSchema.parse('op-attachment-ref'),
+        input: {
+          type: 'prompt',
+          blocks: [
+            { type: 'text', text: 'look at this' },
+            { type: 'attachment_ref', attachmentId },
+          ],
+        },
+      }),
+    );
+  });
+}
+
 describe('dev mock attachment store', () => {
   it('round-trips a multi-chunk upload and dedupes the second copy', async () => {
     const client = await connectedClient();
-    const sessionId = await client.startSession({ kind: 'codex', cwd: '/mock/repo' });
+    await client.startSession({ kind: 'codex', cwd: '/mock/repo' });
     const bytes = new Uint8Array(ATTACHMENT_UPLOAD_CHUNK_BYTES + 17);
     for (let i = 0; i < bytes.byteLength; i++) bytes[i] = i % 251;
 
     const first = await client.putAttachment({ bytes, name: 'shot.bin', attachmentKind: 'file' });
-    const read = await client.getAttachmentBytes(sessionId, first.attachmentId);
-    expect(read.bytes).toEqual(bytes);
-    expect(read.blobId).toBe(first.blobId);
-
     // Same bytes, new record: the mock must answer `exists` and transfer nothing.
     const second = await client.putAttachment({ bytes, name: 'copy.bin', attachmentKind: 'file' });
     expect(second.blobId).toBe(first.blobId);
     expect(second.attachmentId).not.toBe(first.attachmentId);
     client.dispose();
   });
 
-  it('rejects an unknown attachment and a chunk at the wrong offset', async () => {
+  it('reads an attachment only from a session that roots it', async () => {
+    const client = await connectedClient();
+    const sessionId = await client.startSession({ kind: 'codex', cwd: '/mock/repo' });
+    const bytes = new TextEncoder().encode('resource bytes');
+
+    // An upload alone is a draft lease — the daemon's isReachable roots nothing yet.
+    const draft = await client.putAttachment({ bytes, name: 'draft.txt', attachmentKind: 'file' });
+    await expect(client.getAttachmentBytes(sessionId, draft.attachmentId)).rejects.toThrow(
+      'Attachment not found',
+    );
+
+    // A session resource is a root, and carries the attachment its bytes landed in.
+    const resource = await client.uploadSource(
+      sessionId,
+      'brief.txt',
+      btoa('resource bytes'),
+      'text/plain',
+    );
+    const attachmentId = nullthrow(resource.attachmentId, 'resource missing attachmentId');
+    const read = await client.getAttachmentBytes(sessionId, attachmentId);
+    expect(read.bytes).toEqual(bytes);
+
+    const otherSession = await client.startSession({ kind: 'codex', cwd: '/mock/other' });
+    await expect(client.getAttachmentBytes(otherSession, attachmentId)).rejects.toThrow(
+      'Attachment not found',
+    );
+    client.dispose();
+  });
+
+  it('roots a draft attachment once a prompt of that session references it', async () => {
+    const transport = createDevMockTransport();
+    const client = new LinkCodeClient(transport);
+    await client.connect();
+    const sessionId = await client.startSession({ kind: 'codex', cwd: '/mock/repo' });
+    const bytes = new TextEncoder().encode('attached by prompt');
+    const draft = await client.putAttachment({ bytes, name: 'note.txt', attachmentKind: 'file' });
+
+    await expect(client.getAttachmentBytes(sessionId, draft.attachmentId)).rejects.toThrow(
+      'Attachment not found',
+    );
+    await submitPromptRef(transport, sessionId, draft.attachmentId);
+    const read = await client.getAttachmentBytes(sessionId, draft.attachmentId);
+    expect(read.bytes).toEqual(bytes);
+    client.dispose();
+  });
+
+  it('rejects an unknown attachment, a wrong offset, and bytes that are not the declared image', async () => {
     const client = await connectedClient();
     const sessionId = await client.startSession({ kind: 'codex', cwd: '/mock/repo' });
     await expect(
@@ -44,6 +127,15 @@ describe('dev mock attachment store', () => {
     await expect(client.sendAttachmentChunk(begun.uploadId, 8, 'YQ==')).rejects.toThrow(
       'Expected offset 0',
     );
+
+    await expect(
+      client.putAttachment({
+        bytes: new TextEncoder().encode('not a png'),
+        name: 'fake.png',
+        mimeType: 'image/png',
+        attachmentKind: 'image',
+      }),
+    ).rejects.toThrow('File contents are not image/png');
     client.dispose();
   });
 });