@assistant-ui/react-ag-ui 0.0.2

package/dist/useAgUiRuntime.js

@@ -0,0 +1,80 @@
+"use client";
+
+// src/useAgUiRuntime.ts
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import {
+  useExternalStoreRuntime,
+  useRuntimeAdapters
+} from "@assistant-ui/react";
+import { makeLogger } from "./runtime/logger.js";
+import { AgUiThreadRuntimeCore } from "./runtime/AgUiThreadRuntimeCore.js";
+function useAgUiRuntime(options) {
+  const logger = useMemo(() => makeLogger(options.logger), [options.logger]);
+  const [version, setVersion] = useState(0);
+  const notifyUpdate = useCallback(() => setVersion((v) => v + 1), []);
+  const coreRef = useRef(null);
+  const runtimeAdapters = useRuntimeAdapters();
+  const attachmentsAdapter = options.adapters?.attachments ?? runtimeAdapters?.attachments;
+  const historyAdapter = options.adapters?.history ?? runtimeAdapters?.history;
+  const speechAdapter = options.adapters?.speech;
+  const feedbackAdapter = options.adapters?.feedback;
+  if (!coreRef.current) {
+    coreRef.current = new AgUiThreadRuntimeCore({
+      agent: options.agent,
+      logger,
+      showThinking: options.showThinking ?? true,
+      ...options.onError ? { onError: options.onError } : {},
+      ...options.onCancel ? { onCancel: options.onCancel } : {},
+      ...historyAdapter ? { history: historyAdapter } : {},
+      notifyUpdate
+    });
+  }
+  const core = coreRef.current;
+  core.updateOptions({
+    agent: options.agent,
+    logger,
+    showThinking: options.showThinking ?? true,
+    ...options.onError ? { onError: options.onError } : {},
+    ...options.onCancel ? { onCancel: options.onCancel } : {},
+    ...historyAdapter ? { history: historyAdapter } : {}
+  });
+  const adapterAdapters = useMemo(() => {
+    const value = {};
+    if (attachmentsAdapter) value.attachments = attachmentsAdapter;
+    if (speechAdapter) value.speech = speechAdapter;
+    if (feedbackAdapter) value.feedback = feedbackAdapter;
+    return Object.keys(value).length ? value : void 0;
+  }, [attachmentsAdapter, speechAdapter, feedbackAdapter]);
+  const store = useMemo(
+    () => ({
+      messages: core.getMessages(),
+      state: core.getState(),
+      isRunning: core.isRunning(),
+      onNew: (message) => core.append(message),
+      onEdit: (message) => core.edit(message),
+      onReload: (parentId, config) => core.reload(parentId, config),
+      onCancel: () => core.cancel(),
+      onAddToolResult: (options2) => core.addToolResult(options2),
+      onResume: (config) => core.resume(config),
+      setMessages: (messages) => core.applyExternalMessages(messages),
+      onImport: (messages) => core.applyExternalMessages(messages),
+      onLoadExternalState: (state) => core.loadExternalState(state),
+      adapters: adapterAdapters
+    }),
+    // version is intentionally included to trigger re-computation when core state changes via notifyUpdate
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    [adapterAdapters, core, version]
+  );
+  const runtime = useExternalStoreRuntime(store);
+  useEffect(() => {
+    core.attachRuntime(runtime);
+    return () => {
+      core.detachRuntime();
+    };
+  }, [core, runtime]);
+  return runtime;
+}
+export {
+  useAgUiRuntime
+};
+//# sourceMappingURL=useAgUiRuntime.js.map

package/dist/useAgUiRuntime.js.map

@@ -0,0 +1 @@
+{"version":3,"sources":["../src/useAgUiRuntime.ts"],"sourcesContent":["\"use client\";\n\nimport { useCallback, useEffect, useMemo, useRef, useState } from \"react\";\nimport {\n  useExternalStoreRuntime,\n  useRuntimeAdapters,\n} from \"@assistant-ui/react\";\nimport type {\n  AssistantRuntime,\n  AppendMessage,\n  ExternalStoreAdapter,\n  ThreadMessage,\n} from \"@assistant-ui/react\";\nimport type { ReadonlyJSONValue } from \"assistant-stream/utils\";\nimport { makeLogger } from \"./runtime/logger\";\nimport type { UseAgUiRuntimeOptions } from \"./runtime/types\";\nimport { AgUiThreadRuntimeCore } from \"./runtime/AgUiThreadRuntimeCore\";\n\nexport function useAgUiRuntime(\n  options: UseAgUiRuntimeOptions,\n): AssistantRuntime {\n  const logger = useMemo(() => makeLogger(options.logger), [options.logger]);\n  const [version, setVersion] = useState(0);\n  const notifyUpdate = useCallback(() => setVersion((v) => v + 1), []);\n  const coreRef = useRef<AgUiThreadRuntimeCore | null>(null);\n  const runtimeAdapters = useRuntimeAdapters();\n\n  const attachmentsAdapter =\n    options.adapters?.attachments ?? runtimeAdapters?.attachments;\n  const historyAdapter = options.adapters?.history ?? runtimeAdapters?.history;\n  const speechAdapter = options.adapters?.speech;\n  const feedbackAdapter = options.adapters?.feedback;\n\n  if (!coreRef.current) {\n    coreRef.current = new AgUiThreadRuntimeCore({\n      agent: options.agent,\n      logger,\n      showThinking: options.showThinking ?? true,\n      ...(options.onError ? { onError: options.onError } : {}),\n      ...(options.onCancel ? { onCancel: options.onCancel } : {}),\n      ...(historyAdapter ? { history: historyAdapter } : {}),\n      notifyUpdate,\n    });\n  }\n\n  const core = coreRef.current;\n  core.updateOptions({\n    agent: options.agent,\n    logger,\n    showThinking: options.showThinking ?? true,\n    ...(options.onError ? { onError: options.onError } : {}),\n    ...(options.onCancel ? { onCancel: options.onCancel } : {}),\n    ...(historyAdapter ? { history: historyAdapter } : {}),\n  });\n\n  const adapterAdapters = useMemo(() => {\n    const value: NonNullable<ExternalStoreAdapter<ThreadMessage>[\"adapters\"]> =\n      {};\n    if (attachmentsAdapter) value.attachments = attachmentsAdapter;\n    if (speechAdapter) value.speech = speechAdapter;\n    if (feedbackAdapter) value.feedback = feedbackAdapter;\n    return Object.keys(value).length ? value : undefined;\n  }, [attachmentsAdapter, speechAdapter, feedbackAdapter]);\n\n  const store = useMemo(\n    () =>\n      ({\n        messages: core.getMessages(),\n        state: core.getState(),\n        isRunning: core.isRunning(),\n        onNew: (message: AppendMessage) => core.append(message),\n        onEdit: (message: AppendMessage) => core.edit(message),\n        onReload: (parentId: string | null, config: { runConfig?: any }) =>\n          core.reload(parentId, config),\n        onCancel: () => core.cancel(),\n        onAddToolResult: (options) => core.addToolResult(options),\n        onResume: (config) => core.resume(config),\n        setMessages: (messages: readonly ThreadMessage[]) =>\n          core.applyExternalMessages(messages),\n        onImport: (messages: readonly ThreadMessage[]) =>\n          core.applyExternalMessages(messages),\n        onLoadExternalState: (state: ReadonlyJSONValue) =>\n          core.loadExternalState(state),\n        adapters: adapterAdapters,\n      }) satisfies ExternalStoreAdapter<ThreadMessage>,\n    // version is intentionally included to trigger re-computation when core state changes via notifyUpdate\n    // eslint-disable-next-line react-hooks/exhaustive-deps\n    [adapterAdapters, core, version],\n  );\n\n  const runtime = useExternalStoreRuntime(store);\n\n  useEffect(() => {\n    core.attachRuntime(runtime);\n    return () => {\n      core.detachRuntime();\n    };\n  }, [core, runtime]);\n\n  return runtime;\n}\n"],"mappings":";;;AAEA,SAAS,aAAa,WAAW,SAAS,QAAQ,gBAAgB;AAClE;AAAA,EACE;AAAA,EACA;AAAA,OACK;AAQP,SAAS,kBAAkB;AAE3B,SAAS,6BAA6B;AAE/B,SAAS,eACd,SACkB;AAClB,QAAM,SAAS,QAAQ,MAAM,WAAW,QAAQ,MAAM,GAAG,CAAC,QAAQ,MAAM,CAAC;AACzE,QAAM,CAAC,SAAS,UAAU,IAAI,SAAS,CAAC;AACxC,QAAM,eAAe,YAAY,MAAM,WAAW,CAAC,MAAM,IAAI,CAAC,GAAG,CAAC,CAAC;AACnE,QAAM,UAAU,OAAqC,IAAI;AACzD,QAAM,kBAAkB,mBAAmB;AAE3C,QAAM,qBACJ,QAAQ,UAAU,eAAe,iBAAiB;AACpD,QAAM,iBAAiB,QAAQ,UAAU,WAAW,iBAAiB;AACrE,QAAM,gBAAgB,QAAQ,UAAU;AACxC,QAAM,kBAAkB,QAAQ,UAAU;AAE1C,MAAI,CAAC,QAAQ,SAAS;AACpB,YAAQ,UAAU,IAAI,sBAAsB;AAAA,MAC1C,OAAO,QAAQ;AAAA,MACf;AAAA,MACA,cAAc,QAAQ,gBAAgB;AAAA,MACtC,GAAI,QAAQ,UAAU,EAAE,SAAS,QAAQ,QAAQ,IAAI,CAAC;AAAA,MACtD,GAAI,QAAQ,WAAW,EAAE,UAAU,QAAQ,SAAS,IAAI,CAAC;AAAA,MACzD,GAAI,iBAAiB,EAAE,SAAS,eAAe,IAAI,CAAC;AAAA,MACpD;AAAA,IACF,CAAC;AAAA,EACH;AAEA,QAAM,OAAO,QAAQ;AACrB,OAAK,cAAc;AAAA,IACjB,OAAO,QAAQ;AAAA,IACf;AAAA,IACA,cAAc,QAAQ,gBAAgB;AAAA,IACtC,GAAI,QAAQ,UAAU,EAAE,SAAS,QAAQ,QAAQ,IAAI,CAAC;AAAA,IACtD,GAAI,QAAQ,WAAW,EAAE,UAAU,QAAQ,SAAS,IAAI,CAAC;AAAA,IACzD,GAAI,iBAAiB,EAAE,SAAS,eAAe,IAAI,CAAC;AAAA,EACtD,CAAC;AAED,QAAM,kBAAkB,QAAQ,MAAM;AACpC,UAAM,QACJ,CAAC;AACH,QAAI,mBAAoB,OAAM,cAAc;AAC5C,QAAI,cAAe,OAAM,SAAS;AAClC,QAAI,gBAAiB,OAAM,WAAW;AACtC,WAAO,OAAO,KAAK,KAAK,EAAE,SAAS,QAAQ;AAAA,EAC7C,GAAG,CAAC,oBAAoB,eAAe,eAAe,CAAC;AAEvD,QAAM,QAAQ;AAAA,IACZ,OACG;AAAA,MACC,UAAU,KAAK,YAAY;AAAA,MAC3B,OAAO,KAAK,SAAS;AAAA,MACrB,WAAW,KAAK,UAAU;AAAA,MAC1B,OAAO,CAAC,YAA2B,KAAK,OAAO,OAAO;AAAA,MACtD,QAAQ,CAAC,YAA2B,KAAK,KAAK,OAAO;AAAA,MACrD,UAAU,CAAC,UAAyB,WAClC,KAAK,OAAO,UAAU,MAAM;AAAA,MAC9B,UAAU,MAAM,KAAK,OAAO;AAAA,MAC5B,iBAAiB,CAACA,aAAY,KAAK,cAAcA,QAAO;AAAA,MACxD,UAAU,CAAC,WAAW,KAAK,OAAO,MAAM;AAAA,MACxC,aAAa,CAAC,aACZ,KAAK,sBAAsB,QAAQ;AAAA,MACrC,UAAU,CAAC,aACT,KAAK,sBAAsB,QAAQ;AAAA,MACrC,qBAAqB,CAAC,UACpB,KAAK,kBAAkB,KAAK;AAAA,MAC9B,UAAU;AAAA,IACZ;AAAA;AAAA;AAAA,IAGF,CAAC,iBAAiB,MAAM,OAAO;AAAA,EACjC;AAEA,QAAM,UAAU,wBAAwB,KAAK;AAE7C,YAAU,MAAM;AACd,SAAK,cAAc,OAAO;AAC1B,WAAO,MAAM;AACX,WAAK,cAAc;AAAA,IACrB;AAAA,EACF,GAAG,CAAC,MAAM,OAAO,CAAC;AAElB,SAAO;AACT;","names":["options"]}

package/package.json

@@ -0,0 +1,63 @@
+{
+  "name": "@assistant-ui/react-ag-ui",
+  "version": "0.0.2",
+  "license": "MIT",
+  "type": "module",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "default": "./dist/index.js"
+    }
+  },
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "files": [
+    "dist",
+    "src",
+    "README.md"
+  ],
+  "sideEffects": false,
+  "dependencies": {
+    "@ag-ui/client": "^0.0.41",
+    "assistant-stream": "0.2.41"
+  },
+  "peerDependencies": {
+    "@assistant-ui/react": "^0.11.41",
+    "@types/react": "*",
+    "react": "^18 || ^19 || ^19.0.0-rc"
+  },
+  "peerDependenciesMeta": {
+    "@types/react": {
+      "optional": true
+    }
+  },
+  "devDependencies": {
+    "@types/node": "^24.10.1",
+    "@types/react": "^19.2.5",
+    "eslint": "^9",
+    "eslint-config-next": "16.0.3",
+    "react": "19.2.0",
+    "tsx": "^4.20.6",
+    "vitest": "^4.0.10",
+    "jsdom": "^27.2.0",
+    "@assistant-ui/react": "0.11.41",
+    "@assistant-ui/x-buildutils": "0.0.1"
+  },
+  "publishConfig": {
+    "access": "public",
+    "provenance": true
+  },
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/assistant-ui/assistant-ui/tree/main/packages/react-ag-ui"
+  },
+  "bugs": {
+    "url": "https://github.com/assistant-ui/assistant-ui/issues"
+  },
+  "scripts": {
+    "build": "tsx scripts/build.mts",
+    "lint": "eslint .",
+    "test": "vitest",
+    "test:run": "vitest run"
+  }
+}

package/src/index.ts

@@ -0,0 +1,2 @@
+export { useAgUiRuntime } from "./useAgUiRuntime";
+export type { UseAgUiRuntimeOptions } from "./runtime/types";

package/src/runtime/AgUiThreadRuntimeCore.ts

@@ -0,0 +1,484 @@
+"use client";
+
+import { INTERNAL } from "@assistant-ui/react";
+import type {
+  AddToolResultOptions,
+  AppendMessage,
+  AssistantRuntime,
+  ChatModelRunResult,
+  MessageStatus,
+  ThreadAssistantMessage,
+  ThreadHistoryAdapter,
+  ThreadMessage,
+} from "@assistant-ui/react";
+import type { HttpAgent } from "@ag-ui/client";
+import type { Logger } from "./logger";
+import type { AgUiEvent } from "./types";
+import type { ReadonlyJSONValue } from "assistant-stream/utils";
+import { RunAggregator } from "./adapter/run-aggregator";
+import { toAgUiMessages, toAgUiTools } from "./adapter/conversions";
+import { createAgUiSubscriber } from "./adapter/subscriber";
+
+type RunConfig = NonNullable<AppendMessage["runConfig"]>;
+type ResumeRunConfig = {
+  parentId: string | null;
+  sourceId: string | null;
+  runConfig: RunConfig;
+  stream?: unknown;
+};
+
+type CoreOptions = {
+  agent: HttpAgent;
+  logger: Logger;
+  showThinking: boolean;
+  onError?: (error: Error) => void;
+  onCancel?: () => void;
+  history?: ThreadHistoryAdapter;
+  notifyUpdate: () => void;
+};
+
+const FALLBACK_USER_STATUS = { type: "complete", reason: "unknown" } as const;
+
+export class AgUiThreadRuntimeCore {
+  private agent: HttpAgent;
+  private logger: Logger;
+  private showThinking: boolean;
+  private onError: ((error: Error) => void) | undefined;
+  private onCancel: (() => void) | undefined;
+  private readonly notifyUpdate: () => void;
+
+  private runtime: AssistantRuntime | undefined;
+  private messages: ThreadMessage[] = [];
+  private isRunningFlag = false;
+  private abortController: AbortController | null = null;
+  private stateSnapshot: ReadonlyJSONValue | undefined;
+  private pendingError: Error | null = null;
+  private history: ThreadHistoryAdapter | undefined;
+  private lastRunConfig: RunConfig | undefined;
+  private readonly assistantHistoryParents = new Map<string, string | null>();
+  private readonly recordedHistoryIds = new Set<string>();
+
+  constructor(options: CoreOptions) {
+    this.agent = options.agent;
+    this.logger = options.logger;
+    this.showThinking = options.showThinking;
+    this.onError = options.onError;
+    this.onCancel = options.onCancel;
+    this.history = options.history;
+    this.notifyUpdate = options.notifyUpdate;
+  }
+
+  updateOptions(options: Omit<CoreOptions, "notifyUpdate">) {
+    this.agent = options.agent;
+    this.logger = options.logger;
+    this.showThinking = options.showThinking;
+    this.onError = options.onError;
+    this.onCancel = options.onCancel;
+    this.history = options.history;
+  }
+
+  attachRuntime(runtime: AssistantRuntime) {
+    this.runtime = runtime;
+  }
+
+  detachRuntime() {
+    this.runtime = undefined;
+  }
+
+  getMessages(): readonly ThreadMessage[] {
+    return this.messages;
+  }
+
+  getState(): ReadonlyJSONValue | undefined {
+    return this.stateSnapshot;
+  }
+
+  isRunning(): boolean {
+    return this.isRunningFlag;
+  }
+
+  async append(message: AppendMessage): Promise<void> {
+    const startRun = message.startRun ?? message.role === "user";
+    if (message.sourceId) {
+      this.messages = this.messages.filter(
+        (entry) => entry.id !== message.sourceId,
+      );
+    }
+    this.resetHead(message.parentId);
+
+    const threadMessage = this.toThreadMessage(message);
+    this.messages = [...this.messages, threadMessage];
+    this.notifyUpdate();
+    this.recordHistoryEntry(message.parentId ?? null, threadMessage);
+
+    if (!startRun) return;
+    await this.startRun(threadMessage.id, message.runConfig);
+  }
+
+  async edit(message: AppendMessage): Promise<void> {
+    await this.append(message);
+  }
+
+  async reload(
+    parentId: string | null,
+    config: { runConfig?: RunConfig } = {},
+  ): Promise<void> {
+    this.resetHead(parentId);
+    this.notifyUpdate();
+    await this.startRun(parentId, config.runConfig);
+  }
+
+  async cancel(): Promise<void> {
+    if (!this.abortController) return;
+    this.abortController.abort();
+  }
+
+  async resume(config: ResumeRunConfig): Promise<void> {
+    if (config.stream) {
+      this.logger.debug?.(
+        "[agui] resume stream is not supported, falling back to regular run",
+      );
+    }
+    await this.startRun(
+      config.parentId,
+      config.runConfig ?? this.lastRunConfig,
+    );
+  }
+
+  addToolResult(options: AddToolResultOptions): void {
+    let updated = false;
+    this.messages = this.messages.map((message) => {
+      if (message.id !== options.messageId || message.role !== "assistant")
+        return message;
+      updated = true;
+      const assistant = message as ThreadAssistantMessage;
+      const content = assistant.content.map((part) => {
+        if (part.type !== "tool-call" || part.toolCallId !== options.toolCallId)
+          return part;
+        return {
+          ...part,
+          result: options.result,
+          artifact: options.artifact,
+          isError: options.isError,
+        };
+      });
+      return {
+        ...assistant,
+        content,
+      };
+    });
+
+    if (updated) {
+      this.notifyUpdate();
+    }
+  }
+
+  applyExternalMessages(messages: readonly ThreadMessage[]): void {
+    this.assistantHistoryParents.clear();
+    this.messages = [...messages];
+    this.recordedHistoryIds.clear();
+    for (const message of this.messages) {
+      this.recordedHistoryIds.add(message.id);
+    }
+    this.notifyUpdate();
+  }
+
+  loadExternalState(state: ReadonlyJSONValue): void {
+    this.stateSnapshot = state;
+    this.notifyUpdate();
+  }
+
+  private async startRun(
+    parentId: string | null,
+    runConfig?: RunConfig,
+  ): Promise<void> {
+    const normalizedRunConfig = runConfig ?? {};
+    this.lastRunConfig = normalizedRunConfig;
+    this.resetHead(parentId);
+    const historicalMessages = [...this.messages];
+
+    const runId = INTERNAL.generateId();
+    this.pendingError = null;
+    const input = this.buildRunInput(
+      runId,
+      normalizedRunConfig,
+      historicalMessages,
+    );
+    const assistantParentId = parentId ?? this.messages.at(-1)?.id ?? null;
+    let assistantMessageId: string | undefined;
+    const ensureAssistant = () => {
+      if (assistantMessageId) return assistantMessageId;
+      const created = this.insertAssistantPlaceholder();
+      assistantMessageId = created;
+      this.markPendingAssistantHistory(created, assistantParentId ?? null);
+      return created;
+    };
+
+    const aggregator = new RunAggregator({
+      showThinking: this.showThinking,
+      logger: this.logger,
+      emit: (update) => this.updateAssistantMessage(ensureAssistant(), update),
+    });
+    const dispatch = (event: AgUiEvent) => this.handleEvent(aggregator, event);
+
+    const abortController = new AbortController();
+    const abortSignal = abortController.signal;
+    this.abortController = abortController;
+
+    abortSignal.addEventListener(
+      "abort",
+      () => {
+        dispatch({ type: "RUN_CANCELLED" });
+        this.finishRun(abortController);
+        this.onCancel?.();
+      },
+      { once: true },
+    );
+
+    const subscriber = createAgUiSubscriber({
+      dispatch,
+      runId,
+      onRunFailed: (error) => {
+        this.pendingError = error;
+        this.onError?.(error);
+      },
+    });
+
+    aggregator.handle({ type: "RUN_STARTED", runId });
+    this.setRunning(true);
+
+    try {
+      try {
+        (this.agent as any).messages = input.messages;
+        (this.agent as any).threadId = input.threadId;
+      } catch {
+        // ignore
+      }
+      await (this.agent as any).runAgent(input, subscriber, {
+        signal: abortSignal,
+      });
+    } catch (error) {
+      if (!abortSignal.aborted) {
+        const err = error instanceof Error ? error : new Error(String(error));
+        dispatch({ type: "RUN_ERROR", message: err.message });
+        this.onError?.(err);
+        this.pendingError = this.pendingError ?? err;
+      }
+    } finally {
+      this.finishRun(abortController);
+    }
+
+    if (this.pendingError) {
+      const err = this.pendingError;
+      this.pendingError = null;
+      throw err;
+    }
+  }
+
+  private buildRunInput(
+    runId: string,
+    runConfig: RunConfig | undefined,
+    historyMessages: readonly ThreadMessage[] | undefined,
+  ) {
+    const threadId = "main";
+    const messages = toAgUiMessages(historyMessages ?? this.messages);
+    const context = this.runtime?.thread.getModelContext();
+    return {
+      threadId,
+      runId,
+      state: this.stateSnapshot ?? null,
+      messages,
+      tools: toAgUiTools(context?.tools),
+      context: context?.system
+        ? [{ description: "system", value: context.system }]
+        : [],
+      forwardedProps: {
+        ...(context?.callSettings ?? {}),
+        ...(context?.config ?? {}),
+        ...(runConfig?.custom ? { runConfig: runConfig.custom } : {}),
+      },
+    };
+  }
+
+  private setRunning(running: boolean) {
+    this.isRunningFlag = running;
+    this.notifyUpdate();
+  }
+
+  private finishRun(controller: AbortController | null) {
+    if (this.abortController === controller) {
+      this.abortController = null;
+    }
+    this.setRunning(false);
+  }
+
+  private insertAssistantPlaceholder(): string {
+    const id = INTERNAL.generateId();
+    const assistant: ThreadAssistantMessage = {
+      id,
+      role: "assistant",
+      createdAt: new Date(),
+      status: { type: "running" },
+      content: [],
+      metadata: {
+        unstable_state: this.stateSnapshot ?? null,
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+    };
+    this.messages = [...this.messages, assistant];
+    this.notifyUpdate();
+    return id;
+  }
+
+  private updateAssistantMessage(
+    messageId: string,
+    update: ChatModelRunResult,
+  ) {
+    let touched = false;
+    let latestStatus: MessageStatus | undefined;
+    this.messages = this.messages.map((message) => {
+      if (message.id !== messageId || message.role !== "assistant")
+        return message;
+      touched = true;
+      const assistant = message as ThreadAssistantMessage;
+      const metadata = update.metadata
+        ? this.mergeAssistantMetadata(assistant.metadata, update.metadata)
+        : assistant.metadata;
+      latestStatus = update.status ?? assistant.status;
+      return {
+        ...assistant,
+        content: (update.content ??
+          assistant.content) as ThreadAssistantMessage["content"],
+        status: latestStatus,
+        metadata,
+      };
+    });
+    if (touched) {
+      this.notifyUpdate();
+      if (this.isTerminalStatus(latestStatus)) {
+        this.persistAssistantHistory(messageId);
+      }
+    }
+  }
+
+  private mergeAssistantMetadata(
+    current: ThreadAssistantMessage["metadata"],
+    incoming: NonNullable<ChatModelRunResult["metadata"]>,
+  ): ThreadAssistantMessage["metadata"] {
+    const annotations = incoming.unstable_annotations
+      ? [...current.unstable_annotations, ...incoming.unstable_annotations]
+      : current.unstable_annotations;
+    const data = incoming.unstable_data
+      ? [...current.unstable_data, ...incoming.unstable_data]
+      : current.unstable_data;
+    const steps = incoming.steps
+      ? [...current.steps, ...incoming.steps]
+      : current.steps;
+    return {
+      unstable_state:
+        incoming.unstable_state !== undefined
+          ? incoming.unstable_state
+          : current.unstable_state,
+      unstable_annotations: annotations,
+      unstable_data: data,
+      steps,
+      custom: incoming.custom
+        ? { ...current.custom, ...incoming.custom }
+        : current.custom,
+    };
+  }
+
+  private handleEvent(aggregator: RunAggregator, event: AgUiEvent) {
+    switch (event.type) {
+      case "STATE_SNAPSHOT": {
+        this.stateSnapshot = event.snapshot as ReadonlyJSONValue;
+        this.notifyUpdate();
+        return;
+      }
+      case "STATE_DELTA": {
+        this.logger.debug?.("[agui] state delta event ignored", event.delta);
+        return;
+      }
+      case "MESSAGES_SNAPSHOT": {
+        this.importMessagesSnapshot(event.messages);
+        return;
+      }
+      default:
+        aggregator.handle(event);
+    }
+  }
+
+  private importMessagesSnapshot(rawMessages: readonly unknown[]) {
+    try {
+      const converted = rawMessages.map((message) =>
+        INTERNAL.fromThreadMessageLike(
+          message as any,
+          INTERNAL.generateId(),
+          FALLBACK_USER_STATUS,
+        ),
+      );
+      this.applyExternalMessages(converted);
+    } catch (error) {
+      this.logger.error?.("[agui] failed to import messages snapshot", error);
+    }
+  }
+
+  private toThreadMessage(message: AppendMessage): ThreadMessage {
+    return INTERNAL.fromThreadMessageLike(
+      message as any,
+      INTERNAL.generateId(),
+      FALLBACK_USER_STATUS,
+    );
+  }
+
+  private resetHead(parentId: string | null | undefined) {
+    if (!parentId) {
+      if (this.messages.length) {
+        this.messages = [];
+      }
+      return;
+    }
+    const idx = this.messages.findIndex((message) => message.id === parentId);
+    if (idx === -1) return;
+    this.messages = this.messages.slice(0, idx + 1);
+  }
+
+  private isTerminalStatus(status?: MessageStatus): boolean {
+    return status?.type === "complete" || status?.type === "incomplete";
+  }
+
+  private recordHistoryEntry(parentId: string | null, message: ThreadMessage) {
+    this.appendHistoryItem(parentId, message);
+  }
+
+  private markPendingAssistantHistory(
+    messageId: string,
+    parentId: string | null,
+  ) {
+    if (!this.history) return;
+    this.assistantHistoryParents.set(messageId, parentId);
+  }
+
+  private persistAssistantHistory(messageId: string) {
+    if (!this.history) return;
+    const parentId = this.assistantHistoryParents.get(messageId);
+    if (parentId === undefined) return;
+    const message = this.messages.find((m) => m.id === messageId);
+    if (!message || message.role !== "assistant") return;
+    if (!this.isTerminalStatus(message.status)) return;
+    this.assistantHistoryParents.delete(messageId);
+    this.appendHistoryItem(parentId, message);
+  }
+
+  private appendHistoryItem(parentId: string | null, message: ThreadMessage) {
+    if (!this.history || this.recordedHistoryIds.has(message.id)) return;
+    this.recordedHistoryIds.add(message.id);
+    void this.history.append({ parentId, message }).catch((error) => {
+      this.recordedHistoryIds.delete(message.id);
+      this.logger.error?.("[agui] failed to append history entry", error);
+    });
+  }
+}