@copilotkit/shared 1.73.0 → 1.73.2-canary.1790068297

package/src/event-transforms/open-generative-ui-middleware.ts

@@ -0,0 +1,392 @@
+import type {
+  RunAgentInput,
+  AbstractAgent,
+  BaseEvent,
+  ToolCallStartEvent,
+  ToolCallArgsEvent,
+  ActivitySnapshotEventProps,
+  ActivityDeltaEvent,
+} from "@ag-ui/client";
+import { Middleware, EventType } from "@ag-ui/client";
+import { Observable } from "rxjs";
+import clarinet from "clarinet";
+
+const TOOL_NAME = "generateSandboxedUi";
+const ACTIVITY_TYPE = "open-generative-ui";
+
+/**
+ * Parsed parameters from the generateSandboxedUi tool call.
+ */
+export interface GenerateSandboxedUIParams {
+  initialHeight?: number;
+  placeholderMessages?: string[];
+  css?: string;
+  html?: string;
+  jsFunctions?: string;
+  jsExpressions?: string[];
+}
+
+/**
+ * Callback invoked by ArgsParser whenever a parameter (or array item) finishes parsing.
+ */
+export type OnParamEvent = (event: BaseEvent) => void;
+
+/**
+ * Tracks incremental JSON parsing state for a single tool call's arguments.
+ * Emits activity events via the onEvent callback as parameters complete.
+ */
+export class ArgsParser {
+  private parser: ReturnType<typeof clarinet.parser> & {
+    error?: Error | null;
+    textNode?: string;
+  };
+  private currentKey: string | null = null;
+  private depth = 0;
+  private currentArrayKey: string | null = null;
+  private snapshotEmitted = false;
+
+  // Streaming html state — reads parser.textNode to emit incremental chunks
+  private streamingHtmlKey = false;
+  private htmlEmittedLength = 0;
+  private htmlArrayEmitted = false;
+
+  public readonly params: GenerateSandboxedUIParams = {};
+  public readonly messageId: string;
+  private readonly onEvent: OnParamEvent;
+
+  constructor(toolCallId: string, onEvent: OnParamEvent) {
+    this.messageId = `${toolCallId}-activity`;
+    this.onEvent = onEvent;
+    this.parser = clarinet.parser();
+
+    this.parser.onopenobject = (key: string | undefined) => {
+      this.depth++;
+      if (key !== undefined && this.depth === 1) {
+        this.currentKey = key;
+        this.initHtmlStreaming(key);
+      }
+    };
+
+    this.parser.onkey = (key: string) => {
+      if (this.depth === 1) {
+        this.currentKey = key;
+        this.initHtmlStreaming(key);
+      }
+    };
+
+    this.parser.onvalue = (value: string | boolean | number | null) => {
+      if (this.depth === 1 && this.currentKey) {
+        if (this.currentArrayKey) {
+          const strValue = String(value);
+          if (this.currentArrayKey === "jsExpressions") {
+            if (!this.params.jsExpressions) this.params.jsExpressions = [];
+            this.params.jsExpressions.push(strValue);
+          } else if (this.currentArrayKey === "placeholderMessages") {
+            if (!this.params.placeholderMessages)
+              this.params.placeholderMessages = [];
+            this.params.placeholderMessages.push(strValue);
+          }
+          this.emitArrayItemDelta(this.currentArrayKey, strValue);
+        } else if (this.streamingHtmlKey) {
+          // HTML string completed — flush any remaining content immediately + htmlComplete
+          const fullHtml = value != null ? String(value) : "";
+          this.params.html = fullHtml || undefined;
+          this.emitPendingHtml(fullHtml);
+          this.emitParamDelta("htmlComplete", true);
+          this.streamingHtmlKey = false;
+        } else {
+          this.setParam(this.currentKey, value);
+        }
+      }
+    };
+
+    this.parser.onopenarray = () => {
+      if (this.depth === 1 && this.currentKey) {
+        const key = this.currentKey;
+        if (key === "jsExpressions" || key === "placeholderMessages") {
+          this.currentArrayKey = key;
+          if (key === "jsExpressions") this.params.jsExpressions = [];
+          else this.params.placeholderMessages = [];
+          // Emit a delta to create the array in the activity content.
+          // Subsequent "add" ops with path "/<key>/-" append to this array.
+          this.emitParamDelta(key, []);
+        }
+      }
+    };
+
+    this.parser.onclosearray = () => {
+      if (this.depth === 1) {
+        if (this.currentArrayKey === "jsExpressions") {
+          this.emitParamDelta("jsExpressionsComplete", true);
+        }
+        this.currentArrayKey = null;
+      }
+    };
+
+    this.parser.oncloseobject = () => {
+      this.depth--;
+    };
+
+    this.parser.onerror = (err: Error) => {
+      console.warn(
+        "[OpenGenerativeUI] JSON parse error in streaming args, resuming:",
+        err?.message ?? err,
+      );
+      // Reset error state so parsing can continue with the next chunk
+      this.parser.error = null;
+      this.parser.resume();
+    };
+  }
+
+  write(chunk: string): void {
+    this.parser.write(chunk);
+    this.flushHtmlChunks();
+  }
+
+  private initHtmlStreaming(key: string): void {
+    if (key === "html") {
+      this.streamingHtmlKey = true;
+      this.htmlEmittedLength = 0;
+      this.htmlArrayEmitted = false;
+    }
+  }
+
+  /**
+   * Read clarinet's internal textNode buffer to emit html chunks incrementally.
+   * Called after every write() so partial string content is emitted as it streams in.
+   */
+  private flushHtmlChunks(): void {
+    if (!this.streamingHtmlKey) return;
+    const textNode = this.parser.textNode;
+    if (typeof textNode !== "string") return;
+    if (textNode.length === this.htmlEmittedLength) return;
+
+    this.emitPendingHtml(textNode);
+  }
+
+  /**
+   * Emit accumulated html content since the last emission.
+   * Called by flushHtmlChunks and directly when html completes.
+   */
+  private emitPendingHtml(textNode: string): void {
+    const newContent = textNode.slice(this.htmlEmittedLength);
+    if (newContent.length === 0) return;
+
+    if (!this.htmlArrayEmitted) {
+      this.htmlArrayEmitted = true;
+      this.emitParamDelta("html", []);
+    }
+    this.emitArrayItemDelta("html", newContent);
+    this.htmlEmittedLength = textNode.length;
+  }
+
+  private setParam(key: string, value: string | boolean | number | null): void {
+    switch (key) {
+      case "initialHeight":
+        this.params.initialHeight =
+          typeof value === "number" ? value : undefined;
+        if (this.snapshotEmitted) {
+          // Snapshot already went out (another param parsed first) — deliver
+          // the height as a delta instead.
+          this.emitParamDelta("initialHeight", this.params.initialHeight);
+        } else {
+          this.emitSnapshot();
+        }
+        break;
+      case "css":
+        this.params.css = value != null ? String(value) : undefined;
+        this.emitParamDelta("css", this.params.css);
+        this.emitParamDelta("cssComplete", true);
+        break;
+      case "jsFunctions":
+        this.params.jsFunctions = value != null ? String(value) : undefined;
+        this.emitParamDelta("jsFunctions", this.params.jsFunctions);
+        this.emitParamDelta("jsFunctionsComplete", true);
+        break;
+    }
+  }
+
+  private emitSnapshot(): void {
+    if (this.snapshotEmitted) return;
+    this.snapshotEmitted = true;
+
+    const event = {
+      type: EventType.ACTIVITY_SNAPSHOT,
+      messageId: this.messageId,
+      activityType: ACTIVITY_TYPE,
+      content: { initialHeight: this.params.initialHeight, generating: true },
+    } satisfies ActivitySnapshotEventProps;
+    this.onEvent(event);
+  }
+
+  private emitParamDelta(key: string, value: unknown): void {
+    // A JSON Patch "add" requires a value; emitting one without it (e.g. when
+    // the LLM sends a null/empty jsFunctions or css) makes fast-json-patch
+    // reject the whole patch client-side with OPERATION_VALUE_REQUIRED and
+    // drops it. Skip the delta when there's nothing to add.
+    if (value === undefined) return;
+    // The activity message must exist before any delta can be applied —
+    // the client silently drops ACTIVITY_DELTA events whose messageId has
+    // no prior ACTIVITY_SNAPSHOT. The LLM controls the key order of the
+    // streamed args, so the snapshot cannot wait for initialHeight.
+    this.emitSnapshot();
+    const event: ActivityDeltaEvent = {
+      type: EventType.ACTIVITY_DELTA,
+      messageId: this.messageId,
+      activityType: ACTIVITY_TYPE,
+      patch: [{ op: "add", path: `/${key}`, value }],
+    };
+    this.onEvent(event);
+  }
+
+  private emitArrayItemDelta(arrayKey: string, value: string): void {
+    this.emitSnapshot();
+    const event: ActivityDeltaEvent = {
+      type: EventType.ACTIVITY_DELTA,
+      messageId: this.messageId,
+      activityType: ACTIVITY_TYPE,
+      patch: [{ op: "add", path: `/${arrayKey}/-`, value }],
+    };
+    this.onEvent(event);
+  }
+}
+
+/**
+ * Extract EventWithState type from Middleware.runNextWithState return type
+ */
+type ExtractObservableType<T> = T extends Observable<infer U> ? U : never;
+type RunNextWithStateReturn = ReturnType<Middleware["runNextWithState"]>;
+type EventWithState = ExtractObservableType<RunNextWithStateReturn>;
+
+export class OpenGenerativeUIMiddleware extends Middleware {
+  run(input: RunAgentInput, next: AbstractAgent): Observable<BaseEvent> {
+    return this.processStream(this.runNextWithState(input, next));
+  }
+
+  private processStream(
+    source: Observable<EventWithState>,
+  ): Observable<BaseEvent> {
+    return new Observable<BaseEvent>((subscriber) => {
+      let heldRunFinished: EventWithState | null = null;
+      // Track active generateSandboxedUi tool call IDs → their streaming parser
+      const activeParsers = new Map<string, ArgsParser>();
+      // Hold genui tool call events until the first activity event is emitted
+      const heldToolCallEvents = new Map<string, BaseEvent[]>();
+      const flushedToolCalls = new Set<string>();
+
+      const flushHeldEvents = (toolCallId: string) => {
+        if (flushedToolCalls.has(toolCallId)) return;
+        flushedToolCalls.add(toolCallId);
+        const held = heldToolCallEvents.get(toolCallId);
+        if (held) {
+          for (const e of held) {
+            subscriber.next(e);
+          }
+          heldToolCallEvents.delete(toolCallId);
+        }
+      };
+
+      const subscription = source.subscribe({
+        next: (eventWithState) => {
+          const event = eventWithState.event;
+
+          if (heldRunFinished) {
+            subscriber.next(heldRunFinished.event);
+            heldRunFinished = null;
+          }
+
+          if (event.type === EventType.RUN_FINISHED) {
+            heldRunFinished = eventWithState;
+            return;
+          }
+
+          // Hold TOOL_CALL_START for genui until the first activity event
+          if (event.type === EventType.TOOL_CALL_START) {
+            const startEvent = event as ToolCallStartEvent;
+            if (startEvent.toolCallName === TOOL_NAME) {
+              heldToolCallEvents.set(startEvent.toolCallId, [event]);
+              activeParsers.set(
+                startEvent.toolCallId,
+                new ArgsParser(startEvent.toolCallId, (activityEvent) => {
+                  subscriber.next(activityEvent);
+                  flushHeldEvents(startEvent.toolCallId);
+                }),
+              );
+              return;
+            }
+          }
+
+          // Hold or emit TOOL_CALL_ARGS for genui tool calls
+          if (event.type === EventType.TOOL_CALL_ARGS) {
+            const argsEvent = event as ToolCallArgsEvent;
+            const parser = activeParsers.get(argsEvent.toolCallId);
+            if (parser) {
+              if (!flushedToolCalls.has(argsEvent.toolCallId)) {
+                heldToolCallEvents.get(argsEvent.toolCallId)!.push(event);
+              } else {
+                subscriber.next(event);
+              }
+              parser.write(argsEvent.delta);
+              return;
+            }
+          }
+
+          // Hold or emit TOOL_CALL_END for genui tool calls
+          if (event.type === EventType.TOOL_CALL_END) {
+            const endEvent = event as { toolCallId: string } & BaseEvent;
+            const parser = activeParsers.get(endEvent.toolCallId);
+            if (parser) {
+              // Mark generation complete
+              const completeEvent: ActivityDeltaEvent = {
+                type: EventType.ACTIVITY_DELTA,
+                messageId: parser.messageId,
+                activityType: ACTIVITY_TYPE,
+                patch: [{ op: "add", path: "/generating", value: false }],
+              };
+              subscriber.next(completeEvent);
+
+              if (!flushedToolCalls.has(endEvent.toolCallId)) {
+                heldToolCallEvents.get(endEvent.toolCallId)!.push(event);
+              } else {
+                subscriber.next(event);
+              }
+              return;
+            }
+          }
+
+          subscriber.next(event);
+        },
+        error: (err) => {
+          // Flush any held tool call events so downstream sees them before the error
+          for (const [, events] of heldToolCallEvents) {
+            for (const event of events) {
+              subscriber.next(event);
+            }
+          }
+          heldToolCallEvents.clear();
+
+          if (heldRunFinished) {
+            subscriber.next(heldRunFinished.event);
+            heldRunFinished = null;
+          }
+          subscriber.error(err);
+        },
+        complete: () => {
+          // Flush any remaining held tool call events (e.g. parser never emitted)
+          heldToolCallEvents.forEach((_, toolCallId) => {
+            flushHeldEvents(toolCallId);
+          });
+
+          if (heldRunFinished) {
+            subscriber.next(heldRunFinished.event);
+            heldRunFinished = null;
+          }
+          activeParsers.clear();
+          subscriber.complete();
+        },
+      });
+
+      return () => subscription.unsubscribe();
+    });
+  }
+}

package/src/event-transforms/recorded-events.test.ts

@@ -0,0 +1,166 @@
+import { describe, expect, it } from "vitest";
+import { AbstractAgent, EventType } from "@ag-ui/client";
+import type { BaseEvent, RunAgentInput } from "@ag-ui/client";
+import { A2UIMiddleware } from "@ag-ui/a2ui-middleware";
+import { from, firstValueFrom, toArray } from "rxjs";
+import { OpenGenerativeUIMiddleware } from "./open-generative-ui-middleware";
+import { transformRecordedEvents } from "./recorded-events";
+
+const input: RunAgentInput = {
+  threadId: "thread",
+  runId: "run",
+  messages: [],
+  state: {},
+  tools: [],
+  context: [],
+  forwardedProps: {},
+};
+
+function toolRun(name: string, args: string): BaseEvent[] {
+  return [
+    { type: EventType.RUN_STARTED, threadId: "thread", runId: "run" },
+    {
+      type: EventType.TOOL_CALL_START,
+      toolCallId: "call",
+      toolCallName: name,
+      parentMessageId: "message",
+    },
+    { type: EventType.TOOL_CALL_ARGS, toolCallId: "call", delta: args },
+    { type: EventType.TOOL_CALL_END, toolCallId: "call" },
+    {
+      type: EventType.TOOL_CALL_RESULT,
+      toolCallId: "call",
+      messageId: "result",
+      content: "done",
+    },
+    { type: EventType.RUN_FINISHED, threadId: "thread", runId: "run" },
+  ];
+}
+
+class ScriptedAgent extends AbstractAgent {
+  constructor(private readonly events: BaseEvent[]) {
+    super();
+  }
+  run() {
+    return from(this.events);
+  }
+  clone() {
+    return new ScriptedAgent(this.events);
+  }
+}
+
+describe("transformRecordedEvents", () => {
+  it("retains pending tool calls from the recorded input history", async () => {
+    const resumed: RunAgentInput = {
+      ...input,
+      messages: [
+        {
+          id: "prior",
+          role: "assistant",
+          toolCalls: [
+            {
+              id: "pending",
+              type: "function",
+              function: { name: "render_a2ui", arguments: "{}" },
+            },
+          ],
+        },
+      ],
+    };
+    const events: BaseEvent[] = [
+      {
+        type: EventType.TEXT_MESSAGE_START,
+        messageId: "new",
+        role: "assistant",
+      },
+      {
+        type: EventType.TEXT_MESSAGE_CONTENT,
+        messageId: "new",
+        delta: "Hello",
+      },
+      { type: EventType.TEXT_MESSAGE_END, messageId: "new" },
+      { type: EventType.RUN_FINISHED, threadId: "thread", runId: "run" },
+    ];
+    const output = await transformRecordedEvents(resumed, events, {
+      a2ui: {},
+      openGenerativeUI: true,
+    });
+    expect(output).toEqual([
+      ...events.slice(0, -1),
+      expect.objectContaining({
+        type: EventType.TOOL_CALL_RESULT,
+        toolCallId: "pending",
+        content: '{"status":"rendered"}',
+      }),
+      events.at(-1),
+    ]);
+  });
+  it("matches the live middleware chain event-for-event without changing its input", async () => {
+    const events = toolRun(
+      "generateSandboxedUi",
+      JSON.stringify({ initialHeight: 200, html: "<p>Hello</p>" }),
+    );
+    const original = structuredClone(events);
+    const live = new ScriptedAgent(events);
+    live.use(new A2UIMiddleware({}), new OpenGenerativeUIMiddleware());
+    const expected: BaseEvent[] = [];
+    await live.runAgent(
+      { runId: input.runId },
+      {
+        onEvent: ({ event }) => {
+          expected.push(event);
+        },
+      },
+    );
+    const actual = await transformRecordedEvents(input, events, {
+      a2ui: {},
+      openGenerativeUI: true,
+    });
+    expect(actual).toEqual(expected);
+    expect(
+      actual.filter((event) => event.type === EventType.ACTIVITY_SNAPSHOT),
+    ).toEqual([
+      {
+        type: EventType.ACTIVITY_SNAPSHOT,
+        messageId: "call-activity",
+        activityType: "open-generative-ui",
+        content: { initialHeight: 200, generating: true },
+      },
+    ]);
+    expect(events).toEqual(original);
+  });
+
+  it("uses the existing A2UI converter and respects catalog configuration", async () => {
+    const events = toolRun(
+      "render_a2ui",
+      JSON.stringify({
+        surfaceId: "surface",
+        components: [{ id: "root", component: "Text", text: "Hello" }],
+      }),
+    );
+    const config = { defaultCatalogId: "test://catalog" };
+    const expected = await firstValueFrom(
+      new A2UIMiddleware(config)
+        .run(input, new ScriptedAgent(events))
+        .pipe(toArray()),
+    );
+    const actual = await transformRecordedEvents(input, events, {
+      a2ui: config,
+    });
+    expect(actual).toEqual(expected);
+    expect(JSON.stringify(actual)).toContain("test://catalog");
+  });
+
+  it("passes ordinary events and terminal errors through unchanged", async () => {
+    const events: BaseEvent[] = [
+      { type: EventType.RUN_ERROR, message: "recorded failure" },
+    ];
+    expect(
+      await transformRecordedEvents(input, events, {
+        a2ui: {},
+        openGenerativeUI: true,
+      }),
+    ).toEqual(events);
+    expect(await transformRecordedEvents(input, [], {})).toEqual([]);
+  });
+});

package/src/event-transforms/recorded-events.ts

@@ -0,0 +1,74 @@
+import { AbstractAgent } from "@ag-ui/client";
+import type { BaseEvent, RunAgentInput } from "@ag-ui/client";
+import { A2UIMiddleware } from "@ag-ui/a2ui-middleware";
+import type { A2UIMiddlewareConfig } from "@ag-ui/a2ui-middleware";
+import { firstValueFrom, from, toArray } from "rxjs";
+import type { Observable } from "rxjs";
+import { OpenGenerativeUIMiddleware } from "./open-generative-ui-middleware";
+
+export interface RecordedEventTransformOptions {
+  /** Use the original runtime configuration, including its catalog when known. */
+  a2ui?: A2UIMiddlewareConfig;
+  openGenerativeUI?: boolean;
+}
+
+/** An inert event source: it cannot contact an agent or execute tools. */
+class RecordedEventAgent extends AbstractAgent {
+  constructor(
+    private readonly input: RunAgentInput,
+    private readonly stream: (input: RunAgentInput) => Observable<BaseEvent>,
+  ) {
+    super({
+      threadId: input.threadId,
+      initialMessages: input.messages,
+      initialState: input.state,
+    });
+  }
+
+  /** Emit recorded events through this layer without contacting an agent. */
+  run(input: RunAgentInput): Observable<BaseEvent> {
+    return this.stream(input);
+  }
+
+  /** Retain the current replay history and state in an independent agent. */
+  clone(): AbstractAgent {
+    return new RecordedEventAgent(
+      { ...this.input, messages: this.messages, state: this.state },
+      this.stream,
+    );
+  }
+}
+
+/**
+ * Apply CPK's rich-UI middleware to one recorded, pre-middleware run.
+ * Supply complete AG-UI run events and the original input/configuration. This
+ * does not run an agent, fetch resources, or execute tools. Already transformed
+ * histories must not be passed through a second time.
+ *
+ * Returns events in runtime order, including generated activity events and any
+ * synthetic A2UI tool results. Persistence IDs, timestamps, and provenance are
+ * the caller's responsibility (A2UI can generate random result message IDs).
+ */
+export async function transformRecordedEvents(
+  input: RunAgentInput,
+  events: readonly BaseEvent[],
+  options: RecordedEventTransformOptions,
+): Promise<BaseEvent[]> {
+  let agent: AbstractAgent = new RecordedEventAgent(input, () => from(events));
+  // Runtime registers A2UI before Open GenUI; the first middleware is outermost.
+  if (options.openGenerativeUI) {
+    const next = agent;
+    const middleware = new OpenGenerativeUIMiddleware();
+    agent = new RecordedEventAgent(input, (runInput) =>
+      middleware.run(runInput, next),
+    );
+  }
+  if (options.a2ui) {
+    const next = agent;
+    const middleware = new A2UIMiddleware(options.a2ui);
+    agent = new RecordedEventAgent(input, (runInput) =>
+      middleware.run(runInput, next),
+    );
+  }
+  return firstValueFrom(agent.run(input).pipe(toArray()));
+}