@ekodb/ekodb-client 0.12.0 → 0.14.0

package/src/client.ts

@@ -4,7 +4,7 @@
 
 import { encode, decode } from "@msgpack/msgpack";
 import { QueryBuilder, Query as QueryBuilderQuery } from "./query-builder";
-import { SearchQuery, SearchQueryBuilder, SearchResponse } from "./search";
+import { SearchQuery, SearchResponse } from "./search";
 import { Schema, SchemaBuilder, CollectionMetadata } from "./schema";
 import { Script, FunctionResult } from "./functions";
 
@@ -46,8 +46,6 @@ export interface ClientConfig {
   shouldRetry?: boolean;
   /** Maximum number of retry attempts (default: 3) */
   maxRetries?: number;
-  /** Request timeout in milliseconds (default: 30000) */
-  timeout?: number;
   /** Serialization format (default: MessagePack for best performance, use Json for debugging) */
   format?: SerializationFormat;
 }
@@ -134,6 +132,26 @@ export interface BatchDeleteOptions {
   transactionId?: string;
 }
 
+export interface DistinctValuesOptions {
+  /** Optional filter expression (same format as find() filter). */
+  filter?: any;
+  /** Bypass ripple propagation for this query. */
+  bypassRipple?: boolean;
+  /** Bypass cache for this query. */
+  bypassCache?: boolean;
+}
+
+export interface DistinctValuesResponse {
+  /** Collection that was queried. */
+  collection: string;
+  /** Field whose distinct values were returned. */
+  field: string;
+  /** Unique values, sorted alphabetically. */
+  values: any[];
+  /** Number of distinct values. */
+  count: number;
+}
+
 // ========== Chat Interfaces ==========
 
 export interface CollectionConfig {
@@ -298,6 +316,24 @@ export interface EmbedResponse {
   dimensions: number;
 }
 
+/**
+ * Request for stateless raw LLM completion — no session, no history, no RAG.
+ */
+export interface RawCompletionRequest {
+  system_prompt: string;
+  message: string;
+  provider?: string;
+  model?: string;
+  max_tokens?: number;
+}
+
+/**
+ * Response from a raw LLM completion request.
+ */
+export interface RawCompletionResponse {
+  content: string;
+}
+
 /**
  * User function definition - reusable sequence of Functions that can be called by Scripts
  */
@@ -335,9 +371,9 @@ export class EkoDBClient {
   private baseURL: string;
   private apiKey: string;
   private token: string | null = null;
+  private tokenExpiry: number = 0;
   private shouldRetry: boolean;
   private maxRetries: number;
-  private timeout: number;
   private format: SerializationFormat;
   private rateLimitInfo: RateLimitInfo | null = null;
 
@@ -348,14 +384,12 @@ export class EkoDBClient {
       this.apiKey = apiKey!;
       this.shouldRetry = true;
       this.maxRetries = 3;
-      this.timeout = 30000;
       this.format = SerializationFormat.MessagePack; // Default to MessagePack for 2-3x performance
     } else {
       this.baseURL = config.baseURL;
       this.apiKey = config.apiKey;
       this.shouldRetry = config.shouldRetry ?? true;
       this.maxRetries = config.maxRetries ?? 3;
-      this.timeout = config.timeout ?? 30000;
       this.format = config.format ?? SerializationFormat.MessagePack; // Default to MessagePack for 2-3x performance
     }
   }
@@ -386,7 +420,7 @@ export class EkoDBClient {
   /**
    * Refresh the authentication token
    */
-  private async refreshToken(): Promise<void> {
+  async refreshToken(): Promise<void> {
     const response = await fetch(`${this.baseURL}/api/auth/token`, {
       method: "POST",
       headers: { "Content-Type": "application/json" },
@@ -408,6 +442,74 @@ export class EkoDBClient {
 
     const result = (await response.json()) as { token: string };
     this.token = result.token;
+
+    // Extract and cache JWT expiry for proactive refresh
+    const expiry = this.extractJWTExpiry(result.token);
+    this.tokenExpiry = expiry ?? Math.floor(Date.now() / 1000) + 3600; // fallback: 1 hour
+  }
+
+  /**
+   * Get a valid authentication token.
+   *
+   * Returns a cached token if it has more than 60s of validity remaining.
+   * Otherwise fetches a new one via refreshToken(). This means callers
+   * never need to handle token refresh themselves — every getToken() call
+   * returns a token that's valid for at least 60 more seconds.
+   */
+  async getToken(): Promise<string | null> {
+    if (this.token) {
+      const now = Math.floor(Date.now() / 1000);
+      if (now + 60 >= this.tokenExpiry) {
+        // Token is about to expire or already expired — refresh proactively
+        await this.refreshToken();
+      }
+    } else {
+      // No token yet — fetch one
+      await this.refreshToken();
+    }
+    return this.token;
+  }
+
+  /**
+   * Clear the cached authentication token and expiry.
+   * The next request will trigger a fresh token exchange.
+   */
+  clearTokenCache(): void {
+    this.token = null;
+    this.tokenExpiry = 0;
+  }
+
+  /**
+   * Extract the `exp` claim from a JWT without verifying the signature.
+   * Returns the Unix timestamp (seconds) of expiry, or null if parsing fails.
+   */
+  private extractJWTExpiry(token: string): number | null {
+    try {
+      const parts = token.split(".");
+      if (parts.length !== 3) {
+        return null;
+      }
+
+      // Convert base64url to standard base64
+      let payload = parts[1];
+      payload = payload.replace(/-/g, "+").replace(/_/g, "/");
+
+      // Pad to multiple of 4
+      const pad = payload.length % 4;
+      if (pad) {
+        payload += "=".repeat(4 - pad);
+      }
+
+      const decoded = atob(payload);
+      const claims = JSON.parse(decoded);
+
+      if (typeof claims.exp === "number") {
+        return claims.exp;
+      }
+      return null;
+    } catch {
+      return null;
+    }
   }
 
   /**
@@ -660,6 +762,32 @@ export class EkoDBClient {
     return this.makeRequest<Record>("GET", `/api/find/${collection}/${id}`);
   }
 
+  /**
+   * Find a document by ID with field projection
+   * @param collection - Collection name
+   * @param id - Document ID
+   * @param selectFields - Fields to include in the result
+   * @param excludeFields - Fields to exclude from the result
+   */
+  async findByIdWithProjection(
+    collection: string,
+    id: string,
+    selectFields?: string[],
+    excludeFields?: string[],
+  ): Promise<Record> {
+    const params = new URLSearchParams();
+    if (selectFields?.length) {
+      params.append("select_fields", selectFields.join(","));
+    }
+    if (excludeFields?.length) {
+      params.append("exclude_fields", excludeFields.join(","));
+    }
+    const url = params.toString()
+      ? `/api/find/${collection}/${id}?${params.toString()}`
+      : `/api/find/${collection}/${id}`;
+    return this.makeRequest<Record>("GET", url);
+  }
+
   /**
    * Update a document
    * @param collection - Collection name
@@ -688,6 +816,52 @@ export class EkoDBClient {
     return this.makeRequest<Record>("PUT", url, record);
   }
 
+  /**
+   * Apply an atomic field action to a single field of a record.
+   *
+   * Use this instead of `update()` for safe concurrent modifications like
+   * incrementing counters, pushing to arrays, or arithmetic operations.
+   *
+   * @param collection - Collection name
+   * @param id - Record ID
+   * @param action - The atomic action: increment, decrement, multiply, divide, modulo,
+   *                 push, pop, shift, unshift, remove, append, clear
+   * @param field - The field name to apply the action to
+   * @param value - The value for the action (omit for pop/shift/clear)
+   */
+  async updateWithAction(
+    collection: string,
+    id: string,
+    action: string,
+    field: string,
+    value?: any,
+  ): Promise<Record> {
+    const url = `/api/update/${collection}/${id}/action/${action}`;
+    return this.makeRequest<Record>("PUT", url, {
+      field,
+      value: value ?? null,
+    });
+  }
+
+  /**
+   * Apply a sequence of atomic field actions to a record in a single request.
+   *
+   * All actions are applied atomically — the record is fetched once, all actions
+   * run in order, and the result is persisted in a single update.
+   *
+   * @param collection - Collection name
+   * @param id - Record ID
+   * @param actions - Array of [action, field, value] tuples
+   */
+  async updateWithActionSequence(
+    collection: string,
+    id: string,
+    actions: [string, string, any][],
+  ): Promise<Record> {
+    const url = `/api/update/sequence/${collection}/${id}`;
+    return this.makeRequest<Record>("PUT", url, actions);
+  }
+
   /**
    * Delete a document
    * @param collection - Collection name
@@ -1313,6 +1487,51 @@ export class EkoDBClient {
     );
   }
 
+  /**
+   * Get distinct (unique) values for a field across all records in a collection.
+   *
+   * Results are deduplicated and sorted alphabetically. Supports an optional filter
+   * to restrict which records are examined.
+   *
+   * @param collection - Collection name
+   * @param field - Field to get distinct values for
+   * @param options - Optional filter and bypass flags
+   *
+   * @example
+   * // All distinct statuses
+   * const resp = await client.distinctValues("orders", "status");
+   * console.log(resp.values); // ["active", "cancelled", "shipped"]
+   *
+   * // Only statuses for US orders
+   * const resp = await client.distinctValues("orders", "status", {
+   *   filter: { type: "Condition", content: { field: "region", operator: "Eq", value: "us" } }
+   * });
+   */
+  async distinctValues(
+    collection: string,
+    field: string,
+    options: DistinctValuesOptions = {},
+  ): Promise<DistinctValuesResponse> {
+    const body: {
+      filter?: any;
+      bypass_ripple?: boolean;
+      bypass_cache?: boolean;
+    } = {};
+    if (options.filter !== undefined) body.filter = options.filter;
+    if (options.bypassRipple !== undefined)
+      body.bypass_ripple = options.bypassRipple;
+    if (options.bypassCache !== undefined)
+      body.bypass_cache = options.bypassCache;
+
+    return this.makeRequest<DistinctValuesResponse>(
+      "POST",
+      `/api/distinct/${collection}/${field}`,
+      body,
+      0,
+      true, // Force JSON
+    );
+  }
+
   /**
    * Health check - verify the ekoDB server is responding
    */
@@ -1348,6 +1567,143 @@ export class EkoDBClient {
     );
   }
 
+  /**
+   * Stateless raw LLM completion — no session, no history, no RAG.
+   *
+   * Sends a system prompt and user message directly to the LLM via ekoDB
+   * and returns the raw text response without any context injection or
+   * conversation management. Use this for structured-output tasks such as
+   * planning where the response must be parsed programmatically.
+   *
+   * @example
+   * const resp = await client.rawCompletion({
+   *   system_prompt: "You are a helpful assistant.",
+   *   message: "Summarize this in JSON.",
+   *   max_tokens: 2048,
+   * });
+   * console.log(resp.content);
+   */
+  async rawCompletion(
+    request: RawCompletionRequest,
+  ): Promise<RawCompletionResponse> {
+    return this.makeRequest<RawCompletionResponse>(
+      "POST",
+      "/api/chat/complete",
+      request,
+      0,
+      true, // Force JSON
+    );
+  }
+
+  /**
+   * Stateless raw LLM completion via SSE streaming.
+   *
+   * Same as rawCompletion() but uses Server-Sent Events to keep the
+   * connection alive. Preferred for deployed instances where reverse proxies
+   * may kill idle HTTP connections before the LLM responds.
+   */
+  async rawCompletionStream(
+    request: RawCompletionRequest,
+  ): Promise<RawCompletionResponse> {
+    let token = await this.getToken();
+    const url = `${this.baseURL}/api/chat/complete/stream`;
+
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Accept: "text/event-stream",
+        Authorization: `Bearer ${token}`,
+      },
+      body: JSON.stringify(request),
+    });
+
+    if (!response.ok) {
+      const body = await response.text();
+      throw new Error(
+        `SSE raw completion failed (${response.status}): ${body}`,
+      );
+    }
+
+    const body = await response.text();
+    let content = "";
+    let lastError: string | null = null;
+
+    for (const line of body.split("\n")) {
+      if (line.startsWith("data:")) {
+        const dataStr = line.slice(5).trim();
+        if (!dataStr) continue;
+        try {
+          const eventData = JSON.parse(dataStr);
+          if (eventData.token) content += eventData.token;
+          if (eventData.content) content = eventData.content;
+          if (eventData.error) lastError = eventData.error;
+        } catch {
+          // skip malformed SSE data
+        }
+      }
+    }
+
+    if (lastError) throw new Error(lastError);
+    return { content };
+  }
+
+  /**
+   * Stateless raw LLM completion via SSE streaming with token-level progress.
+   *
+   * Same as rawCompletionStream() but invokes `onToken` with each token as it
+   * arrives, allowing callers to show real-time progress.
+   */
+  async rawCompletionStreamWithProgress(
+    request: RawCompletionRequest,
+    onToken: (token: string) => void,
+  ): Promise<RawCompletionResponse> {
+    let token = await this.getToken();
+    const url = `${this.baseURL}/api/chat/complete/stream`;
+
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Accept: "text/event-stream",
+        Authorization: `Bearer ${token}`,
+      },
+      body: JSON.stringify(request),
+    });
+
+    if (!response.ok) {
+      const body = await response.text();
+      throw new Error(
+        `SSE raw completion failed (${response.status}): ${body}`,
+      );
+    }
+
+    const body = await response.text();
+    let content = "";
+    let lastError: string | null = null;
+
+    for (const line of body.split("\n")) {
+      if (line.startsWith("data:")) {
+        const dataStr = line.slice(5).trim();
+        if (!dataStr) continue;
+        try {
+          const eventData = JSON.parse(dataStr);
+          if (eventData.token) {
+            content += eventData.token;
+            onToken(eventData.token);
+          }
+          if (eventData.content) content = eventData.content;
+          if (eventData.error) lastError = eventData.error;
+        } catch {
+          // skip malformed SSE data
+        }
+      }
+    }
+
+    if (lastError) throw new Error(lastError);
+    return { content };
+  }
+
   /**
    * Send a message in an existing chat session
    */
@@ -1364,6 +1720,96 @@ export class EkoDBClient {
     );
   }
 
+  /**
+   * Send a message in an existing chat session via SSE streaming.
+   *
+   * Returns an EventStream that emits ChatStreamEvent objects as they arrive:
+   * - `{ type: "chunk", content: "..." }` for each token
+   * - `{ type: "end", messageId, executionTimeMs, tokenUsage?, contextWindow? }` when complete
+   * - `{ type: "error", error: "..." }` on failure
+   *
+   * Preferred over chatMessage() for long-running responses where reverse
+   * proxies may kill idle HTTP connections before the LLM responds.
+   */
+  chatMessageStream(
+    chatId: string,
+    request: ChatMessageRequest,
+  ): EventStream<ChatStreamEvent> {
+    const stream = new EventStream<ChatStreamEvent>();
+
+    (async () => {
+      try {
+        let token = this.getToken();
+        if (!token) {
+          await this.refreshToken();
+          token = this.getToken();
+        }
+        const url = `${this.baseURL}/api/chat/${chatId}/messages/stream`;
+
+        const response = await fetch(url, {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+            Accept: "text/event-stream",
+            Authorization: `Bearer ${token}`,
+          },
+          body: JSON.stringify(request),
+        });
+
+        if (!response.ok) {
+          const body = await response.text();
+          stream.emit("event", {
+            type: "error",
+            error: `SSE chat message stream failed (${response.status}): ${body}`,
+          } as ChatStreamEvent);
+          stream.close();
+          return;
+        }
+
+        const body = await response.text();
+        for (const line of body.split("\n")) {
+          if (!line.startsWith("data:")) continue;
+          const dataStr = line.slice(5).trim();
+          if (!dataStr) continue;
+          try {
+            const eventData = JSON.parse(dataStr);
+            if (eventData.error) {
+              stream.emit("event", {
+                type: "error",
+                error: eventData.error,
+              } as ChatStreamEvent);
+            } else if (eventData.content && eventData.message_id) {
+              // Done event — has full content + message_id
+              stream.emit("event", {
+                type: "end",
+                messageId: eventData.message_id,
+                executionTimeMs: eventData.execution_time_ms ?? 0,
+                tokenUsage: eventData.token_usage,
+                contextWindow: eventData.context_window,
+              } as ChatStreamEvent);
+            } else if (eventData.token) {
+              stream.emit("event", {
+                type: "chunk",
+                content: eventData.token,
+              } as ChatStreamEvent);
+            }
+          } catch {
+            // skip malformed SSE data
+          }
+        }
+        stream.close();
+      } catch (err: any) {
+        stream.emit("event", {
+          type: "error",
+          error: err.message ?? String(err),
+        } as ChatStreamEvent);
+        stream.close();
+      }
+    })();
+
+    return stream;
+  }
+
   /**
    * Get a chat session by ID
    */
@@ -1560,6 +2006,21 @@ export class EkoDBClient {
     );
   }
 
+  /**
+   * Get all built-in server-side chat tool definitions.
+   * Returns a list of tool objects with name, description, and parameters fields.
+   * Used by planning agents to discover available tools dynamically.
+   */
+  async getChatTools(): Promise<object[]> {
+    return this.makeRequest<object[]>(
+      "GET",
+      "/api/chat/tools",
+      undefined,
+      0,
+      true, // Force JSON
+    );
+  }
+
   /**
    * Get available models for a specific provider
    * @param provider - Provider name (e.g., "openai", "anthropic", "perplexity")
@@ -1734,74 +2195,585 @@ export class EkoDBClient {
   }
 
   // ========================================================================
-  // COLLECTION UTILITIES
+  // GOAL API
   // ========================================================================
 
-  /**
-   * Check if a collection exists
-   * @param collection - Collection name to check
-   * @returns true if the collection exists, false otherwise
-   */
-  async collectionExists(collection: string): Promise<boolean> {
-    try {
-      const collections = await this.listCollections();
-      return collections.includes(collection);
-    } catch {
-      return false;
-    }
+  /** Create a new goal */
+  async goalCreate(data: Record): Promise<Record> {
+    return this.makeRequest<Record>("POST", "/api/chat/goals", data, 0, true);
   }
 
-  /**
-   * Count documents in a collection
-   * @param collection - Collection name
-   * @returns Number of documents in the collection
-   */
-  async countDocuments(collection: string): Promise<number> {
-    const query = new QueryBuilder().limit(100000).build();
-    const records = await this.find(collection, query);
-    return records.length;
+  /** List all goals */
+  async goalList(): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      "/api/chat/goals",
+      undefined,
+      0,
+      true,
+    );
   }
 
-  /**
-   * Create a WebSocket client
-   */
-  websocket(wsURL: string): WebSocketClient {
-    return new WebSocketClient(wsURL, this.token!);
+  /** Get a goal by ID */
+  async goalGet(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/chat/goals/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
   }
 
-  // ========== RAG Helper Methods ==========
+  /** Update a goal by ID */
+  async goalUpdate(id: string, data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "PUT",
+      `/api/chat/goals/${encodeURIComponent(id)}`,
+      data,
+      0,
+      true,
+    );
+  }
 
-  /**
-   * Generate embeddings for a single text
-   *
-   * @param text - The text to generate embeddings for
-   * @param model - The embedding model to use (e.g., "text-embedding-3-small")
-   * @returns Array of floats representing the embedding vector
-   *
-   * @example
-   * ```typescript
-   * const embedding = await client.embed(
-   *   "Hello world",
-   *   "text-embedding-3-small"
-   * );
-   * console.log(`Generated ${embedding.length} dimensions`);
-   * ```
-   */
-  async embed(text: string, model: string): Promise<number[]> {
-    const response = await this.embedRequest({ text, model });
-    if (response.embeddings.length === 0) {
-      throw new Error("No embedding returned");
-    }
-    return response.embeddings[0];
+  /** Delete a goal by ID */
+  async goalDelete(id: string): Promise<void> {
+    await this.makeRequest<void>(
+      "DELETE",
+      `/api/chat/goals/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
   }
 
-  /**
-   * Generate embeddings for multiple texts in a single batch request
-   *
-   * @param texts - Array of texts to generate embeddings for
-   * @param model - The embedding model to use
-   * @returns Array of embedding vectors
-   */
+  // ── Goal Templates ──
+
+  /** Create a new goal template */
+  async goalTemplateCreate(data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      "/api/chat/goal-templates",
+      data,
+      0,
+      true,
+    );
+  }
+
+  /** List all goal templates */
+  async goalTemplateList(): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      "/api/chat/goal-templates",
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Get a goal template by ID */
+  async goalTemplateGet(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/chat/goal-templates/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Update a goal template by ID */
+  async goalTemplateUpdate(id: string, data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "PUT",
+      `/api/chat/goal-templates/${encodeURIComponent(id)}`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  /** Delete a goal template by ID */
+  async goalTemplateDelete(id: string): Promise<void> {
+    await this.makeRequest<void>(
+      "DELETE",
+      `/api/chat/goal-templates/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Search goals */
+  async goalSearch(query: string): Promise<Record> {
+    const params = new URLSearchParams({ q: query });
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/chat/goals/search?${params}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Mark a goal as complete (status -> pending_review) */
+  async goalComplete(id: string, data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/goals/${encodeURIComponent(id)}/complete`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  /** Approve a goal (status -> in_progress) */
+  async goalApprove(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/goals/${encodeURIComponent(id)}/approve`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Reject a goal (status -> failed) */
+  async goalReject(id: string, data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/goals/${encodeURIComponent(id)}/reject`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  /** Start a goal step (status -> in_progress) */
+  async goalStepStart(id: string, stepIndex: number): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/goals/${encodeURIComponent(id)}/steps/${stepIndex}/start`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Complete a goal step with result */
+  async goalStepComplete(
+    id: string,
+    stepIndex: number,
+    data: Record,
+  ): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/goals/${encodeURIComponent(id)}/steps/${stepIndex}/complete`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  /** Fail a goal step with error */
+  async goalStepFail(
+    id: string,
+    stepIndex: number,
+    data: Record,
+  ): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/goals/${encodeURIComponent(id)}/steps/${stepIndex}/fail`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  // ========================================================================
+  // TASK API
+  // ========================================================================
+
+  /** Create a new scheduled task */
+  async taskCreate(data: Record): Promise<Record> {
+    return this.makeRequest<Record>("POST", "/api/chat/tasks", data, 0, true);
+  }
+
+  /** List all scheduled tasks */
+  async taskList(): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      "/api/chat/tasks",
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Get a task by ID */
+  async taskGet(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/chat/tasks/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Update a task by ID */
+  async taskUpdate(id: string, data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "PUT",
+      `/api/chat/tasks/${encodeURIComponent(id)}`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  /** Delete a task by ID */
+  async taskDelete(id: string): Promise<void> {
+    await this.makeRequest<void>(
+      "DELETE",
+      `/api/chat/tasks/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Get tasks that are due at the given time */
+  async taskDue(now: string): Promise<Record> {
+    const params = new URLSearchParams({ now });
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/chat/tasks/due?${params}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Start a task (status -> running) */
+  async taskStart(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/tasks/${encodeURIComponent(id)}/start`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Mark a task as succeeded */
+  async taskSucceed(id: string, data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/tasks/${encodeURIComponent(id)}/succeed`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  /** Mark a task as failed */
+  async taskFail(id: string, data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/tasks/${encodeURIComponent(id)}/fail`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  /** Pause a task */
+  async taskPause(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/tasks/${encodeURIComponent(id)}/pause`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Resume a paused task */
+  async taskResume(id: string, data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/chat/tasks/${encodeURIComponent(id)}/resume`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  // ========================================================================
+  // AGENT API
+  // ========================================================================
+
+  /** Create a new agent */
+  async agentCreate(data: Record): Promise<Record> {
+    return this.makeRequest<Record>("POST", "/api/chat/agents", data, 0, true);
+  }
+
+  /** List all agents */
+  async agentList(): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      "/api/chat/agents",
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Get an agent by ID */
+  async agentGet(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/chat/agents/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Get an agent by name */
+  async agentGetByName(name: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/chat/agents/by-name/${encodeURIComponent(name)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Update an agent by ID */
+  async agentUpdate(id: string, data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "PUT",
+      `/api/chat/agents/${encodeURIComponent(id)}`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  /** Delete an agent by ID */
+  async agentDelete(id: string): Promise<void> {
+    await this.makeRequest<void>(
+      "DELETE",
+      `/api/chat/agents/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Get agents by deployment ID */
+  async agentsByDeployment(deploymentId: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/chat/agents/by-deployment/${encodeURIComponent(deploymentId)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  // ========================================================================
+  // KV DOCUMENT LINKING
+  // ========================================================================
+
+  /** Get documents linked to a KV key */
+  async kvGetLinks(key: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/kv/links/${encodeURIComponent(key)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Link a document to a KV key */
+  async kvLink(
+    key: string,
+    collection: string,
+    documentId: string,
+  ): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/kv/link`,
+      { key, collection, document_id: documentId },
+      0,
+      true,
+    );
+  }
+
+  /** Unlink a document from a KV key */
+  async kvUnlink(
+    key: string,
+    collection: string,
+    documentId: string,
+  ): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/kv/unlink`,
+      { key, collection, document_id: documentId },
+      0,
+      true,
+    );
+  }
+
+  // ========================================================================
+  // SCHEDULE MANAGEMENT
+  // ========================================================================
+
+  /** Create a new schedule */
+  async createSchedule(data: Record): Promise<Record> {
+    return this.makeRequest<Record>("POST", `/api/schedules`, data, 0, true);
+  }
+
+  /** List all schedules */
+  async listSchedules(): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/schedules`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Get a schedule by ID */
+  async getSchedule(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/schedules/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Update a schedule */
+  async updateSchedule(id: string, data: Record): Promise<Record> {
+    return this.makeRequest<Record>(
+      "PUT",
+      `/api/schedules/${encodeURIComponent(id)}`,
+      data,
+      0,
+      true,
+    );
+  }
+
+  /** Delete a schedule */
+  async deleteSchedule(id: string): Promise<void> {
+    await this.makeRequest<void>(
+      "DELETE",
+      `/api/schedules/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Pause a schedule */
+  async pauseSchedule(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/schedules/${encodeURIComponent(id)}/pause`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Resume a schedule */
+  async resumeSchedule(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "POST",
+      `/api/schedules/${encodeURIComponent(id)}/resume`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  // ========================================================================
+  // COLLECTION UTILITIES
+  // ========================================================================
+
+  /**
+   * Check if a collection exists
+   * @param collection - Collection name to check
+   * @returns true if the collection exists, false otherwise
+   */
+  async collectionExists(collection: string): Promise<boolean> {
+    try {
+      const collections = await this.listCollections();
+      return collections.includes(collection);
+    } catch {
+      return false;
+    }
+  }
+
+  /**
+   * Count documents in a collection
+   * @param collection - Collection name
+   * @returns Number of documents in the collection
+   */
+  async countDocuments(collection: string): Promise<number> {
+    const query = new QueryBuilder().limit(100000).build();
+    const records = await this.find(collection, query);
+    return records.length;
+  }
+
+  /**
+   * Create a WebSocket client
+   */
+  websocket(wsURL: string): WebSocketClient {
+    return new WebSocketClient(wsURL, this.token!);
+  }
+
+  // ========== RAG Helper Methods ==========
+
+  /**
+   * Generate embeddings for a single text
+   *
+   * @param text - The text to generate embeddings for
+   * @param model - The embedding model to use (e.g., "text-embedding-3-small")
+   * @returns Array of floats representing the embedding vector
+   *
+   * @example
+   * ```typescript
+   * const embedding = await client.embed(
+   *   "Hello world",
+   *   "text-embedding-3-small"
+   * );
+   * console.log(`Generated ${embedding.length} dimensions`);
+   * ```
+   */
+  async embed(text: string, model: string): Promise<number[]> {
+    const response = await this.embedRequest({ text, model });
+    if (response.embeddings.length === 0) {
+      throw new Error("No embedding returned");
+    }
+    return response.embeddings[0];
+  }
+
+  /**
+   * Generate embeddings for multiple texts in a single batch request
+   *
+   * @param texts - Array of texts to generate embeddings for
+   * @param model - The embedding model to use
+   * @returns Array of embedding vectors
+   */
   async embedBatch(texts: string[], model: string): Promise<number[][]> {
     const response = await this.embedRequest({ texts, model });
     return response.embeddings;
@@ -1916,26 +2888,132 @@ export class EkoDBClient {
   }
 }
 
+/** Mutation notification from a subscription. */
+export interface MutationNotification {
+  collection: string;
+  event: string;
+  recordIds: string[];
+  records?: any;
+  timestamp: string;
+}
+
+/** A chunk/event from a streaming chat response. */
+export type ChatStreamEvent =
+  | { type: "chunk"; content: string }
+  | {
+      type: "end";
+      messageId: string;
+      tokenUsage?: any;
+      toolCallHistory?: any;
+      executionTimeMs: number;
+      /** Model's context window size in tokens. */
+      contextWindow?: number;
+    }
+  | {
+      type: "toolCall";
+      chatId: string;
+      callId: string;
+      toolName: string;
+      arguments: any;
+    }
+  | { type: "error"; error: string };
+
+/** Definition for a client-side tool the LLM can call. */
+export interface ClientToolDefinition {
+  name: string;
+  description: string;
+  parameters: any;
+}
+
+/** Options for chatSend. */
+export interface ChatSendOptions {
+  bypassRipple?: boolean;
+  clientTools?: ClientToolDefinition[];
+  maxIterations?: number;
+  confirmTools?: string[];
+  excludeTools?: string[];
+}
+
+/** Options for subscribe. */
+export interface SubscribeOptions {
+  filterField?: string;
+  filterValue?: string;
+}
+
+/** EventEmitter-like interface for subscriptions and chat streams. */
+export class EventStream<_T = unknown> {
+  private listeners: Map<string, Array<(data: any) => void>> = new Map();
+  private _closed = false;
+
+  on(event: string, listener: (data: any) => void): this {
+    if (!this.listeners.has(event)) {
+      this.listeners.set(event, []);
+    }
+    this.listeners.get(event)!.push(listener);
+    return this;
+  }
+
+  /** @internal */
+  emit(event: string, data?: any): void {
+    const handlers = this.listeners.get(event);
+    if (handlers) {
+      for (const handler of handlers) {
+        handler(data);
+      }
+    }
+  }
+
+  get closed(): boolean {
+    return this._closed;
+  }
+
+  /** @internal */
+  close(): void {
+    this._closed = true;
+    this.emit("close");
+  }
+}
+
 /**
- * WebSocket client for real-time queries
+ * WebSocket client for real-time queries, subscriptions, and chat streaming.
  */
 export class WebSocketClient {
   private wsURL: string;
   private token: string;
   private ws: any = null;
+  private dispatcherRunning = false;
+
+  // Dispatcher state
+  private pendingRequests: Map<
+    string,
+    { resolve: (value: any) => void; reject: (reason: any) => void }
+  > = new Map();
+  private subscriptions: Map<string, EventStream<MutationNotification>> =
+    new Map();
+  private chatStreams: Map<string, EventStream<ChatStreamEvent>> = new Map();
+  private registerToolsAck: {
+    resolve: (value: any) => void;
+    reject: (reason: any) => void;
+  } | null = null;
 
   constructor(wsURL: string, token: string) {
     this.wsURL = wsURL;
     this.token = token;
   }
 
+  private messageCounter = 0;
+
+  private genMessageId(): string {
+    const counter = this.messageCounter++;
+    return `${Date.now()}-${counter}-${Math.random().toString(36).slice(2, 8)}`;
+  }
+
   /**
-   * Connect to WebSocket
+   * Connect and start the dispatcher.
    */
-  private async connect(): Promise<void> {
-    if (this.ws) return;
+  private async ensureConnected(): Promise<void> {
+    if (this.ws && this.dispatcherRunning) return;
 
-    // Dynamic import for Node.js WebSocket
     const WebSocket = (await import("ws")).default;
 
     let url = this.wsURL;
@@ -1949,49 +3027,369 @@ export class WebSocketClient {
       },
     });
 
-    return new Promise((resolve, reject) => {
+    await new Promise<void>((resolve, reject) => {
       this.ws.on("open", () => resolve());
       this.ws.on("error", (err: Error) => reject(err));
     });
+
+    this.spawnDispatcher();
+  }
+
+  private spawnDispatcher(): void {
+    if (this.dispatcherRunning) return;
+    this.dispatcherRunning = true;
+
+    this.ws.on("message", (data: Buffer) => {
+      try {
+        const msg = JSON.parse(data.toString());
+        this.routeMessage(msg);
+      } catch {
+        // Ignore malformed messages
+      }
+    });
+
+    this.ws.on("close", () => {
+      this.dispatcherRunning = false;
+      // Notify all pending requests
+      for (const [, pending] of this.pendingRequests) {
+        pending.reject(new Error("WebSocket connection closed"));
+      }
+      this.pendingRequests.clear();
+      // Close all chat streams
+      for (const [, stream] of this.chatStreams) {
+        stream.emit("event", { type: "error", error: "Connection closed" });
+        stream.close();
+      }
+      this.chatStreams.clear();
+      // Close all subscriptions
+      for (const [, stream] of this.subscriptions) {
+        stream.close();
+      }
+      this.subscriptions.clear();
+      this.ws = null;
+    });
+  }
+
+  private routeMessage(msg: any): void {
+    switch (msg.type) {
+      case "Success":
+      case "Error": {
+        // Try messageId from top-level, then from payload
+        const messageId =
+          msg.messageId ||
+          msg.message_id ||
+          msg.payload?.message_id ||
+          msg.payload?.messageId;
+        let matched = false;
+        if (messageId && this.pendingRequests.has(messageId)) {
+          const pending = this.pendingRequests.get(messageId)!;
+          this.pendingRequests.delete(messageId);
+          if (msg.type === "Error") {
+            pending.reject(new Error(msg.message || "Unknown error"));
+          } else {
+            pending.resolve(msg.payload);
+          }
+          matched = true;
+        }
+        if (!matched && this.registerToolsAck) {
+          const ack = this.registerToolsAck;
+          this.registerToolsAck = null;
+          if (msg.type === "Error") {
+            ack.reject(new Error(msg.message || "Tool registration failed"));
+          } else {
+            ack.resolve(msg.payload);
+          }
+          matched = true;
+        }
+        // Server doesn't echo messageId — if there's exactly one pending
+        // request, deliver the response to it (sequential request/response).
+        if (!matched && this.pendingRequests.size === 1) {
+          const entry = this.pendingRequests.entries().next().value!;
+          const key = entry[0];
+          const pending = entry[1];
+          this.pendingRequests.delete(key);
+          if (msg.type === "Error") {
+            pending.reject(new Error(msg.message || "Unknown error"));
+          } else {
+            pending.resolve(msg.payload);
+          }
+        }
+        break;
+      }
+
+      case "MutationNotification": {
+        const payload = msg.payload;
+        const notification: MutationNotification = {
+          collection: payload.collection,
+          event: payload.event,
+          recordIds: payload.record_ids || payload.recordIds || [],
+          records: payload.records,
+          timestamp: payload.timestamp,
+        };
+        for (const [collection, stream] of this.subscriptions) {
+          if (collection === notification.collection) {
+            stream.emit("mutation", notification);
+          }
+        }
+        break;
+      }
+
+      case "ChatStreamChunk": {
+        const chatId = msg.payload?.chat_id || msg.payload?.chatId;
+        const stream = this.chatStreams.get(chatId);
+        if (stream) {
+          stream.emit("event", {
+            type: "chunk",
+            content: msg.payload.content,
+          } as ChatStreamEvent);
+        }
+        break;
+      }
+
+      case "ChatStreamEnd": {
+        const chatId = msg.payload?.chat_id || msg.payload?.chatId;
+        const stream = this.chatStreams.get(chatId);
+        if (stream) {
+          stream.emit("event", {
+            type: "end",
+            messageId: msg.payload.message_id || msg.payload.messageId || "",
+            tokenUsage: msg.payload.token_usage || msg.payload.tokenUsage,
+            toolCallHistory:
+              msg.payload.tool_call_history || msg.payload.toolCallHistory,
+            executionTimeMs:
+              msg.payload.execution_time_ms || msg.payload.executionTimeMs || 0,
+            contextWindow:
+              msg.payload.context_window || msg.payload.contextWindow,
+          } as ChatStreamEvent);
+          this.chatStreams.delete(chatId);
+          stream.close();
+        }
+        break;
+      }
+
+      case "ChatStreamError": {
+        const chatId = msg.payload?.chat_id || msg.payload?.chatId;
+        const stream = this.chatStreams.get(chatId);
+        if (stream) {
+          stream.emit("event", {
+            type: "error",
+            error: msg.payload.error || msg.payload.message || "Unknown error",
+          } as ChatStreamEvent);
+          this.chatStreams.delete(chatId);
+          stream.close();
+        }
+        break;
+      }
+
+      case "ClientToolCall": {
+        const chatId = msg.payload?.chat_id || msg.payload?.chatId;
+        const stream = this.chatStreams.get(chatId);
+        if (stream) {
+          stream.emit("event", {
+            type: "toolCall",
+            chatId,
+            callId: msg.payload.call_id || msg.payload.callId,
+            toolName: msg.payload.tool_name || msg.payload.toolName,
+            arguments: msg.payload.arguments,
+          } as ChatStreamEvent);
+        }
+        break;
+      }
+    }
+  }
+
+  private async sendRequest(request: any): Promise<any> {
+    await this.ensureConnected();
+    const messageId = request.messageId || request.message_id;
+
+    return new Promise((resolve, reject) => {
+      this.pendingRequests.set(messageId, { resolve, reject });
+      try {
+        this.ws.send(JSON.stringify(request));
+      } catch (err) {
+        this.pendingRequests.delete(messageId);
+        reject(err);
+      }
+    });
   }
 
   /**
-   * Find all records in a collection via WebSocket
+   * Find all records in a collection via WebSocket.
    */
   async findAll(collection: string): Promise<Record[]> {
-    await this.connect();
-
-    const messageId = Date.now().toString();
-    const request = {
+    const messageId = this.genMessageId();
+    const payload = await this.sendRequest({
       type: "FindAll",
       messageId,
       payload: { collection },
+    });
+    return payload?.data || [];
+  }
+
+  /**
+   * Subscribe to mutation notifications on a collection.
+   * Returns an EventStream that emits "mutation" events.
+   */
+  async subscribe(
+    collection: string,
+    options?: SubscribeOptions,
+  ): Promise<EventStream<MutationNotification>> {
+    await this.ensureConnected();
+
+    if (this.subscriptions.has(collection)) {
+      throw new Error(`Already subscribed to collection "${collection}"`);
+    }
+
+    const messageId = this.genMessageId();
+    const stream = new EventStream<MutationNotification>();
+    this.subscriptions.set(collection, stream);
+
+    const request: any = {
+      type: "Subscribe",
+      messageId,
+      payload: {
+        collection,
+        ...(options?.filterField && { filter_field: options.filterField }),
+        ...(options?.filterValue && { filter_value: options.filterValue }),
+      },
     };
 
-    return new Promise((resolve, reject) => {
+    // Send subscribe request and wait for ack
+    try {
+      await this.sendRequest(request);
+    } catch (err) {
+      this.subscriptions.delete(collection);
+      throw err;
+    }
+    return stream;
+  }
+
+  /**
+   * Send a chat message and receive a streaming response.
+   * Returns an EventStream that emits "event" with ChatStreamEvent objects.
+   */
+  async chatSend(
+    chatId: string,
+    message: string,
+    options?: ChatSendOptions,
+  ): Promise<EventStream<ChatStreamEvent>> {
+    await this.ensureConnected();
+
+    if (this.chatStreams.has(chatId)) {
+      throw new Error(`Chat stream already active for chatId "${chatId}"`);
+    }
+
+    const stream = new EventStream<ChatStreamEvent>();
+    this.chatStreams.set(chatId, stream);
+
+    const request: any = {
+      type: "ChatSend",
+      payload: {
+        chat_id: chatId,
+        message,
+        ...(options?.bypassRipple != null && {
+          bypass_ripple: options.bypassRipple,
+        }),
+        ...(options?.clientTools && { client_tools: options.clientTools }),
+        ...(options?.maxIterations != null && {
+          max_iterations: options.maxIterations,
+        }),
+        ...(options?.confirmTools && { confirm_tools: options.confirmTools }),
+        ...(options?.excludeTools && { exclude_tools: options.excludeTools }),
+      },
+    };
+
+    this.ws.send(JSON.stringify(request));
+    return stream;
+  }
+
+  /**
+   * Register client-side tools for a chat session.
+   */
+  async registerClientTools(
+    chatId: string,
+    tools: ClientToolDefinition[],
+  ): Promise<void> {
+    await this.ensureConnected();
+
+    const request = {
+      type: "RegisterClientTools",
+      payload: {
+        chat_id: chatId,
+        tools,
+      },
+    };
+
+    await new Promise<void>((resolve, reject) => {
+      this.registerToolsAck = {
+        resolve: () => resolve(),
+        reject: (err) => reject(err),
+      };
       this.ws.send(JSON.stringify(request));
+    });
+  }
 
-      this.ws.once("message", (data: Buffer) => {
-        const response = JSON.parse(data.toString());
+  /**
+   * Send a tool result back to the server during a chat stream.
+   */
+  async sendToolResult(
+    chatId: string,
+    callId: string,
+    success: boolean,
+    result?: any,
+    error?: string,
+  ): Promise<void> {
+    await this.ensureConnected();
 
-        if (response.type === "Error") {
-          reject(new Error(response.message));
-        } else {
-          resolve(response.payload?.data || []);
-        }
-      });
+    const request = {
+      type: "ClientToolResult",
+      payload: {
+        chat_id: chatId,
+        call_id: callId,
+        success,
+        ...(result !== undefined && { result }),
+        ...(error !== undefined && { error }),
+      },
+    };
+
+    this.ws.send(JSON.stringify(request));
+  }
 
-      this.ws.once("error", reject);
+  /**
+   * Stateless raw LLM completion via WebSocket.
+   *
+   * Sends a RawComplete message and waits for the Success response.
+   * Preferred over HTTP for deployed instances: the persistent WSS
+   * connection is already authenticated and won't be killed by reverse
+   * proxy timeouts.
+   */
+  async rawCompletion(
+    request: RawCompletionRequest,
+  ): Promise<RawCompletionResponse> {
+    await this.ensureConnected();
+    const messageId = this.genMessageId();
+    const payload = await this.sendRequest({
+      type: "RawComplete",
+      messageId,
+      payload: {
+        system_prompt: request.system_prompt,
+        message: request.message,
+        ...(request.provider && { provider: request.provider }),
+        ...(request.model && { model: request.model }),
+        ...(request.max_tokens != null && { max_tokens: request.max_tokens }),
+      },
     });
+    return { content: payload?.data?.content || "" };
   }
 
   /**
-   * Close the WebSocket connection
+   * Close the WebSocket connection.
    */
   close(): void {
     if (this.ws) {
       this.ws.close();
       this.ws = null;
+      this.dispatcherRunning = false;
     }
   }
 }