fix(workbench): bound Agent Document responses and keep document markdown images controlled

.changeset/agent-document-review-fixes.md

@@ -0,0 +1,5 @@
+---
+"agent-bundle": patch
+---
+
+Bound decoded Agent Document responses and prevent remote Markdown images from loading in the Workbench document stage.

packages/agent-bundle/src/dev/runtime-routes.ts

@@ -19,6 +19,8 @@ import { freezeJsonValue, type JsonValue } from './types.ts';
 
 const bodyLimit = 64 * 1024;
 const assetLimit = 4 * 1024 * 1024;
+// 16 MiB leaves room for rich timelines while bounding retained replacement snapshots.
+const agentDocumentResponseLimit = 16 * 1024 * 1024;
 
 interface RequestDiagnostic {
   readonly code: string;
@@ -418,18 +420,35 @@ export class RuntimeRoutes {
         ));
       }
       try {
+        const abortController = new AbortController();
         const flight = new ReadableStream<Uint8Array>({
           start(controller) {
             controller.enqueue(asset.body);
             controller.close();
           },
         });
         const events: AgentRenderEvent[] = [];
-        for await (const event of runtime.decodeAgentFlightStream(flight, { limits: runtime.DEFAULT_AGENT_RENDER_LIMITS })) {
+        let responseBytes = Buffer.byteLength('{"events":[]}');
+        for await (const event of runtime.decodeAgentFlightStream(flight, {
+          limits: runtime.DEFAULT_AGENT_RENDER_LIMITS,
+          signal: abortController.signal,
+        })) {
+          const eventBytes = Buffer.byteLength(JSON.stringify(event));
+          const separatorBytes = events.length === 0 ? 0 : 1;
+          if (responseBytes + separatorBytes + eventBytes > agentDocumentResponseLimit) {
+            abortController.abort();
+            throw requestError(diagnostic(
+              'AB8209',
+              'Decoded Agent Document exceeds the 16 MiB response limit.',
+              413,
+            ));
+          }
+          responseBytes += separatorBytes + eventBytes;
           events.push(event);
         }
         return responseJson(response, { events });
-      } catch {
+      } catch (error) {
+        if (isRequestDiagnostic(error)) throw error;
         throw requestError(diagnostic('AB8208', 'Stored Flight could not be decoded as an Agent Document.', 409));
       }
     }

packages/agent-bundle/tests/runtime-routes.test.ts

@@ -322,6 +322,56 @@ it('decodes stored Flight into bounded Agent Document events in the foreground p
   }
 });
 
+it('aborts decoding when Agent Document events exceed the aggregate response budget', async () => {
+  let aborted = false;
+  const largeText = 'x'.repeat(512 * 1024);
+  const server = await start(new MemoryRuntime(), {
+    loadAgentDocumentRuntime: async () => ({
+      DEFAULT_AGENT_RENDER_LIMITS,
+      decodeAgentFlightStream: (_flight, options) => {
+        let sequence = 0;
+        return new ReadableStream({
+          start(controller) {
+            options?.signal?.addEventListener('abort', () => {
+              aborted = true;
+              controller.error(new DOMException('Agent render was aborted', 'AbortError'));
+            }, { once: true });
+          },
+          pull(controller) {
+            if (sequence === 40) {
+              controller.close();
+              return;
+            }
+            controller.enqueue({
+              document: {
+                root: { children: [{ kind: 'markdown', text: largeText }], kind: 'result' },
+                status: 'success',
+                version: 1,
+              },
+              sequence,
+              type: 'shell',
+            });
+            sequence += 1;
+          },
+        });
+      },
+    }),
+  });
+  try {
+    const response = await fetch(`${server.url}/api/runtime/runs/run-a/document`, { headers: authenticated(server) });
+    expect(response.status).toBe(413);
+    await expect(response.json()).resolves.toEqual({
+      diagnostic: {
+        code: 'AB8209',
+        message: 'Decoded Agent Document exceeds the 16 MiB response limit.',
+      },
+    });
+    expect(aborted).toBe(true);
+  } finally {
+    await server.close();
+  }
+});
+
 it('returns honest diagnostics when the Agent runtime is absent or stored Flight cannot decode', async () => {
   const absent = await start(new MemoryRuntime(), {
     loadAgentDocumentRuntime: async () => {

packages/workbench/src/runtime/agent-document-stage.tsx

@@ -1,6 +1,7 @@
 import React, { useEffect, useState } from 'react';
 
 import { MarkdownProjector } from '../skill-markdown.tsx';
+import { allowedExternalResourceUrl } from '../skills-model.ts';
 import type {
   AgentDocument,
   AgentDocumentNode,
@@ -74,6 +75,12 @@ const progressLabel = (progress: Readonly<{
   return progress.message === undefined ? amount : `${progress.message} · ${amount}`;
 };
 
+const agentDocumentImageUrl = (reference: string): string | undefined =>
+  /^data:/iu.test(reference) ? reference : undefined;
+
+const agentDocumentLinkUrl = (reference: string): string | undefined =>
+  reference.startsWith('#') ? reference : allowedExternalResourceUrl(reference);
+
 const AgentDocumentNodeView = ({ node, path }: Readonly<{
   readonly node: AgentDocumentNode;
   readonly path: string;
@@ -88,7 +95,11 @@ const AgentDocumentNodeView = ({ node, path }: Readonly<{
         </div>
       </section>;
     case 'markdown':
-      return <MarkdownProjector body={node.text} />;
+      return <MarkdownProjector
+        body={node.text}
+        resolveImage={agentDocumentImageUrl}
+        resolveLink={agentDocumentLinkUrl}
+      />;
     case 'text':
       return <p className="agent-document-text">{node.text}</p>;
     case 'context':

packages/workbench/src/skill-markdown.tsx

@@ -1,4 +1,4 @@
-import React, { lazy, Suspense, type ComponentProps, type ComponentPropsWithoutRef, type ReactNode } from 'react';
+import React, { lazy, Suspense, type ComponentPropsWithoutRef, type ReactNode } from 'react';
 import ReactMarkdown from 'react-markdown';
 import remarkGfm from 'remark-gfm';
 
@@ -20,7 +20,8 @@ export interface SkillMarkdownProps {
 
 export interface MarkdownProjectorProps {
   readonly body: string;
-  readonly components?: ComponentProps<typeof ReactMarkdown>['components'];
+  readonly resolveImage: (reference: string) => string | undefined;
+  readonly resolveLink: (reference: string) => string | undefined;
 }
 
 type MarkdownElementProps<Tag extends keyof React.JSX.IntrinsicElements> =
@@ -68,47 +69,61 @@ const SkillCode = ({ children, className }: ComponentPropsWithoutRef<'code'>) =>
   </Suspense>;
 };
 
-const SkillLink = ({
-  base,
+const MarkdownLink = ({
   children,
   href,
   node: _node,
-  resources,
+  resolve,
   ...properties
-}: MarkdownElementProps<'a'> & Pick<SkillMarkdownProps, 'base' | 'resources'>) => {
-  const resolved = typeof href === 'string' ? resourceUrlFor(base, href, resources) : undefined;
+}: MarkdownElementProps<'a'> & Readonly<{
+  readonly resolve: MarkdownProjectorProps['resolveLink'];
+}>) => {
+  const resolved = typeof href === 'string' ? resolve(href) : undefined;
   if (resolved === undefined) return <span className="skill-broken-link">{children}</span>;
   const external = /^https?:|^mailto:/u.test(resolved);
   return <a {...properties} href={resolved} {...(external ? { rel: 'noreferrer', target: '_blank' } : {})}>{children}</a>;
 };
 
-const SkillImage = ({
+const MarkdownImage = ({
   alt,
-  base,
   node: _node,
-  resources,
+  resolve,
   src,
   ...properties
-}: MarkdownElementProps<'img'> & Pick<SkillMarkdownProps, 'base' | 'resources'>) => {
-  const resolved = typeof src === 'string' ? resourceUrlFor(base, src, resources) : undefined;
-  if (resolved === undefined) return <span className="skill-broken-image" role="img">{alt ?? 'Image unavailable'}</span>;
+}: MarkdownElementProps<'img'> & Readonly<{
+  readonly resolve: MarkdownProjectorProps['resolveImage'];
+}>) => {
+  const source = typeof src === 'string' ? src : undefined;
+  const resolved = source === undefined ? undefined : resolve(source);
+  if (resolved === undefined) {
+    return <span className="skill-broken-image" role="img">
+      {alt ?? 'Image unavailable'}
+      {source === undefined ? undefined : <> · <code>{source}</code></>}
+    </span>;
+  }
   return <img {...properties} alt={alt ?? ''} src={resolved} />;
 };
 
 /** The audited inert-HTML/GFM projector shared by Skills and Agent Documents. */
-export const MarkdownProjector = ({ body, components }: MarkdownProjectorProps) => (
+export const MarkdownProjector = ({
+  body,
+  resolveImage,
+  resolveLink,
+}: MarkdownProjectorProps) => (
   <div className="skill-markdown">
     <ReactMarkdown
       components={{
+        a: (properties) => <MarkdownLink {...properties} resolve={resolveLink} />,
         code: SkillCode,
         h1: ({ node: _node, ...properties }) => <h1 className="skill-heading skill-heading--one" {...properties} />,
         h2: ({ node: _node, ...properties }) => <h2 className="skill-heading skill-heading--two" {...properties} />,
         h3: ({ node: _node, ...properties }) => <h3 className="skill-heading skill-heading--three" {...properties} />,
+        img: (properties) => <MarkdownImage {...properties} resolve={resolveImage} />,
         pre: ({ children }) => <>{children}</>,
         table: ({ node: _node, ...properties }) => <div className="skill-table-wrap"><table {...properties} /></div>,
-        ...components,
       }}
       remarkPlugins={[remarkGfm, inertHtml]}
+      urlTransform={(url) => url}
     >
       {body}
     </ReactMarkdown>
@@ -119,9 +134,7 @@ export const MarkdownProjector = ({ body, components }: MarkdownProjectorProps)
 export const SkillMarkdown = ({ base, body, resources }: SkillMarkdownProps) => (
   <MarkdownProjector
     body={body}
-    components={{
-      a: (properties) => <SkillLink {...properties} base={base} resources={resources} />,
-      img: (properties) => <SkillImage {...properties} base={base} resources={resources} />,
-    }}
+    resolveImage={(reference) => resourceUrlFor(base, reference, resources)}
+    resolveLink={(reference) => resourceUrlFor(base, reference, resources)}
   />
 );

packages/workbench/src/skills-model.ts

@@ -23,12 +23,14 @@ const splitReference = (reference: string): Readonly<{ readonly fragment: string
     : Object.freeze({ fragment: reference.slice(index), path: reference.slice(0, index) });
 };
 
-const isAllowedExternal = (value: string): boolean => {
+export const allowedExternalResourceUrl = (value: string): string | undefined => {
   try {
     const protocol = new URL(value).protocol;
-    return protocol === 'http:' || protocol === 'https:' || protocol === 'mailto:';
+    return protocol === 'http:' || protocol === 'https:' || protocol === 'mailto:'
+      ? value
+      : undefined;
   } catch {
-    return false;
+    return undefined;
   }
 };
 
@@ -45,7 +47,8 @@ export const resourceUrlFor = (
   resources: readonly string[],
 ): string | undefined => {
   if (reference.startsWith('#')) return reference;
-  if (isAllowedExternal(reference)) return reference;
+  const external = allowedExternalResourceUrl(reference);
+  if (external !== undefined) return external;
   const { fragment, path } = splitReference(reference);
   const segments = localSegments(path);
   if (segments === undefined) return undefined;

packages/workbench/tests/agent-document-stage.test.ts

@@ -110,4 +110,40 @@ describe('Agent Document stage', () => {
     expect(markup).toContain('Shell · #0');
     expect(markup).toContain('Complete · #3');
   });
+
+  it('keeps remote Markdown images inert while rendering data URI images', () => {
+    const projected: AgentDocument = {
+      root: {
+        children: [{
+          kind: 'markdown',
+          text: [
+            '![Remote tracker](https://example.invalid/track)',
+            '![Protocol-relative tracker](//example.invalid/track)',
+            '![Inline image](data:image/png;base64,iVBORw0KGgo=)',
+            '[External guide](https://example.com/guide)',
+          ].join('\n\n'),
+        }],
+        kind: 'result',
+      },
+      status: 'success',
+      version: 1,
+    };
+
+    const markup = renderToStaticMarkup(createElement(AgentDocumentStage, {
+      events: [{ document: projected, sequence: 0, type: 'complete' }],
+    }));
+
+    expect(markup).toContain('class="skill-broken-image"');
+    expect(markup).toContain('Remote tracker');
+    expect(markup).toContain('https://example.invalid/track');
+    expect(markup).not.toContain('src="https://example.invalid/track"');
+    expect(markup).toContain('Protocol-relative tracker');
+    expect(markup).toContain('//example.invalid/track');
+    expect(markup).not.toContain('src="//example.invalid/track"');
+    expect(markup.match(/<img/gu)).toHaveLength(1);
+    expect(markup).toContain('src="data:image/png;base64,iVBORw0KGgo="');
+    expect(markup).toContain('href="https://example.com/guide"');
+    expect(markup).toContain('rel="noreferrer"');
+    expect(markup).toContain('target="_blank"');
+  });
 });