@copilotkit/runtime 1.54.1-next.6 → 1.55.0-next.7

package/src/v2/runtime/middleware-sse-parser.ts

@@ -0,0 +1,200 @@
+import { logger } from "@copilotkit/shared";
+
+export interface ParsedSSEResult {
+  messages: Message[];
+  threadId?: string;
+  runId?: string;
+}
+
+/** Minimal message shape reconstructed from AG-UI events. */
+export interface Message {
+  id: string;
+  role: string;
+  content?: string;
+  toolCalls?: ToolCall[];
+  toolCallId?: string;
+}
+
+interface ToolCall {
+  id: string;
+  name: string;
+  args: string;
+}
+
+/**
+ * Parse a cloned SSE Response body into structured messages.
+ * Returns empty results for non-SSE responses.
+ */
+export async function parseSSEResponse(
+  response: Response,
+): Promise<ParsedSSEResult> {
+  const contentType = response.headers.get("content-type") ?? "";
+  if (!contentType.includes("text/event-stream")) {
+    return { messages: [] };
+  }
+
+  let text: string;
+  try {
+    text = await response.text();
+  } catch {
+    logger.warn("Failed to read SSE response body in afterRequestMiddleware");
+    return { messages: [] };
+  }
+
+  if (!text.trim()) {
+    return { messages: [] };
+  }
+
+  let threadId: string | undefined;
+  let runId: string | undefined;
+  const messagesById = new Map<string, Message>();
+  const toolCallsById = new Map<string, ToolCall>();
+  const toolCallParent = new Map<string, string>(); // toolCallId → messageId
+  let snapshotMessages: Message[] | undefined;
+
+  for (const line of text.split("\n")) {
+    const trimmed = line.trim();
+    if (!trimmed.startsWith("data:")) continue;
+
+    let event: Record<string, any>;
+    try {
+      event = JSON.parse(trimmed.slice(5).trim());
+    } catch {
+      continue;
+    }
+
+    switch (event.type) {
+      case "RUN_STARTED":
+        threadId = event.threadId;
+        runId = event.runId;
+        break;
+
+      case "MESSAGES_SNAPSHOT":
+        if (Array.isArray(event.messages)) {
+          snapshotMessages = event.messages;
+        }
+        break;
+
+      case "TEXT_MESSAGE_START":
+        messagesById.set(event.messageId, {
+          id: event.messageId,
+          role: event.role ?? "assistant",
+          content: "",
+        });
+        break;
+
+      case "TEXT_MESSAGE_CONTENT": {
+        const msg = messagesById.get(event.messageId);
+        if (msg) {
+          msg.content = (msg.content ?? "") + (event.delta ?? "");
+        }
+        break;
+      }
+
+      case "TEXT_MESSAGE_CHUNK": {
+        // Chunk format: combined start+content. First chunk creates the
+        // message, subsequent chunks append delta to content.
+        if (event.messageId) {
+          const existing = messagesById.get(event.messageId);
+          if (existing) {
+            existing.content = (existing.content ?? "") + (event.delta ?? "");
+          } else {
+            messagesById.set(event.messageId, {
+              id: event.messageId,
+              role: event.role ?? "assistant",
+              content: event.delta ?? "",
+            });
+          }
+        }
+        break;
+      }
+
+      case "TOOL_CALL_START": {
+        const tc: ToolCall = {
+          id: event.toolCallId,
+          name: event.toolCallName,
+          args: "",
+        };
+        toolCallsById.set(event.toolCallId, tc);
+        if (event.parentMessageId) {
+          toolCallParent.set(event.toolCallId, event.parentMessageId);
+        }
+        break;
+      }
+
+      case "TOOL_CALL_ARGS": {
+        const tc = toolCallsById.get(event.toolCallId);
+        if (tc) {
+          tc.args += event.delta ?? "";
+        }
+        break;
+      }
+
+      case "TOOL_CALL_CHUNK": {
+        // Chunk format: combined start+args. First chunk for a given
+        // toolCallId creates the tool call, subsequent chunks append delta.
+        if (event.toolCallId) {
+          let tc = toolCallsById.get(event.toolCallId);
+          if (!tc) {
+            tc = {
+              id: event.toolCallId,
+              name: event.toolCallName ?? "",
+              args: "",
+            };
+            toolCallsById.set(event.toolCallId, tc);
+            if (event.parentMessageId) {
+              toolCallParent.set(event.toolCallId, event.parentMessageId);
+            }
+          }
+          if (event.toolCallName) {
+            tc.name = event.toolCallName;
+          }
+          tc.args += event.delta ?? "";
+        }
+        break;
+      }
+
+      case "TOOL_CALL_END": {
+        const tc = toolCallsById.get(event.toolCallId);
+        const parentId = toolCallParent.get(event.toolCallId);
+        if (tc && parentId) {
+          const parent = messagesById.get(parentId);
+          if (parent) {
+            parent.toolCalls = parent.toolCalls ?? [];
+            parent.toolCalls.push(tc);
+          }
+        }
+        break;
+      }
+
+      case "TOOL_CALL_RESULT":
+        messagesById.set(event.messageId, {
+          id: event.messageId,
+          role: "tool",
+          content: event.content,
+          toolCallId: event.toolCallId,
+        });
+        break;
+    }
+  }
+
+  // Attach any tool calls not yet linked to their parent message.
+  // This handles TOOL_CALL_CHUNK flows which don't emit TOOL_CALL_END.
+  for (const [toolCallId, tc] of toolCallsById) {
+    const parentId = toolCallParent.get(toolCallId);
+    if (!parentId) continue;
+    const parent = messagesById.get(parentId);
+    if (!parent) continue;
+    const alreadyAttached = parent.toolCalls?.some((t) => t.id === tc.id);
+    if (!alreadyAttached) {
+      parent.toolCalls = parent.toolCalls ?? [];
+      parent.toolCalls.push(tc);
+    }
+  }
+
+  // Prefer MESSAGES_SNAPSHOT if present (contains full history).
+  // Otherwise reconstruct from individual events.
+  const messages = snapshotMessages ?? [...messagesById.values()];
+
+  return { messages, threadId, runId };
+}

package/src/v2/runtime/middleware.ts

@@ -0,0 +1,115 @@
+/**
+ * Middleware support for CopilotKit Runtime.
+ *
+ * A middleware hook can be provided as either:
+ *   1. A **callback function** executed in-process.
+ *   2. A **webhook URL** (http/https).  The runtime will `POST` a JSON payload
+ *      to the URL and, for *before* hooks, accept an optional modified
+ *      `Request` object in the response body.
+ *
+ * Two lifecycle hooks are available:
+ *   • `BEFORE_REQUEST` – runs *before* the request handler.
+ *   • `AFTER_REQUEST`  – runs *after* the handler returns a `Response`.
+ */
+
+import type { CopilotRuntimeLike } from "./runtime";
+import type { MaybePromise } from "@copilotkit/shared";
+import { logger } from "@copilotkit/shared";
+import { parseSSEResponse } from "./middleware-sse-parser";
+import type { Message } from "./middleware-sse-parser";
+
+/* ------------------------------------------------------------------------------------------------
+ * Public types
+ * --------------------------------------------------------------------------------------------- */
+
+export interface BeforeRequestMiddlewareParameters {
+  runtime: CopilotRuntimeLike;
+  request: Request;
+  path: string;
+}
+export interface AfterRequestMiddlewareParameters {
+  runtime: CopilotRuntimeLike;
+  response: Response;
+  path: string;
+  /** Reconstructed messages from the SSE stream (empty for non-SSE responses). */
+  messages?: Message[];
+  /** Thread ID extracted from the RUN_STARTED event. */
+  threadId?: string;
+  /** Run ID extracted from the RUN_STARTED event. */
+  runId?: string;
+}
+
+export type BeforeRequestMiddlewareFn = (
+  params: BeforeRequestMiddlewareParameters,
+) => MaybePromise<Request | void>;
+export type AfterRequestMiddlewareFn = (
+  params: AfterRequestMiddlewareParameters,
+) => MaybePromise<void>;
+
+/**
+ * A middleware value can be either a callback function or a webhook URL.
+ */
+export type BeforeRequestMiddleware = BeforeRequestMiddlewareFn;
+export type AfterRequestMiddleware = AfterRequestMiddlewareFn;
+
+/** Lifecycle events emitted to webhook middleware. */
+export enum CopilotKitMiddlewareEvent {
+  BeforeRequest = "BEFORE_REQUEST",
+  AfterRequest = "AFTER_REQUEST",
+}
+
+/** Stages used by the Middleware Webhook Protocol */
+/** Stages used by the CopilotKit webhook protocol */
+export enum WebhookStage {
+  BeforeRequest = "before_request",
+  AfterRequest = "after_request",
+}
+
+/* ------------------------------------------------------------------------------------------------
+ * Internal helpers – (de)serialisation
+ * --------------------------------------------------------------------------------------------- */
+
+export async function callBeforeRequestMiddleware({
+  runtime,
+  request,
+  path,
+}: BeforeRequestMiddlewareParameters): Promise<Request | void> {
+  const mw = runtime.beforeRequestMiddleware;
+  if (!mw) return;
+
+  // Function-based middleware (in-process)
+  if (typeof mw === "function") {
+    return (mw as BeforeRequestMiddlewareFn)({ runtime, request, path });
+  }
+
+  logger.warn({ mw }, "Unsupported beforeRequestMiddleware value – skipped");
+  return;
+}
+
+export async function callAfterRequestMiddleware({
+  runtime,
+  response,
+  path,
+}: {
+  runtime: CopilotRuntimeLike;
+  response: Response;
+  path: string;
+}): Promise<void> {
+  const mw = runtime.afterRequestMiddleware;
+  if (!mw) return;
+
+  const { messages, threadId, runId } = await parseSSEResponse(response);
+
+  if (typeof mw === "function") {
+    return (mw as AfterRequestMiddlewareFn)({
+      runtime,
+      response,
+      path,
+      messages,
+      threadId,
+      runId,
+    });
+  }
+
+  logger.warn({ mw }, "Unsupported afterRequestMiddleware value – skipped");
+}

package/src/v2/runtime/runner/__tests__/finalize-events.test.ts

@@ -0,0 +1,109 @@
+import { describe, expect, it } from "vitest";
+import {
+  BaseEvent,
+  EventType,
+  ToolCallResultEvent,
+  RunErrorEvent,
+} from "@ag-ui/client";
+import { finalizeRunEvents } from "@copilotkit/shared";
+
+const createTextStart = (messageId: string): BaseEvent =>
+  ({
+    type: EventType.TEXT_MESSAGE_START,
+    messageId,
+  }) as BaseEvent;
+
+const createToolStart = (toolCallId: string): BaseEvent =>
+  ({
+    type: EventType.TOOL_CALL_START,
+    toolCallId,
+  }) as BaseEvent;
+
+describe("finalizeRunEvents", () => {
+  it("closes streams with a RUN_FINISHED event when a stop was requested", () => {
+    const events: BaseEvent[] = [
+      createTextStart("msg-1"),
+      createToolStart("tool-1"),
+    ];
+
+    const appended = finalizeRunEvents(events, { stopRequested: true });
+
+    expect(appended.map((event) => event.type)).toEqual([
+      EventType.TEXT_MESSAGE_END,
+      EventType.TOOL_CALL_END,
+      EventType.TOOL_CALL_RESULT,
+      EventType.RUN_FINISHED,
+    ]);
+
+    const resultEvent = appended.find(
+      (event): event is ToolCallResultEvent =>
+        event.type === EventType.TOOL_CALL_RESULT,
+    );
+    expect(JSON.parse(resultEvent?.content ?? "")).toEqual(
+      expect.objectContaining({
+        status: "stopped",
+        reason: "stop_requested",
+      }),
+    );
+
+    expect(events.at(-1)?.type).toBe(EventType.RUN_FINISHED);
+  });
+
+  it("emits a RUN_ERROR with meaningful payload when the stream ends abruptly", () => {
+    const events: BaseEvent[] = [
+      createTextStart("msg-1"),
+      createToolStart("tool-1"),
+    ];
+
+    const appended = finalizeRunEvents(events);
+
+    expect(appended.map((event) => event.type)).toEqual([
+      EventType.TEXT_MESSAGE_END,
+      EventType.TOOL_CALL_END,
+      EventType.TOOL_CALL_RESULT,
+      EventType.RUN_ERROR,
+    ]);
+
+    const resultEvent = appended.find(
+      (event): event is ToolCallResultEvent =>
+        event.type === EventType.TOOL_CALL_RESULT,
+    );
+    expect(JSON.parse(resultEvent?.content ?? "")).toEqual(
+      expect.objectContaining({
+        status: "error",
+        reason: "missing_terminal_event",
+      }),
+    );
+
+    const errorEvent = appended.find(
+      (event): event is RunErrorEvent => event.type === EventType.RUN_ERROR,
+    );
+    expect(errorEvent?.code).toBe("INCOMPLETE_STREAM");
+    expect(errorEvent?.message).toContain("terminal event");
+  });
+
+  it("only appends structural fixes when a terminal event already exists", () => {
+    const events: BaseEvent[] = [
+      createTextStart("msg-1"),
+      createToolStart("tool-1"),
+      { type: EventType.RUN_FINISHED } as BaseEvent,
+    ];
+
+    const appended = finalizeRunEvents(events);
+
+    expect(appended.map((event) => event.type)).toEqual([
+      EventType.TEXT_MESSAGE_END,
+      EventType.TOOL_CALL_END,
+    ]);
+
+    expect(
+      appended.some((event) => event.type === EventType.TOOL_CALL_RESULT),
+    ).toBe(false);
+    expect(appended.some((event) => event.type === EventType.RUN_ERROR)).toBe(
+      false,
+    );
+    expect(
+      appended.some((event) => event.type === EventType.RUN_FINISHED),
+    ).toBe(false);
+  });
+});