Fix various sandbox issues.

packages/web-ui/src/tools/artifacts/HtmlArtifact.ts

@@ -5,7 +5,8 @@ import { customElement, property, state } from "lit/decorators.js";
 import { createRef, type Ref, ref } from "lit/directives/ref.js";
 import { unsafeHTML } from "lit/directives/unsafe-html.js";
 import type { SandboxIframe } from "../../components/SandboxedIframe.js";
-import type { Attachment } from "../../utils/attachment-utils.js";
+import { type MessageConsumer, SANDBOX_MESSAGE_ROUTER } from "../../components/sandbox/SandboxMessageRouter.js";
+import type { SandboxRuntimeProvider } from "../../components/sandbox/SandboxRuntimeProvider.js";
 import { i18n } from "../../utils/i18n.js";
 import "../../components/SandboxedIframe.js";
 import { ArtifactElement } from "./ArtifactElement.js";
@@ -16,7 +17,7 @@ import "./Console.js";
 export class HtmlArtifact extends ArtifactElement {
 	@property() override filename = "";
 	@property({ attribute: false }) override displayTitle = "";
-	@property({ attribute: false }) attachments: Attachment[] = [];
+	@property({ attribute: false }) runtimeProviders: SandboxRuntimeProvider[] = [];
 	@property({ attribute: false }) sandboxUrlProvider?: () => string;
 
 	private _content = "";
@@ -26,9 +27,6 @@ export class HtmlArtifact extends ArtifactElement {
 	private sandboxIframeRef: Ref<SandboxIframe> = createRef();
 	private consoleRef: Ref<Console> = createRef();
 
-	// Store message handler so we can remove it
-	private messageHandler?: (e: MessageEvent) => void;
-
 	@state() private viewMode: "preview" | "code" = "preview";
 
 	private setViewMode(mode: "preview" | "code") {
@@ -47,11 +45,17 @@ export class HtmlArtifact extends ArtifactElement {
 		copyButton.title = i18n("Copy HTML");
 		copyButton.showText = false;
 
+		// Generate standalone HTML with all runtime code injected for download
+		const sandbox = this.sandboxIframeRef.value;
+		const sandboxId = `artifact-${this.filename}`;
+		const downloadContent =
+			sandbox?.prepareHtmlDocument(sandboxId, this._content, this.runtimeProviders || []) || this._content;
+
 		return html`
 			<div class="flex items-center gap-2">
 				${toggle}
 				${copyButton}
-				${DownloadButton({ content: this._content, filename: this.filename, mimeType: "text/html", title: i18n("Download HTML") })}
+				${DownloadButton({ content: downloadContent, filename: this.filename, mimeType: "text/html", title: i18n("Download HTML") })}
 			</div>
 		`;
 	}
@@ -79,33 +83,29 @@ export class HtmlArtifact extends ArtifactElement {
 			sandbox.sandboxUrlProvider = this.sandboxUrlProvider;
 		}
 
-		// Remove previous message handler if it exists
-		if (this.messageHandler) {
-			window.removeEventListener("message", this.messageHandler);
-		}
-
 		const sandboxId = `artifact-${this.filename}`;
 
-		// Set up message listener to collect logs
-		this.messageHandler = (e: MessageEvent) => {
-			if (e.data.sandboxId !== sandboxId) return;
-
-			if (e.data.type === "console") {
-				// Create new array reference for Lit reactivity
-				this.logs = [
-					...this.logs,
-					{
-						type: e.data.method === "error" ? "error" : "log",
-						text: e.data.text,
-					},
-				];
-				this.requestUpdate(); // Re-render to show console
-			}
+		// Create consumer for console messages
+		const consumer: MessageConsumer = {
+			handleMessage: (message: any): boolean => {
+				if (message.type === "console") {
+					// Create new array reference for Lit reactivity
+					this.logs = [
+						...this.logs,
+						{
+							type: message.method === "error" ? "error" : "log",
+							text: message.text,
+						},
+					];
+					this.requestUpdate(); // Re-render to show console
+					return true;
+				}
+				return false;
+			},
 		};
-		window.addEventListener("message", this.messageHandler);
 
-		// Load content (iframe persists, doesn't get removed)
-		sandbox.loadContent(sandboxId, html, this.attachments);
+		// Load content - this handles sandbox registration, consumer registration, and iframe creation
+		sandbox.loadContent(sandboxId, html, this.runtimeProviders, [consumer]);
 	}
 
 	override get content(): string {
@@ -114,11 +114,9 @@ export class HtmlArtifact extends ArtifactElement {
 
 	override disconnectedCallback() {
 		super.disconnectedCallback();
-		// Clean up message handler when element is removed from DOM
-		if (this.messageHandler) {
-			window.removeEventListener("message", this.messageHandler);
-			this.messageHandler = undefined;
-		}
+		// Unregister sandbox when element is removed from DOM
+		const sandboxId = `artifact-${this.filename}`;
+		SANDBOX_MESSAGE_ROUTER.unregisterSandbox(sandboxId);
 	}
 
 	override firstUpdated() {

packages/web-ui/src/tools/artifacts/artifacts-tool-renderer.ts

@@ -93,7 +93,11 @@ export class ArtifactsToolRenderer implements ToolRenderer<ArtifactsParams, unde
 
 			// For create/update/rewrite errors, show code block + console/error
 			if (command === "create" || command === "update" || command === "rewrite") {
-				const content = command === "update" ? params?.new_str || params?.old_str || "" : params?.content || "";
+				const content = params?.content || "";
+				const { old_str, new_str } = params || {};
+				const isDiff = command === "update";
+				const diffContent =
+					old_str !== undefined && new_str !== undefined ? Diff({ oldText: old_str, newText: new_str }) : "";
 
 				const isHtml = filename?.endsWith(".html");
 
@@ -101,7 +105,7 @@ export class ArtifactsToolRenderer implements ToolRenderer<ArtifactsParams, unde
 					<div>
 						${renderCollapsibleHeader(state, FileCode2, renderHeaderWithPill(headerText, filename), contentRef, chevronRef, false)}
 						<div ${ref(contentRef)} class="max-h-0 overflow-hidden transition-all duration-300 space-y-3">
-							${content ? html`<code-block .code=${content} language=${getLanguageFromFilename(filename)}></code-block>` : ""}
+							${isDiff ? diffContent : content ? html`<code-block .code=${content} language=${getLanguageFromFilename(filename)}></code-block>` : ""}
 							${
 								isHtml
 									? html`<console-block .content=${result.output || i18n("An error occurred")} variant="error"></console-block>`
@@ -156,7 +160,7 @@ export class ArtifactsToolRenderer implements ToolRenderer<ArtifactsParams, unde
 			}
 
 			// CREATE/UPDATE/REWRITE: always show code block, + console block for .html files
-			if (command === "create" || command === "update" || command === "rewrite") {
+			if (command === "create" || command === "rewrite") {
 				const codeContent = content || "";
 				const isHtml = filename?.endsWith(".html");
 				const logs = result.output || "";
@@ -172,6 +176,20 @@ export class ArtifactsToolRenderer implements ToolRenderer<ArtifactsParams, unde
 				`;
 			}
 
+			if (command === "update") {
+				const isHtml = filename?.endsWith(".html");
+				const logs = result.output || "";
+				return html`
+					<div>
+						${renderCollapsibleHeader(state, FileCode2, renderHeaderWithPill(headerText, filename), contentRef, chevronRef, false)}
+						<div ${ref(contentRef)} class="max-h-0 overflow-hidden transition-all duration-300 space-y-3">
+							${Diff({ oldText: params.old_str || "", newText: params.new_str || "" })}
+							${isHtml && logs ? html`<console-block .content=${logs}></console-block>` : ""}
+						</div>
+					</div>
+				`;
+			}
+
 			// For DELETE, just show header
 			return html`
 				<div class="space-y-3">

packages/web-ui/src/tools/artifacts/artifacts.ts

@@ -6,7 +6,7 @@ import { html, LitElement, type TemplateResult } from "lit";
 import { customElement, property, state } from "lit/decorators.js";
 import { createRef, type Ref, ref } from "lit/directives/ref.js";
 import { X } from "lucide";
-import type { Attachment } from "../../utils/attachment-utils.js";
+import type { SandboxRuntimeProvider } from "../../components/sandbox/SandboxRuntimeProvider.js";
 import { i18n } from "../../utils/i18n.js";
 import type { ArtifactElement } from "./ArtifactElement.js";
 import { HtmlArtifact } from "./HtmlArtifact.js";
@@ -45,8 +45,8 @@ export class ArtifactsPanel extends LitElement {
 	private artifactElements = new Map<string, ArtifactElement>();
 	private contentRef: Ref<HTMLDivElement> = createRef();
 
-	// External provider for attachments (decouples panel from AgentInterface)
-	@property({ attribute: false }) attachmentsProvider?: () => Attachment[];
+	// External factory for runtime providers (decouples panel from AgentInterface)
+	@property({ attribute: false }) runtimeProvidersFactory?: () => SandboxRuntimeProvider[];
 	// Sandbox URL provider for browser extensions (optional)
 	@property({ attribute: false }) sandboxUrlProvider?: () => string;
 	// Callbacks
@@ -108,7 +108,8 @@ export class ArtifactsPanel extends LitElement {
 			const type = this.getFileType(filename);
 			if (type === "html") {
 				element = new HtmlArtifact();
-				(element as HtmlArtifact).attachments = this.attachmentsProvider?.() || [];
+				const runtimeProviders = this.runtimeProvidersFactory?.() || [];
+				(element as HtmlArtifact).runtimeProviders = runtimeProviders;
 				if (this.sandboxUrlProvider) {
 					(element as HtmlArtifact).sandboxUrlProvider = this.sandboxUrlProvider;
 				}
@@ -144,7 +145,8 @@ export class ArtifactsPanel extends LitElement {
 			element.content = content;
 			element.displayTitle = title;
 			if (element instanceof HtmlArtifact) {
-				element.attachments = this.attachmentsProvider?.() || [];
+				const runtimeProviders = this.runtimeProvidersFactory?.() || [];
+				element.runtimeProviders = runtimeProviders;
 			}
 		}
 
@@ -414,46 +416,11 @@ CRITICAL REMINDER FOR ALL ARTIFACTS:
 		}
 
 		return new Promise((resolve) => {
-			let resolved = false;
-
-			// Listen for the execution-complete message
-			const messageHandler = (event: MessageEvent) => {
-				if (event.data?.type === "execution-complete" && event.data?.artifactId === filename) {
-					if (!resolved) {
-						resolved = true;
-						window.removeEventListener("message", messageHandler);
-
-						// Get the logs from the element
-						const logs = element.getLogs();
-						if (logs.includes("[error]")) {
-							resolve(`\n\nExecution completed with errors:\n${logs}`);
-						} else if (logs !== `No logs for ${filename}`) {
-							resolve(`\n\nExecution logs:\n${logs}`);
-						} else {
-							resolve("");
-						}
-					}
-				}
-			};
-
-			window.addEventListener("message", messageHandler);
-
-			// Fallback timeout in case the message never arrives
+			// Fallback timeout - just get logs after execution should complete
 			setTimeout(() => {
-				if (!resolved) {
-					resolved = true;
-					window.removeEventListener("message", messageHandler);
-
-					// Get whatever logs we have so far
-					const logs = element.getLogs();
-					if (logs.includes("[error]")) {
-						resolve(`\n\nExecution timed out with errors:\n${logs}`);
-					} else if (logs !== `No logs for ${filename}`) {
-						resolve(`\n\nExecution timed out. Partial logs:\n${logs}`);
-					} else {
-						resolve("");
-					}
-				}
+				// Get whatever logs we have
+				const logs = element.getLogs();
+				resolve(logs);
 			}, 1500);
 		});
 	}
@@ -568,7 +535,7 @@ CRITICAL REMINDER FOR ALL ARTIFACTS:
 		this.showArtifact(params.filename);
 
 		// For HTML files, wait for execution
-		let result = `Rewrote file ${params.filename}`;
+		let result = "";
 		if (this.getFileType(params.filename) === "html" && !options.skipWait) {
 			const logs = await this.waitForHtmlExecution(params.filename);
 			result += logs;

packages/web-ui/src/tools/javascript-repl.ts

@@ -4,15 +4,15 @@ import { type Static, Type } from "@sinclair/typebox";
 import { createRef, ref } from "lit/directives/ref.js";
 import { Code } from "lucide";
 import { type SandboxFile, SandboxIframe, type SandboxResult } from "../components/SandboxedIframe.js";
+import type { SandboxRuntimeProvider } from "../components/sandbox/SandboxRuntimeProvider.js";
 import type { Attachment } from "../utils/attachment-utils.js";
-
 import { registerToolRenderer, renderCollapsibleHeader, renderHeader } from "./renderer-registry.js";
 import type { ToolRenderer } from "./types.js";
 
 // Execute JavaScript code with attachments using SandboxedIframe
 export async function executeJavaScript(
 	code: string,
-	attachments: Attachment[] = [],
+	runtimeProviders: SandboxRuntimeProvider[],
 	signal?: AbortSignal,
 	sandboxUrlProvider?: () => string,
 ): Promise<{ output: string; files?: SandboxFile[] }> {
@@ -34,8 +34,11 @@ export async function executeJavaScript(
 	document.body.appendChild(sandbox);
 
 	try {
-		const sandboxId = `repl-${Date.now()}`;
-		const result: SandboxResult = await sandbox.execute(sandboxId, code, attachments, signal);
+		const sandboxId = `repl-${Date.now()}-${Math.random().toString(36).substring(7)}`;
+
+		// Pass providers to execute (router handles all message routing)
+		// No additional consumers needed - execute() has its own internal consumer
+		const result: SandboxResult = await sandbox.execute(sandboxId, code, runtimeProviders, [], signal);
 
 		// Remove the sandbox iframe after execution
 		sandbox.remove();
@@ -114,13 +117,13 @@ interface JavaScriptReplResult {
 }
 
 export function createJavaScriptReplTool(): AgentTool<typeof javascriptReplSchema, JavaScriptReplToolResult> & {
-	attachmentsProvider?: () => Attachment[];
+	runtimeProvidersFactory?: () => SandboxRuntimeProvider[];
 	sandboxUrlProvider?: () => string;
 } {
 	return {
 		label: "JavaScript REPL",
 		name: "javascript_repl",
-		attachmentsProvider: () => [], // default to empty array
+		runtimeProvidersFactory: () => [], // default to empty array
 		sandboxUrlProvider: undefined, // optional, for browser extensions
 		description: `Execute JavaScript code in a sandboxed browser environment with full modern browser capabilities.
 
@@ -196,8 +199,12 @@ Global variables:
 - All standard browser globals (window, document, fetch, etc.)`,
 		parameters: javascriptReplSchema,
 		execute: async function (_toolCallId: string, args: Static<typeof javascriptReplSchema>, signal?: AbortSignal) {
-			const attachments = this.attachmentsProvider?.() || [];
-			const result = await executeJavaScript(args.code, attachments, signal, this.sandboxUrlProvider);
+			const result = await executeJavaScript(
+				args.code,
+				this.runtimeProvidersFactory?.() ?? [],
+				signal,
+				this.sandboxUrlProvider,
+			);
 			// Convert files to JSON-serializable with base64 payloads
 			const files = (result.files || []).map((f) => {
 				const toBase64 = (input: string | Uint8Array): { base64: string; size: number } => {