@cloudflare/ai-chat 0.0.1 → 0.0.4

package/src/tests/test-utils.ts

@@ -0,0 +1,39 @@
+import { createExecutionContext, env } from "cloudflare:test";
+import { expect } from "vitest";
+import { MessageType, type OutgoingMessage } from "../types";
+import worker from "./worker";
+
+/**
+ * Connects to the chat agent and returns the WebSocket and execution context
+ */
+export async function connectChatWS(
+  path: string
+): Promise<{ ws: WebSocket; ctx: ExecutionContext }> {
+  const ctx = createExecutionContext();
+  const req = new Request(`http://example.com${path}`, {
+    headers: { Upgrade: "websocket" }
+  });
+  const res = await worker.fetch(req, env, ctx);
+  expect(res.status).toBe(101);
+  const ws = res.webSocket as WebSocket;
+  expect(ws).toBeDefined();
+  ws.accept();
+  return { ws, ctx };
+}
+
+/**
+ * Type guard for CF_AGENT_USE_CHAT_RESPONSE messages
+ */
+export function isUseChatResponseMessage(
+  m: unknown
+): m is Extract<
+  OutgoingMessage,
+  { type: MessageType.CF_AGENT_USE_CHAT_RESPONSE }
+> {
+  return (
+    typeof m === "object" &&
+    m !== null &&
+    "type" in m &&
+    m.type === MessageType.CF_AGENT_USE_CHAT_RESPONSE
+  );
+}

package/src/tests/tsconfig.json

@@ -0,0 +1,10 @@
+{
+  "compilerOptions": {
+    "types": [
+      "@cloudflare/workers-types/experimental",
+      "@cloudflare/vitest-pool-workers"
+    ]
+  },
+  "extends": "../../../../tsconfig.base.json",
+  "include": ["./**/*.ts"]
+}

package/src/tests/vitest.config.ts

@@ -0,0 +1,28 @@
+import { defineWorkersConfig } from "@cloudflare/vitest-pool-workers/config";
+
+export default defineWorkersConfig({
+  test: {
+    deps: {
+      optimizer: {
+        ssr: {
+          include: [
+            // vitest can't seem to properly import
+            // `require('./path/to/anything.json')` files,
+            // which ajv uses (by way of @modelcontextprotocol/sdk)
+            // the workaround is to add the package to the include list
+            "ajv"
+          ]
+        }
+      }
+    },
+    poolOptions: {
+      workers: {
+        isolatedStorage: false,
+        singleWorker: true,
+        wrangler: {
+          configPath: "./wrangler.jsonc"
+        }
+      }
+    }
+  }
+});

package/src/tests/worker.ts

@@ -0,0 +1,258 @@
+import { AIChatAgent } from "../";
+import type { UIMessage as ChatMessage } from "ai";
+import { callable, getCurrentAgent, routeAgentRequest } from "agents";
+
+// Type helper for tool call parts - extracts from ChatMessage parts
+type TestToolCallPart = Extract<
+  ChatMessage["parts"][number],
+  { type: `tool-${string}` }
+>;
+
+export type Env = {
+  TestChatAgent: DurableObjectNamespace<TestChatAgent>;
+};
+
+export class TestChatAgent extends AIChatAgent<Env> {
+  observability = undefined;
+  // Store captured context for testing
+  private _capturedContext: {
+    hasAgent: boolean;
+    hasConnection: boolean;
+    connectionId: string | undefined;
+  } | null = null;
+  // Store context captured from nested async function (simulates tool execute)
+  private _nestedContext: {
+    hasAgent: boolean;
+    hasConnection: boolean;
+    connectionId: string | undefined;
+  } | null = null;
+
+  async onChatMessage() {
+    // Capture getCurrentAgent() context for testing
+    const { agent, connection } = getCurrentAgent();
+    this._capturedContext = {
+      hasAgent: agent !== undefined,
+      hasConnection: connection !== undefined,
+      connectionId: connection?.id
+    };
+
+    // Simulate what happens inside a tool's execute function:
+    // It's a nested async function called from within onChatMessage
+    await this._simulateToolExecute();
+
+    // Simple echo response for testing
+    return new Response("Hello from chat agent!", {
+      headers: { "Content-Type": "text/plain" }
+    });
+  }
+
+  // This simulates an AI SDK tool's execute function being called
+  private async _simulateToolExecute(): Promise<void> {
+    // Add a small delay to ensure we're in a new microtask (like real tool execution)
+    await Promise.resolve();
+
+    // Capture context inside the "tool execute" function
+    const { agent, connection } = getCurrentAgent();
+    this._nestedContext = {
+      hasAgent: agent !== undefined,
+      hasConnection: connection !== undefined,
+      connectionId: connection?.id
+    };
+  }
+
+  @callable()
+  getCapturedContext(): {
+    hasAgent: boolean;
+    hasConnection: boolean;
+    connectionId: string | undefined;
+  } | null {
+    return this._capturedContext;
+  }
+
+  @callable()
+  getNestedContext(): {
+    hasAgent: boolean;
+    hasConnection: boolean;
+    connectionId: string | undefined;
+  } | null {
+    return this._nestedContext;
+  }
+
+  @callable()
+  clearCapturedContext(): void {
+    this._capturedContext = null;
+    this._nestedContext = null;
+  }
+
+  @callable()
+  getPersistedMessages(): ChatMessage[] {
+    const rawMessages = (
+      this.sql`select * from cf_ai_chat_agent_messages order by created_at` ||
+      []
+    ).map((row) => {
+      return JSON.parse(row.message as string);
+    });
+    return rawMessages;
+  }
+
+  @callable()
+  async testPersistToolCall(messageId: string, toolName: string) {
+    const toolCallPart: TestToolCallPart = {
+      type: `tool-${toolName}`,
+      toolCallId: `call_${messageId}`,
+      state: "input-available",
+      input: { location: "London" }
+    };
+
+    const messageWithToolCall: ChatMessage = {
+      id: messageId,
+      role: "assistant",
+      parts: [toolCallPart] as ChatMessage["parts"]
+    };
+    await this.persistMessages([messageWithToolCall]);
+    return messageWithToolCall;
+  }
+
+  @callable()
+  async testPersistToolResult(
+    messageId: string,
+    toolName: string,
+    output: string
+  ) {
+    const toolResultPart: TestToolCallPart = {
+      type: `tool-${toolName}`,
+      toolCallId: `call_${messageId}`,
+      state: "output-available",
+      input: { location: "London" },
+      output
+    };
+
+    const messageWithToolOutput: ChatMessage = {
+      id: messageId,
+      role: "assistant",
+      parts: [toolResultPart] as ChatMessage["parts"]
+    };
+    await this.persistMessages([messageWithToolOutput]);
+    return messageWithToolOutput;
+  }
+
+  // Resumable streaming test helpers
+
+  @callable()
+  testStartStream(requestId: string): string {
+    return this._startStream(requestId);
+  }
+
+  @callable()
+  testStoreStreamChunk(streamId: string, body: string): void {
+    this._storeStreamChunk(streamId, body);
+  }
+
+  @callable()
+  testFlushChunkBuffer(): void {
+    this._flushChunkBuffer();
+  }
+
+  @callable()
+  testCompleteStream(streamId: string): void {
+    this._completeStream(streamId);
+  }
+
+  @callable()
+  testMarkStreamError(streamId: string): void {
+    this._markStreamError(streamId);
+  }
+
+  @callable()
+  getActiveStreamId(): string | null {
+    return this._activeStreamId;
+  }
+
+  @callable()
+  getActiveRequestId(): string | null {
+    return this._activeRequestId;
+  }
+
+  @callable()
+  getStreamChunks(
+    streamId: string
+  ): Array<{ body: string; chunk_index: number }> {
+    return (
+      this.sql<{ body: string; chunk_index: number }>`
+        select body, chunk_index from cf_ai_chat_stream_chunks 
+        where stream_id = ${streamId} 
+        order by chunk_index asc
+      ` || []
+    );
+  }
+
+  @callable()
+  getStreamMetadata(
+    streamId: string
+  ): { status: string; request_id: string } | null {
+    const result = this.sql<{ status: string; request_id: string }>`
+      select status, request_id from cf_ai_chat_stream_metadata 
+      where id = ${streamId}
+    `;
+    return result && result.length > 0 ? result[0] : null;
+  }
+
+  @callable()
+  getAllStreamMetadata(): Array<{
+    id: string;
+    status: string;
+    request_id: string;
+    created_at: number;
+  }> {
+    return (
+      this.sql<{
+        id: string;
+        status: string;
+        request_id: string;
+        created_at: number;
+      }>`select id, status, request_id, created_at from cf_ai_chat_stream_metadata` ||
+      []
+    );
+  }
+
+  @callable()
+  testInsertStaleStream(
+    streamId: string,
+    requestId: string,
+    ageMs: number
+  ): void {
+    const createdAt = Date.now() - ageMs;
+    this.sql`
+      insert into cf_ai_chat_stream_metadata (id, request_id, status, created_at)
+      values (${streamId}, ${requestId}, 'streaming', ${createdAt})
+    `;
+  }
+
+  @callable()
+  testRestoreActiveStream(): void {
+    this._restoreActiveStream();
+  }
+}
+
+export default {
+  async fetch(request: Request, env: Env, _ctx: ExecutionContext) {
+    const url = new URL(request.url);
+
+    if (url.pathname === "/500") {
+      return new Response("Internal Server Error", { status: 500 });
+    }
+
+    return (
+      (await routeAgentRequest(request, env)) ||
+      new Response("Not found", { status: 404 })
+    );
+  },
+
+  async email(
+    _message: ForwardableEmailMessage,
+    _env: Env,
+    _ctx: ExecutionContext
+  ) {
+    // Bring this in when we write tests for the complete email handler flow
+  }
+};

package/src/tests/wrangler.jsonc

@@ -0,0 +1,26 @@
+{
+  "compatibility_date": "2025-04-17",
+  "compatibility_flags": [
+    "nodejs_compat",
+    // adding these flags since the vitest runner needs them
+    "enable_nodejs_tty_module",
+    "enable_nodejs_fs_module",
+    "enable_nodejs_http_modules",
+    "enable_nodejs_perf_hooks_module"
+  ],
+  "durable_objects": {
+    "bindings": [
+      {
+        "class_name": "TestChatAgent",
+        "name": "TestChatAgent"
+      }
+    ]
+  },
+  "main": "worker.ts",
+  "migrations": [
+    {
+      "new_sqlite_classes": ["TestChatAgent"],
+      "tag": "v1"
+    }
+  ]
+}

package/src/types.ts

@@ -0,0 +1,122 @@
+import type { UIMessage } from "ai";
+
+/**
+ * Enum for message types to improve type safety and maintainability
+ */
+export enum MessageType {
+  CF_AGENT_CHAT_MESSAGES = "cf_agent_chat_messages",
+  CF_AGENT_USE_CHAT_REQUEST = "cf_agent_use_chat_request",
+  CF_AGENT_USE_CHAT_RESPONSE = "cf_agent_use_chat_response",
+  CF_AGENT_CHAT_CLEAR = "cf_agent_chat_clear",
+  CF_AGENT_CHAT_REQUEST_CANCEL = "cf_agent_chat_request_cancel",
+
+  /** Sent by server when client connects and there's an active stream to resume */
+  CF_AGENT_STREAM_RESUMING = "cf_agent_stream_resuming",
+  /** Sent by client to acknowledge stream resuming notification and request chunks */
+  CF_AGENT_STREAM_RESUME_ACK = "cf_agent_stream_resume_ack",
+
+  /** Client sends tool result to server (for client-side tools) */
+  CF_AGENT_TOOL_RESULT = "cf_agent_tool_result",
+  /** Server notifies client that a message was updated (e.g., tool result applied) */
+  CF_AGENT_MESSAGE_UPDATED = "cf_agent_message_updated"
+}
+
+/**
+ * Types of messages sent from the Agent to clients
+ */
+export type OutgoingMessage<ChatMessage extends UIMessage = UIMessage> =
+  | {
+      /** Indicates this message is a command to clear chat history */
+      type: MessageType.CF_AGENT_CHAT_CLEAR;
+    }
+  | {
+      /** Indicates this message contains updated chat messages */
+      type: MessageType.CF_AGENT_CHAT_MESSAGES;
+      /** Array of chat messages */
+      messages: ChatMessage[];
+    }
+  | {
+      /** Indicates this message is a response to a chat request */
+      type: MessageType.CF_AGENT_USE_CHAT_RESPONSE;
+      /** Unique ID of the request this response corresponds to */
+      id: string;
+      /** Content body of the response */
+      body: string;
+      /** Whether this is the final chunk of the response */
+      done: boolean;
+      /** Whether this response contains an error */
+      error?: boolean;
+      /** Whether this is a continuation (append to last assistant message) */
+      continuation?: boolean;
+    }
+  | {
+      /** Indicates the server is resuming an active stream */
+      type: MessageType.CF_AGENT_STREAM_RESUMING;
+      /** The request ID of the stream being resumed */
+      id: string;
+    }
+  | {
+      /** Server notifies client that a message was updated (e.g., tool result applied) */
+      type: MessageType.CF_AGENT_MESSAGE_UPDATED;
+      /** The updated message */
+      message: ChatMessage;
+    };
+
+/**
+ * Types of messages sent from clients to the Agent
+ */
+export type IncomingMessage<ChatMessage extends UIMessage = UIMessage> =
+  | {
+      /** Indicates this message is a command to clear chat history */
+      type: MessageType.CF_AGENT_CHAT_CLEAR;
+    }
+  | {
+      /** Indicates this message is a request to the chat API */
+      type: MessageType.CF_AGENT_USE_CHAT_REQUEST;
+      /** Unique ID for this request */
+      id: string;
+      /** Request initialization options */
+      init: Pick<
+        RequestInit,
+        | "method"
+        | "keepalive"
+        | "headers"
+        | "body"
+        | "redirect"
+        | "integrity"
+        | "credentials"
+        | "mode"
+        | "referrer"
+        | "referrerPolicy"
+        | "window"
+      >;
+    }
+  | {
+      /** Indicates this message contains updated chat messages */
+      type: MessageType.CF_AGENT_CHAT_MESSAGES;
+      /** Array of chat messages */
+      messages: ChatMessage[];
+    }
+  | {
+      /** Indicates the user wants to stop generation of this message */
+      type: MessageType.CF_AGENT_CHAT_REQUEST_CANCEL;
+      id: string;
+    }
+  | {
+      /** Client acknowledges stream resuming notification and is ready to receive chunks */
+      type: MessageType.CF_AGENT_STREAM_RESUME_ACK;
+      /** The request ID of the stream being resumed */
+      id: string;
+    }
+  | {
+      /** Client sends tool result to server (for client-side tools) */
+      type: MessageType.CF_AGENT_TOOL_RESULT;
+      /** The tool call ID this result is for */
+      toolCallId: string;
+      /** The name of the tool */
+      toolName: string;
+      /** The output from the tool execution */
+      output: unknown;
+      /** Whether server should auto-continue the conversation after applying result */
+      autoContinue?: boolean;
+    };

package/tsconfig.json

@@ -0,0 +1,4 @@
+{
+  "exclude": ["src/tests/**/*.ts", "src/e2e/**/*.ts"],
+  "extends": "../../tsconfig.base.json"
+}