diff --git a/.changeset/agent-document-review-fixes.md b/.changeset/agent-document-review-fixes.md new file mode 100644 index 000000000..7436c816b --- /dev/null +++ b/.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. diff --git a/packages/agent-bundle/src/dev/runtime-routes.ts b/packages/agent-bundle/src/dev/runtime-routes.ts index b21151677..c4958dfd7 100644 --- a/packages/agent-bundle/src/dev/runtime-routes.ts +++ b/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,6 +420,7 @@ export class RuntimeRoutes { )); } try { + const abortController = new AbortController(); const flight = new ReadableStream({ start(controller) { controller.enqueue(asset.body); @@ -425,11 +428,27 @@ export class RuntimeRoutes { }, }); 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)); } } diff --git a/packages/agent-bundle/tests/runtime-routes.test.ts b/packages/agent-bundle/tests/runtime-routes.test.ts index 13cb3a126..fb350cf9c 100644 --- a/packages/agent-bundle/tests/runtime-routes.test.ts +++ b/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 () => { diff --git a/packages/workbench/src/runtime/agent-document-stage.tsx b/packages/workbench/src/runtime/agent-document-stage.tsx index 0e731909f..d4f15d54c 100644 --- a/packages/workbench/src/runtime/agent-document-stage.tsx +++ b/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<{ ; case 'markdown': - return ; + return ; case 'text': return

{node.text}

; case 'context': diff --git a/packages/workbench/src/skill-markdown.tsx b/packages/workbench/src/skill-markdown.tsx index 73a37de7a..a6d6db3f3 100644 --- a/packages/workbench/src/skill-markdown.tsx +++ b/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['components']; + readonly resolveImage: (reference: string) => string | undefined; + readonly resolveLink: (reference: string) => string | undefined; } type MarkdownElementProps = @@ -68,47 +69,61 @@ const SkillCode = ({ children, className }: ComponentPropsWithoutRef<'code'>) => ; }; -const SkillLink = ({ - base, +const MarkdownLink = ({ children, href, node: _node, - resources, + resolve, ...properties -}: MarkdownElementProps<'a'> & Pick) => { - 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 {children}; const external = /^https?:|^mailto:/u.test(resolved); return {children}; }; -const SkillImage = ({ +const MarkdownImage = ({ alt, - base, node: _node, - resources, + resolve, src, ...properties -}: MarkdownElementProps<'img'> & Pick) => { - const resolved = typeof src === 'string' ? resourceUrlFor(base, src, resources) : undefined; - if (resolved === undefined) return {alt ?? 'Image unavailable'}; +}: 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 + {alt ?? 'Image unavailable'} + {source === undefined ? undefined : <> · {source}} + ; + } return {alt; }; /** 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) => (
, code: SkillCode, h1: ({ node: _node, ...properties }) =>

, h2: ({ node: _node, ...properties }) =>

, h3: ({ node: _node, ...properties }) =>

, + img: (properties) => , pre: ({ children }) => <>{children}, table: ({ node: _node, ...properties }) =>
, - ...components, }} remarkPlugins={[remarkGfm, inertHtml]} + urlTransform={(url) => url} > {body} @@ -119,9 +134,7 @@ export const MarkdownProjector = ({ body, components }: MarkdownProjectorProps) export const SkillMarkdown = ({ base, body, resources }: SkillMarkdownProps) => ( , - img: (properties) => , - }} + resolveImage={(reference) => resourceUrlFor(base, reference, resources)} + resolveLink={(reference) => resourceUrlFor(base, reference, resources)} /> ); diff --git a/packages/workbench/src/skills-model.ts b/packages/workbench/src/skills-model.ts index 12987825f..56f36958e 100644 --- a/packages/workbench/src/skills-model.ts +++ b/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; diff --git a/packages/workbench/tests/agent-document-stage.test.ts b/packages/workbench/tests/agent-document-stage.test.ts index 84c37cb02..139d74f7f 100644 --- a/packages/workbench/tests/agent-document-stage.test.ts +++ b/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(/