@ekodb/ekodb-client 0.19.0 → 0.21.0

package/src/client.ts

@@ -3,7 +3,7 @@
  */
 
 import { encode, decode } from "@msgpack/msgpack";
-import { QueryBuilder, Query as QueryBuilderQuery } from "./query-builder";
+import { QueryBuilder, Query } from "./query-builder";
 import { SearchQuery, SearchResponse } from "./search";
 import { Schema, SchemaBuilder, CollectionMetadata } from "./schema";
 import { UserFunction, FunctionResult } from "./functions";
@@ -65,11 +65,11 @@ export class RateLimitError extends Error {
   }
 }
 
-export interface Query {
-  limit?: number;
-  offset?: number;
-  filter?: Record;
-}
+// `Query` is the canonical find/query body shape — the server's FindBody —
+// re-exported from the query builder so the whole client shares a single type
+// (`filter`, `sort`, `limit`, `skip`, `join`, `select_fields`, `exclude_fields`,
+// matching the server exactly). `QueryBuilder.build()` returns this same type.
+export type { Query };
 
 export interface BatchOperationResult {
   successful: string[];
@@ -115,6 +115,24 @@ export interface FindOptions {
   bypassCache?: boolean;
   selectFields?: string[];
   excludeFields?: string[];
+  /**
+   * Read within a transaction (read-your-writes). When set, the read is served
+   * from the transaction's own view — its uncommitted staged writes, else the
+   * committed store — and recorded in the transaction's read set for
+   * commit-time conflict detection. Omit for an ordinary committed read.
+   */
+  transactionId?: string;
+}
+
+/**
+ * Options for a point read by id. `transactionId` enables read-your-writes
+ * within a transaction (see {@link FindOptions.transactionId}).
+ */
+export interface FindByIdOptions {
+  selectFields?: string[];
+  excludeFields?: string[];
+  bypassRipple?: boolean;
+  transactionId?: string;
 }
 
 export interface BatchInsertOptions {
@@ -379,6 +397,20 @@ export interface FunctionStageConfig {
   [key: string]: any;
 }
 
+/**
+ * Strip trailing slashes from a base URL so path concatenation
+ * (`${base}/api/...`) never yields a double-slash path. Uses a linear scan
+ * rather than a regex like `/\/+$/`, which CodeQL flags as polynomial-time
+ * backtracking on caller-supplied input.
+ */
+function stripTrailingSlashes(url: string): string {
+  let end = url.length;
+  while (end > 0 && url.charCodeAt(end - 1) === 47 /* "/" */) {
+    end--;
+  }
+  return end === url.length ? url : url.slice(0, end);
+}
+
 export class EkoDBClient {
   private baseURL: string;
   private apiKey: string;
@@ -392,13 +424,15 @@ export class EkoDBClient {
   constructor(config: string | ClientConfig, apiKey?: string) {
     // Support both old (baseURL, apiKey) and new (config object) signatures
     if (typeof config === "string") {
-      this.baseURL = config;
+      // Strip trailing slashes so `${baseURL}/api/...` never produces a
+      // double-slash path (some servers/proxies reject `//api/...`).
+      this.baseURL = stripTrailingSlashes(config);
       this.apiKey = apiKey!;
       this.shouldRetry = true;
       this.maxRetries = 3;
       this.format = SerializationFormat.MessagePack; // Default to MessagePack for 2-3x performance
     } else {
-      this.baseURL = config.baseURL;
+      this.baseURL = stripTrailingSlashes(config.baseURL);
       this.apiKey = config.apiKey;
       this.shouldRetry = config.shouldRetry ?? true;
       this.maxRetries = config.maxRetries ?? 3;
@@ -555,6 +589,41 @@ export class EkoDBClient {
     return new Promise((resolve) => setTimeout(resolve, seconds * 1000));
   }
 
+  /**
+   * Parse a `Retry-After` header into a non-negative delay in seconds.
+   *
+   * Per RFC 9110 the value is either delay-seconds (an integer) or an
+   * HTTP-date. Anything that doesn't resolve to a finite, non-negative number
+   * (missing header, garbage, a past date) falls back to `defaultSecs`.
+   */
+  private parseRetryAfter(header: string | null, defaultSecs = 60): number {
+    if (!header) return defaultSecs;
+
+    // delay-seconds form: a bare integer.
+    const secs = Number(header.trim());
+    if (Number.isFinite(secs)) return Math.max(0, secs);
+
+    // HTTP-date form: compute the delay from now.
+    const dateMs = Date.parse(header);
+    if (Number.isFinite(dateMs)) {
+      return Math.max(0, (dateMs - Date.now()) / 1000);
+    }
+
+    return defaultSecs;
+  }
+
+  /**
+   * Backoff delay (in seconds) for a 0-indexed retry attempt: a capped
+   * exponential schedule (0.2s → 5s) with full jitter, so concurrent clients
+   * don't retry in lockstep. Returns a value in [d/2, d].
+   */
+  private backoffSeconds(attempt: number): number {
+    const base = 0.2;
+    const max = 5;
+    const d = Math.min(base * Math.pow(2, Math.max(0, attempt)), max);
+    return d / 2 + Math.random() * (d / 2);
+  }
+
   /**
    * Helper to determine if a path should use JSON
    * Only CRUD operations (insert/update/delete/batch) use MessagePack
@@ -640,14 +709,16 @@ export class EkoDBClient {
 
       // Handle rate limiting (429)
       if (response.status === 429) {
-        const retryAfter = parseInt(
-          response.headers.get("retry-after") || "60",
-          10,
+        const retryAfter = this.parseRetryAfter(
+          response.headers.get("retry-after"),
         );
 
         if (this.shouldRetry && attempt < this.maxRetries) {
-          console.log(`Rate limited. Retrying after ${retryAfter} seconds...`);
-          await this.sleep(retryAfter);
+          // Honor the server's Retry-After, but cap it so a hostile/large value
+          // can't pin the client for minutes.
+          const wait = Math.min(retryAfter, 60);
+          console.log(`Rate limited. Retrying after ${wait} seconds...`);
+          await this.sleep(wait);
           return this.makeRequest<T>(
             method,
             path,
@@ -674,9 +745,9 @@ export class EkoDBClient {
         this.shouldRetry &&
         attempt < this.maxRetries
       ) {
-        const retryDelay = 10;
+        const retryDelay = this.backoffSeconds(attempt);
         console.log(
-          `Service unavailable. Retrying after ${retryDelay} seconds...`,
+          `Service unavailable. Retrying after ${retryDelay.toFixed(2)}s...`,
         );
         await this.sleep(retryDelay);
         return this.makeRequest<T>(method, path, data, attempt + 1, forceJson);
@@ -692,8 +763,10 @@ export class EkoDBClient {
         this.shouldRetry &&
         attempt < this.maxRetries
       ) {
-        const retryDelay = 3;
-        console.log(`Network error. Retrying after ${retryDelay} seconds...`);
+        const retryDelay = this.backoffSeconds(attempt);
+        console.log(
+          `Network error. Retrying after ${retryDelay.toFixed(2)}s...`,
+        );
         await this.sleep(retryDelay);
         return this.makeRequest<T>(method, path, data, attempt + 1, forceJson);
       }
@@ -727,8 +800,8 @@ export class EkoDBClient {
     }
 
     const url = params.toString()
-      ? `/api/insert/${collection}?${params.toString()}`
-      : `/api/insert/${collection}`;
+      ? `/api/insert/${encodeURIComponent(collection)}?${params.toString()}`
+      : `/api/insert/${encodeURIComponent(collection)}`;
 
     return this.makeRequest<Record>("POST", url, data);
   }
@@ -758,20 +831,64 @@ export class EkoDBClient {
   async find(
     collection: string,
     query: Query | QueryBuilder = {},
+    options?: { bypassRipple?: boolean; transactionId?: string },
   ): Promise<Record[]> {
     const queryObj = query instanceof QueryBuilder ? query.build() : query;
-    return this.makeRequest<Record[]>(
-      "POST",
-      `/api/find/${collection}`,
-      queryObj,
-    );
+    // bypass_ripple and transaction_id are query parameters — the same way every
+    // other method (insert/update/findById) carries bypass_ripple — not part of
+    // the FindBody. Hoist any bypass_ripple carried on the query object (e.g. from
+    // QueryBuilder.bypassRipple()) out of the body so it is ALWAYS sent as a query
+    // param; an explicit options.bypassRipple wins.
+    let body: unknown = queryObj;
+    let bypassRipple = options?.bypassRipple;
+    if (body && typeof body === "object" && "bypass_ripple" in body) {
+      const { bypass_ripple, ...rest } = body as Record & {
+        bypass_ripple?: boolean;
+      };
+      body = rest;
+      if (bypassRipple === undefined) bypassRipple = bypass_ripple;
+    }
+    const params = new URLSearchParams();
+    if (options?.transactionId)
+      params.append("transaction_id", options.transactionId);
+    if (bypassRipple !== undefined)
+      params.append("bypass_ripple", String(bypassRipple));
+    const qs = params.toString();
+    const url = qs
+      ? `/api/find/${encodeURIComponent(collection)}?${qs}`
+      : `/api/find/${encodeURIComponent(collection)}`;
+    return this.makeRequest<Record[]>("POST", url, body);
   }
 
   /**
-   * Find a document by ID
+   * Find a document by ID.
+   * @param options - Optional read options. `transactionId` reads within a
+   *   transaction (read-your-writes); see {@link FindByIdOptions}.
    */
-  async findById(collection: string, id: string): Promise<Record> {
-    return this.makeRequest<Record>("GET", `/api/find/${collection}/${id}`);
+  async findById(
+    collection: string,
+    id: string,
+    options?: FindByIdOptions,
+  ): Promise<Record> {
+    const params = new URLSearchParams();
+    if (options?.selectFields?.length) {
+      params.append("select_fields", options.selectFields.join(","));
+    }
+    if (options?.excludeFields?.length) {
+      params.append("exclude_fields", options.excludeFields.join(","));
+    }
+    // bypass_ripple is a GET query param, the same way the non-transactional
+    // findById carries it; it rides alongside transaction_id when both are set.
+    if (options?.bypassRipple !== undefined) {
+      params.append("bypass_ripple", String(options.bypassRipple));
+    }
+    if (options?.transactionId) {
+      params.append("transaction_id", options.transactionId);
+    }
+    const url = params.toString()
+      ? `/api/find/${encodeURIComponent(collection)}/${encodeURIComponent(id)}?${params.toString()}`
+      : `/api/find/${encodeURIComponent(collection)}/${encodeURIComponent(id)}`;
+    return this.makeRequest<Record>("GET", url);
   }
 
   /**
@@ -780,12 +897,14 @@ export class EkoDBClient {
    * @param id - Document ID
    * @param selectFields - Fields to include in the result
    * @param excludeFields - Fields to exclude from the result
+   * @param transactionId - Read within a transaction (read-your-writes)
    */
   async findByIdWithProjection(
     collection: string,
     id: string,
     selectFields?: string[],
     excludeFields?: string[],
+    transactionId?: string,
   ): Promise<Record> {
     const params = new URLSearchParams();
     if (selectFields?.length) {
@@ -794,9 +913,12 @@ export class EkoDBClient {
     if (excludeFields?.length) {
       params.append("exclude_fields", excludeFields.join(","));
     }
+    if (transactionId) {
+      params.append("transaction_id", transactionId);
+    }
     const url = params.toString()
-      ? `/api/find/${collection}/${id}?${params.toString()}`
-      : `/api/find/${collection}/${id}`;
+      ? `/api/find/${encodeURIComponent(collection)}/${encodeURIComponent(id)}?${params.toString()}`
+      : `/api/find/${encodeURIComponent(collection)}/${encodeURIComponent(id)}`;
     return this.makeRequest<Record>("GET", url);
   }
 
@@ -822,8 +944,8 @@ export class EkoDBClient {
     }
 
     const url = params.toString()
-      ? `/api/update/${collection}/${id}?${params.toString()}`
-      : `/api/update/${collection}/${id}`;
+      ? `/api/update/${encodeURIComponent(collection)}/${encodeURIComponent(id)}?${params.toString()}`
+      : `/api/update/${encodeURIComponent(collection)}/${encodeURIComponent(id)}`;
 
     return this.makeRequest<Record>("PUT", url, record);
   }
@@ -848,7 +970,7 @@ export class EkoDBClient {
     field: string,
     value?: any,
   ): Promise<Record> {
-    const url = `/api/update/${collection}/${id}/action/${action}`;
+    const url = `/api/update/${encodeURIComponent(collection)}/${encodeURIComponent(id)}/action/${encodeURIComponent(action)}`;
     return this.makeRequest<Record>("PUT", url, {
       field,
       value: value ?? null,
@@ -870,7 +992,7 @@ export class EkoDBClient {
     id: string,
     actions: [string, string, any][],
   ): Promise<Record> {
-    const url = `/api/update/sequence/${collection}/${id}`;
+    const url = `/api/update/sequence/${encodeURIComponent(collection)}/${encodeURIComponent(id)}`;
     return this.makeRequest<Record>("PUT", url, actions);
   }
 
@@ -894,8 +1016,8 @@ export class EkoDBClient {
     }
 
     const url = params.toString()
-      ? `/api/delete/${collection}/${id}?${params.toString()}`
-      : `/api/delete/${collection}/${id}`;
+      ? `/api/delete/${encodeURIComponent(collection)}/${encodeURIComponent(id)}?${params.toString()}`
+      : `/api/delete/${encodeURIComponent(collection)}/${encodeURIComponent(id)}`;
 
     await this.makeRequest<void>("DELETE", url);
   }
@@ -921,8 +1043,8 @@ export class EkoDBClient {
 
     const inserts = records.map((data) => ({ data }));
     const url = params.toString()
-      ? `/api/batch/insert/${collection}?${params.toString()}`
-      : `/api/batch/insert/${collection}`;
+      ? `/api/batch/insert/${encodeURIComponent(collection)}?${params.toString()}`
+      : `/api/batch/insert/${encodeURIComponent(collection)}`;
 
     return this.makeRequest<BatchOperationResult>("POST", url, { inserts });
   }
@@ -941,7 +1063,7 @@ export class EkoDBClient {
     }));
     return this.makeRequest<BatchOperationResult>(
       "PUT",
-      `/api/batch/update/${collection}`,
+      `/api/batch/update/${encodeURIComponent(collection)}`,
       { updates: formattedUpdates },
     );
   }
@@ -960,7 +1082,7 @@ export class EkoDBClient {
     }));
     return this.makeRequest<BatchOperationResult>(
       "DELETE",
-      `/api/batch/delete/${collection}`,
+      `/api/batch/delete/${encodeURIComponent(collection)}`,
       { deletes },
     );
   }
@@ -1012,6 +1134,19 @@ export class EkoDBClient {
     );
   }
 
+  /**
+   * Clear the entire KV store (all keys in the namespace).
+   */
+  async kvClear(): Promise<void> {
+    await this.makeRequest<void>(
+      "DELETE",
+      "/api/kv/clear",
+      undefined,
+      0,
+      true, // Force JSON for KV operations
+    );
+  }
+
   /**
    * Batch get multiple keys
    * @param keys - Array of keys to retrieve
@@ -1116,7 +1251,18 @@ export class EkoDBClient {
   // ============================================================================
 
   /**
-   * Begin a new transaction
+   * Begin a new transaction.
+   *
+   * Transactions are buffered: statements issued with this `transactionId`
+   * (passed via the `transactionId` option on insert/update/delete/find/…) are
+   * staged and applied atomically only at {@link commitTransaction}. They are
+   * invisible to everyone else until commit, and visible to this transaction's
+   * own reads (read-your-writes) only when those reads also carry the
+   * `transactionId`. {@link rollbackTransaction} discards the staged writes.
+   * `commitTransaction` may reject with a conflict (HTTP 409) if a record this
+   * transaction read or wrote was changed by another committed transaction —
+   * retry the transaction in that case.
+   *
    * @param isolationLevel - Transaction isolation level (default: "ReadCommitted")
    * @returns Transaction ID
    */
@@ -1165,7 +1311,7 @@ export class EkoDBClient {
   }
 
   /**
-   * Rollback a transaction
+   * Rollback a transaction (discards all staged writes; nothing was applied).
    * @param transactionId - The transaction ID to rollback
    */
   async rollbackTransaction(transactionId: string): Promise<void> {
@@ -1178,6 +1324,49 @@ export class EkoDBClient {
     );
   }
 
+  /**
+   * Create a savepoint within a transaction. A later
+   * {@link rollbackToSavepoint} discards everything staged after it.
+   */
+  async createSavepoint(transactionId: string, name: string): Promise<void> {
+    await this.makeRequest<void>(
+      "POST",
+      `/api/transactions/${encodeURIComponent(transactionId)}/savepoints`,
+      { name },
+      0,
+      true,
+    );
+  }
+
+  /**
+   * Roll the transaction back to a savepoint, discarding writes staged after it.
+   */
+  async rollbackToSavepoint(
+    transactionId: string,
+    name: string,
+  ): Promise<void> {
+    await this.makeRequest<void>(
+      "POST",
+      `/api/transactions/${encodeURIComponent(transactionId)}/savepoints/${encodeURIComponent(name)}/rollback`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /**
+   * Release (forget) a savepoint. Staged work is unaffected.
+   */
+  async releaseSavepoint(transactionId: string, name: string): Promise<void> {
+    await this.makeRequest<void>(
+      "DELETE",
+      `/api/transactions/${encodeURIComponent(transactionId)}/savepoints/${encodeURIComponent(name)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
   // ============================================================================
   // Convenience Methods
   // ============================================================================
@@ -1324,7 +1513,7 @@ export class EkoDBClient {
   ): Promise<Record[]> {
     // Page 1 = skip 0, Page 2 = skip pageSize, etc.
     const skip = page > 0 ? (page - 1) * pageSize : 0;
-    const query: QueryBuilderQuery = {
+    const query: Query = {
       limit: pageSize,
       skip: skip,
     };
@@ -1345,13 +1534,27 @@ export class EkoDBClient {
     return result.collections;
   }
 
+  /**
+   * List collections, excluding internal chat/system collections.
+   */
+  async listUserCollections(): Promise<string[]> {
+    const result = await this.makeRequest<{ collections: string[] }>(
+      "GET",
+      "/api/collections?exclude_internal=true",
+      undefined,
+      0,
+      true, // Force JSON for metadata operations
+    );
+    return result.collections;
+  }
+
   /**
    * Delete a collection
    */
   async deleteCollection(collection: string): Promise<void> {
     await this.makeRequest<void>(
       "DELETE",
-      `/api/collections/${collection}`,
+      `/api/collections/${encodeURIComponent(collection)}`,
       undefined,
       0,
       true, // Force JSON for metadata operations
@@ -1369,7 +1572,7 @@ export class EkoDBClient {
   async restoreRecord(collection: string, id: string): Promise<boolean> {
     const result = await this.makeRequest<{ status: string }>(
       "POST",
-      `/api/trash/${collection}/${id}`,
+      `/api/trash/${encodeURIComponent(collection)}/${encodeURIComponent(id)}`,
       undefined,
       0,
       true,
@@ -1391,7 +1594,13 @@ export class EkoDBClient {
       status: string;
       collection: string;
       records_restored: number;
-    }>("POST", `/api/trash/${collection}`, undefined, 0, true);
+    }>(
+      "POST",
+      `/api/trash/${encodeURIComponent(collection)}`,
+      undefined,
+      0,
+      true,
+    );
     return { recordsRestored: result.records_restored };
   }
 
@@ -1418,7 +1627,7 @@ export class EkoDBClient {
     const schemaObj = schema instanceof SchemaBuilder ? schema.build() : schema;
     await this.makeRequest<void>(
       "POST",
-      `/api/collections/${collection}`,
+      `/api/collections/${encodeURIComponent(collection)}`,
       schemaObj,
       0,
       true, // Force JSON for metadata operations
@@ -1434,7 +1643,7 @@ export class EkoDBClient {
   async getCollection(collection: string): Promise<CollectionMetadata> {
     return this.makeRequest<CollectionMetadata>(
       "GET",
-      `/api/collections/${collection}`,
+      `/api/collections/${encodeURIComponent(collection)}`,
       undefined,
       0,
       true, // Force JSON for metadata operations
@@ -1492,7 +1701,7 @@ export class EkoDBClient {
     // Ensure all parameters from SearchQuery are sent to server
     return this.makeRequest<SearchResponse>(
       "POST",
-      `/api/search/${collection}`,
+      `/api/search/${encodeURIComponent(collection)}`,
       query,
       0,
       true, // Force JSON for search operations
@@ -1537,7 +1746,7 @@ export class EkoDBClient {
 
     return this.makeRequest<DistinctValuesResponse>(
       "POST",
-      `/api/distinct/${collection}/${field}`,
+      `/api/distinct/${encodeURIComponent(collection)}/${encodeURIComponent(field)}`,
       body,
       0,
       true, // Force JSON
@@ -1773,7 +1982,7 @@ export class EkoDBClient {
   ): Promise<ChatResponse> {
     return this.makeRequest<ChatResponse>(
       "POST",
-      `/api/chat/${sessionId}/messages`,
+      `/api/chat/${encodeURIComponent(sessionId)}/messages`,
       request,
       0,
       true, // Force JSON for chat operations
@@ -1793,7 +2002,7 @@ export class EkoDBClient {
   ): Promise<void> {
     await this.makeRequest(
       "POST",
-      `/api/chat/${chatId}/tool-result`,
+      `/api/chat/${encodeURIComponent(chatId)}/tool-result`,
       {
         call_id: callId,
         success,
@@ -1824,12 +2033,12 @@ export class EkoDBClient {
 
     (async () => {
       try {
-        let token = this.getToken();
+        let token = await this.getToken();
         if (!token) {
           await this.refreshToken();
-          token = this.getToken();
+          token = await this.getToken();
         }
-        const url = `${this.baseURL}/api/chat/${chatId}/messages/stream`;
+        const url = `${this.baseURL}/api/chat/${encodeURIComponent(chatId)}/messages/stream`;
 
         const response = await fetch(url, {
           method: "POST",
@@ -1851,11 +2060,10 @@ export class EkoDBClient {
           return;
         }
 
-        const body = await response.text();
-        for (const line of body.split("\n")) {
-          if (!line.startsWith("data:")) continue;
+        const emitLine = (line: string) => {
+          if (!line.startsWith("data:")) return;
           const dataStr = line.slice(5).trim();
-          if (!dataStr) continue;
+          if (!dataStr) return;
           try {
             const eventData = JSON.parse(dataStr);
             if (eventData.error) {
@@ -1881,6 +2089,30 @@ export class EkoDBClient {
           } catch {
             // skip malformed SSE data
           }
+        };
+
+        const reader = response.body?.getReader?.();
+        if (reader) {
+          // True incremental streaming: decode and emit each SSE line as soon as
+          // it arrives, rather than buffering the entire response body first.
+          const decoder = new TextDecoder();
+          let buffer = "";
+          for (;;) {
+            const { done, value } = await reader.read();
+            if (done) break;
+            buffer += decoder.decode(value, { stream: true });
+            let nl: number;
+            while ((nl = buffer.indexOf("\n")) >= 0) {
+              emitLine(buffer.slice(0, nl));
+              buffer = buffer.slice(nl + 1);
+            }
+          }
+          buffer += decoder.decode();
+          if (buffer) emitLine(buffer);
+        } else {
+          // Fallback for environments/tests without a readable body stream.
+          const body = await response.text();
+          for (const line of body.split("\n")) emitLine(line);
         }
         stream.close();
       } catch (err: any) {
@@ -1901,7 +2133,7 @@ export class EkoDBClient {
   async getChatSession(sessionId: string): Promise<ChatSessionResponse> {
     return this.makeRequest<ChatSessionResponse>(
       "GET",
-      `/api/chat/${sessionId}`,
+      `/api/chat/${encodeURIComponent(sessionId)}`,
       undefined,
       0,
       true, // Force JSON for chat operations
@@ -1944,8 +2176,8 @@ export class EkoDBClient {
 
     const queryString = params.toString();
     const path = queryString
-      ? `/api/chat/${sessionId}/messages?${queryString}`
-      : `/api/chat/${sessionId}/messages`;
+      ? `/api/chat/${encodeURIComponent(sessionId)}/messages?${queryString}`
+      : `/api/chat/${encodeURIComponent(sessionId)}/messages`;
     return this.makeRequest<GetMessagesResponse>(
       "GET",
       path,
@@ -1964,7 +2196,7 @@ export class EkoDBClient {
   ): Promise<ChatSessionResponse> {
     return this.makeRequest<ChatSessionResponse>(
       "PUT",
-      `/api/chat/${sessionId}`,
+      `/api/chat/${encodeURIComponent(sessionId)}`,
       request,
       0,
       true, // Force JSON for chat operations
@@ -1992,7 +2224,7 @@ export class EkoDBClient {
   async deleteChatSession(sessionId: string): Promise<void> {
     await this.makeRequest<void>(
       "DELETE",
-      `/api/chat/${sessionId}`,
+      `/api/chat/${encodeURIComponent(sessionId)}`,
       undefined,
       0,
       true, // Force JSON for chat operations
@@ -2008,7 +2240,7 @@ export class EkoDBClient {
   ): Promise<ChatResponse> {
     return this.makeRequest<ChatResponse>(
       "POST",
-      `/api/chat/${sessionId}/messages/${messageId}/regenerate`,
+      `/api/chat/${encodeURIComponent(sessionId)}/messages/${encodeURIComponent(messageId)}/regenerate`,
       undefined,
       0,
       true, // Force JSON for chat operations
@@ -2025,7 +2257,7 @@ export class EkoDBClient {
   ): Promise<void> {
     await this.makeRequest<void>(
       "PUT",
-      `/api/chat/${sessionId}/messages/${messageId}`,
+      `/api/chat/${encodeURIComponent(sessionId)}/messages/${encodeURIComponent(messageId)}`,
       { content },
       0,
       true, // Force JSON for chat operations
@@ -2038,7 +2270,7 @@ export class EkoDBClient {
   async deleteChatMessage(sessionId: string, messageId: string): Promise<void> {
     await this.makeRequest<void>(
       "DELETE",
-      `/api/chat/${sessionId}/messages/${messageId}`,
+      `/api/chat/${encodeURIComponent(sessionId)}/messages/${encodeURIComponent(messageId)}`,
       undefined,
       0,
       true, // Force JSON for chat operations
@@ -2055,7 +2287,7 @@ export class EkoDBClient {
   ): Promise<void> {
     await this.makeRequest<void>(
       "PATCH",
-      `/api/chat/${sessionId}/messages/${messageId}/forgotten`,
+      `/api/chat/${encodeURIComponent(sessionId)}/messages/${encodeURIComponent(messageId)}/forgotten`,
       { forgotten },
       0,
       true, // Force JSON for chat operations
@@ -2082,7 +2314,7 @@ export class EkoDBClient {
     }
     return this.makeRequest<CompactChatResponse>(
       "POST",
-      `/api/chat/${chatId}/compact`,
+      `/api/chat/${encodeURIComponent(chatId)}/compact`,
       body,
       0,
       true, // Force JSON for chat operations
@@ -2157,7 +2389,7 @@ export class EkoDBClient {
   async getChatMessage(sessionId: string, messageId: string): Promise<Record> {
     return this.makeRequest<Record>(
       "GET",
-      `/api/chat/${sessionId}/messages/${messageId}`,
+      `/api/chat/${encodeURIComponent(sessionId)}/messages/${encodeURIComponent(messageId)}`,
       undefined,
       0,
       true, // Force JSON for chat operations
@@ -2184,7 +2416,10 @@ export class EkoDBClient {
    * Get a function by ID
    */
   async getFunction(id: string): Promise<UserFunction> {
-    return this.makeRequest<UserFunction>("GET", `/api/functions/${id}`);
+    return this.makeRequest<UserFunction>(
+      "GET",
+      `/api/functions/${encodeURIComponent(id)}`,
+    );
   }
 
   /**
@@ -2199,14 +2434,21 @@ export class EkoDBClient {
    * Update an existing function by ID
    */
   async updateFunction(id: string, script: UserFunction): Promise<void> {
-    await this.makeRequest<void>("PUT", `/api/functions/${id}`, script);
+    await this.makeRequest<void>(
+      "PUT",
+      `/api/functions/${encodeURIComponent(id)}`,
+      script,
+    );
   }
 
   /**
    * Delete a function by ID
    */
   async deleteFunction(id: string): Promise<void> {
-    await this.makeRequest<void>("DELETE", `/api/functions/${id}`);
+    await this.makeRequest<void>(
+      "DELETE",
+      `/api/functions/${encodeURIComponent(id)}`,
+    );
   }
 
   /**
@@ -2218,7 +2460,7 @@ export class EkoDBClient {
   ): Promise<FunctionResult> {
     return this.makeRequest<FunctionResult>(
       "POST",
-      `/api/functions/${idOrLabel}`,
+      `/api/functions/${encodeURIComponent(idOrLabel)}`,
       params || {},
     );
   }
@@ -2465,7 +2707,7 @@ export class EkoDBClient {
   async goalStepStart(id: string, stepIndex: number): Promise<Record> {
     return this.makeRequest<Record>(
       "POST",
-      `/api/chat/goals/${encodeURIComponent(id)}/steps/${stepIndex}/start`,
+      `/api/chat/goals/${encodeURIComponent(id)}/steps/${encodeURIComponent(String(stepIndex))}/start`,
       undefined,
       0,
       true,
@@ -2480,7 +2722,7 @@ export class EkoDBClient {
   ): Promise<Record> {
     return this.makeRequest<Record>(
       "POST",
-      `/api/chat/goals/${encodeURIComponent(id)}/steps/${stepIndex}/complete`,
+      `/api/chat/goals/${encodeURIComponent(id)}/steps/${encodeURIComponent(String(stepIndex))}/complete`,
       data,
       0,
       true,
@@ -2495,7 +2737,7 @@ export class EkoDBClient {
   ): Promise<Record> {
     return this.makeRequest<Record>(
       "POST",
-      `/api/chat/goals/${encodeURIComponent(id)}/steps/${stepIndex}/fail`,
+      `/api/chat/goals/${encodeURIComponent(id)}/steps/${encodeURIComponent(String(stepIndex))}/fail`,
       data,
       0,
       true,
@@ -2956,10 +3198,18 @@ export class EkoDBClient {
   }
 
   /**
-   * Create a WebSocket client
+   * Create a WebSocket client.
+   *
+   * The token is supplied as a provider bound to this client's
+   * {@link getToken}, so every (re)connect re-evaluates (and proactively
+   * refreshes) the auth token instead of snapshotting it once. This means a
+   * reconnect after a token rotation uses the current token.
+   *
+   * @param wsURL - The WebSocket URL (e.g. `wss://host`); `/api/ws` is appended if absent.
+   * @param options - Optional reconnect/timeout tunables.
    */
-  websocket(wsURL: string): WebSocketClient {
-    return new WebSocketClient(wsURL, this.token!);
+  websocket(wsURL: string, options?: WebSocketClientOptions): WebSocketClient {
+    return new WebSocketClient(wsURL, () => this.getToken(), options);
   }
 
   // ========== RAG Helper Methods ==========
@@ -3164,6 +3414,39 @@ export interface SubscribeOptions {
   filterValue?: string;
 }
 
+/**
+ * A token provider: either a static token string, or a (possibly async)
+ * function that returns a fresh token. When a function is supplied it is
+ * re-invoked on every (re)connect, so a rotated/refreshed token is always
+ * used for the new socket instead of a stale snapshot captured once.
+ */
+export type TokenProvider =
+  | string
+  | (() => string | null | Promise<string | null>);
+
+/** Tunables for the WebSocket client's reconnect + request-timeout behavior. */
+export interface WebSocketClientOptions {
+  /**
+   * Auto-reconnect after an unexpected socket close/error (not an explicit
+   * `close()`/unsubscribe). Defaults to true.
+   */
+  autoReconnect?: boolean;
+  /** Initial backoff delay in ms before the first reconnect attempt. Default 200. */
+  reconnectInitialDelayMs?: number;
+  /** Maximum backoff delay in ms (the cap for exponential growth). Default 5000. */
+  reconnectMaxDelayMs?: number;
+  /**
+   * Maximum number of consecutive reconnect attempts before giving up.
+   * 0 or undefined means unlimited. Default unlimited.
+   */
+  reconnectMaxAttempts?: number;
+  /**
+   * Per-request timeout in ms for request/response WS calls. If no response
+   * arrives in this window the pending promise rejects. Default 30000.
+   */
+  requestTimeoutMs?: number;
+}
+
 /** EventEmitter-like interface for subscriptions and chat streams. */
 export class EventStream<_T = unknown> {
   private listeners: Map<string, Array<(data: any) => void>> = new Map();
@@ -3305,13 +3588,17 @@ export function extractRecordId(
   for (const key of extraCandidates) {
     const val = record[key];
     if (typeof val === "string") return val;
-    if (val && typeof val === "object" && "value" in val)
+    // Unwrap only a genuine typed wrapper (both "type" and "value"), matching
+    // getValue's rule so a user object like { value: 1 } isn't mistaken for one.
+    if (val && typeof val === "object" && "type" in val && "value" in val)
       return String(val.value);
   }
   for (const key of ["id", "_id"]) {
     const val = record[key];
     if (typeof val === "string") return val;
-    if (val && typeof val === "object" && "value" in val)
+    // Unwrap only a genuine typed wrapper (both "type" and "value"), matching
+    // getValue's rule so a user object like { value: 1 } isn't mistaken for one.
+    if (val && typeof val === "object" && "type" in val && "value" in val)
       return String(val.value);
   }
   return undefined;
@@ -3319,27 +3606,72 @@ export function extractRecordId(
 
 export class WebSocketClient {
   private wsURL: string;
-  private token: string;
+  private tokenProvider: () => string | null | Promise<string | null>;
   private ws: any = null;
   private dispatcherRunning = false;
   private schemaCache: SchemaCache | null = null;
+  /**
+   * Per-connection wire format, set by negotiateFormat() on every (re)connect:
+   * true once the server has Welcomed msgpack, so frames are sent/received as
+   * binary msgpack; false (JSON text) otherwise, including against an older
+   * server that never Welcomes. Keeps the transport fully back-compatible.
+   */
+  private binary = false;
+
+  // Reconnect config
+  private autoReconnect: boolean;
+  private reconnectInitialDelayMs: number;
+  private reconnectMaxDelayMs: number;
+  private reconnectMaxAttempts: number;
+  private requestTimeoutMs: number;
+
+  // Reconnect state
+  /** Set while close() is in progress so the close handler doesn't reconnect. */
+  private closed = false;
+  private reconnectAttempts = 0;
+  private reconnecting = false;
+  private connectPromise: Promise<void> | null = null;
 
   // Dispatcher state
   private pendingRequests: Map<
     string,
-    { resolve: (value: any) => void; reject: (reason: any) => void }
+    {
+      resolve: (value: any) => void;
+      reject: (reason: any) => void;
+      timer?: ReturnType<typeof setTimeout>;
+    }
   > = new Map();
   private subscriptions: Map<string, EventStream<MutationNotification>> =
     new Map();
+  /** Bookkeeping so subscriptions can be replayed on reconnect. */
+  private subscriptionParams: Map<string, SubscribeOptions | undefined> =
+    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;
+  /**
+   * @param wsURL - WebSocket URL; `/api/ws` is appended if absent.
+   * @param token - A static token string OR a {@link TokenProvider} function
+   *   re-evaluated on every (re)connect (so a refreshed token is used after a drop).
+   * @param options - Optional reconnect/timeout tunables.
+   */
+  constructor(
+    wsURL: string,
+    token: TokenProvider,
+    options: WebSocketClientOptions = {},
+  ) {
+    // Strip trailing slashes so appending `/api/ws` can't yield `//api/ws`,
+    // which warp's exact path match (`api / ws`) would reject.
+    this.wsURL = stripTrailingSlashes(wsURL);
+    this.tokenProvider = typeof token === "function" ? token : () => token;
+    this.autoReconnect = options.autoReconnect ?? true;
+    this.reconnectInitialDelayMs = options.reconnectInitialDelayMs ?? 200;
+    this.reconnectMaxDelayMs = options.reconnectMaxDelayMs ?? 5000;
+    this.reconnectMaxAttempts = options.reconnectMaxAttempts ?? 0;
+    this.requestTimeoutMs = options.requestTimeoutMs ?? 30000;
   }
 
   private messageCounter = 0;
@@ -3350,11 +3682,39 @@ export class WebSocketClient {
   }
 
   /**
-   * Connect and start the dispatcher.
+   * Compute the capped exponential backoff (with jitter) for a reconnect
+   * attempt. attempt 0 -> ~initial, growing x2 each time up to the max cap.
+   * Jitter is +/-25% to avoid thundering-herd reconnect storms.
+   * @internal exposed for testing
+   */
+  computeBackoff(attempt: number): number {
+    const base = Math.min(
+      this.reconnectInitialDelayMs * 2 ** attempt,
+      this.reconnectMaxDelayMs,
+    );
+    const jitter = base * 0.25 * (Math.random() * 2 - 1);
+    return Math.max(0, Math.round(base + jitter));
+  }
+
+  /**
+   * Connect and start the dispatcher. Re-evaluates the token provider so the
+   * current/refreshed token is used for this socket.
    */
   private async ensureConnected(): Promise<void> {
     if (this.ws && this.dispatcherRunning) return;
+    // Coalesce concurrent connect attempts onto a single in-flight promise.
+    if (this.connectPromise) return this.connectPromise;
+    // Clear the intentional-close flag only for user-initiated connects. During
+    // a reconnect cycle this stays untouched so a concurrent close() can't be
+    // undone and have the reconnect proceed against the user's intent.
+    if (!this.reconnecting) this.closed = false;
+    this.connectPromise = this.openSocket().finally(() => {
+      this.connectPromise = null;
+    });
+    return this.connectPromise;
+  }
 
+  private async openSocket(): Promise<void> {
     const WebSocket = (await import("ws")).default;
 
     let url = this.wsURL;
@@ -3362,9 +3722,19 @@ export class WebSocketClient {
       url += "/api/ws";
     }
 
+    // Re-evaluate the token on every (re)connect — never a stale snapshot.
+    const token = await this.tokenProvider();
+    if (!token) {
+      // Fail fast with a clear error instead of sending `Bearer null`, which
+      // would surface as a confusing 401 from the server.
+      throw new Error(
+        "WebSocket auth token is unavailable (the token provider returned null/empty)",
+      );
+    }
+
     this.ws = new WebSocket(url, {
       headers: {
-        Authorization: `Bearer ${this.token}`,
+        Authorization: `Bearer ${token}`,
       },
     });
 
@@ -3373,42 +3743,245 @@ export class WebSocketClient {
       this.ws.on("error", (err: Error) => reject(err));
     });
 
+    // Negotiate the wire format before the dispatcher starts so the Welcome is
+    // consumed here (not by routeMessage), and before any real frame is sent
+    // (resubscribeAll runs only after this resolves).
+    await this.negotiateFormat(this.ws);
+
     this.spawnDispatcher();
   }
 
+  /**
+   * Additive capability handshake: offer msgpack and, if the server Welcomes
+   * it, switch this connection to binary msgpack frames; otherwise stay on JSON
+   * text. The Welcome (a text frame) is read with a one-shot listener and a
+   * timeout so an older server that never answers — or answers with an Error —
+   * simply leaves the connection on JSON. Best-effort and never throws: JSON
+   * always works.
+   */
+  private async negotiateFormat(socket: any): Promise<void> {
+    this.binary = false;
+    const welcome = await new Promise<any | null>((resolve) => {
+      const onMsg = (data: Buffer) => {
+        clearTimeout(timer);
+        try {
+          resolve(JSON.parse(data.toString()));
+        } catch {
+          resolve(null);
+        }
+      };
+      // Only caps the wait when no Welcome comes (a silent/old server); the
+      // listener resolves immediately when it does arrive. 2s comfortably exceeds
+      // the handshake round-trip even on high-latency links.
+      const timer = setTimeout(() => {
+        socket.off("message", onMsg);
+        resolve(null);
+      }, 2000);
+      socket.once("message", onMsg);
+      try {
+        socket.send(
+          JSON.stringify({
+            type: "Hello",
+            payload: { formats: ["msgpack", "json"] },
+          }),
+        );
+      } catch {
+        clearTimeout(timer);
+        socket.off("message", onMsg);
+        resolve(null);
+      }
+    });
+    if (
+      welcome &&
+      welcome.type === "Welcome" &&
+      welcome.payload?.format === "msgpack"
+    ) {
+      this.binary = true;
+    }
+  }
+
+  /**
+   * Send a request object on the active socket using the negotiated format:
+   * binary msgpack when the server Welcomed it, JSON text otherwise. The single
+   * write point so every request honors the negotiated transport.
+   */
+  private sendFrame(obj: any): void {
+    this.ws.send(this.binary ? encode(obj) : JSON.stringify(obj));
+  }
+
   private spawnDispatcher(): void {
     if (this.dispatcherRunning) return;
     this.dispatcherRunning = true;
 
-    this.ws.on("message", (data: Buffer) => {
+    // Capture the socket this dispatcher is bound to. After a reconnect, the old
+    // socket may still emit late close/error events; ignore them so they don't
+    // tear down the replacement connection.
+    const socket = this.ws;
+
+    socket.on("message", (data: Buffer, isBinary: boolean) => {
+      if (this.ws !== socket) return;
       try {
-        const msg = JSON.parse(data.toString());
+        // A binary frame is msgpack (the server only sends binary once it has
+        // Welcomed msgpack); a text frame is JSON. Decode by frame type so the
+        // routed value is identical regardless of negotiated transport.
+        const msg = isBinary
+          ? (decode(data) as any)
+          : 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
+    // Both "close" and "error" mean this socket is dead. ws typically emits
+    // "error" followed by "close", so route both through one handler and let the
+    // identity check dedupe: the first to fire nulls this.ws, the second no-ops.
+    const onDown = () => {
+      if (this.ws !== socket) return;
+      this.handleDisconnect();
+    };
+    socket.on("close", onDown);
+    socket.on("error", onDown);
+  }
+
+  /**
+   * Reject in-flight requests and tear down the dead socket. If the close was
+   * unexpected (not an explicit `close()`) and auto-reconnect is enabled,
+   * schedule a reconnect that re-sends the active subscriptions.
+   */
+  private handleDisconnect(): void {
+    this.dispatcherRunning = false;
+    this.ws = null;
+
+    // Reject all in-flight pending requests so callers don't hang forever.
+    for (const [, pending] of this.pendingRequests) {
+      if (pending.timer) clearTimeout(pending.timer);
+      pending.reject(new Error("WebSocket connection closed"));
+    }
+    this.pendingRequests.clear();
+    if (this.registerToolsAck) {
+      this.registerToolsAck.reject(new Error("WebSocket connection closed"));
+      this.registerToolsAck = null;
+    }
+    // Close all chat streams (they are one-shot; not replayed on reconnect).
+    for (const [, stream] of this.chatStreams) {
+      stream.emit("event", { type: "error", error: "Connection closed" });
+      stream.close();
+    }
+    this.chatStreams.clear();
+
+    const shouldReconnect =
+      this.autoReconnect && !this.closed && this.subscriptionParams.size > 0;
+
+    if (shouldReconnect) {
+      this.scheduleReconnect();
+    } else {
+      // No reconnect: tear down subscriptions too.
       for (const [, stream] of this.subscriptions) {
         stream.close();
       }
       this.subscriptions.clear();
-      this.ws = null;
-    });
+      this.subscriptionParams.clear();
+    }
+  }
+
+  /**
+   * Reconnect with capped exponential backoff + jitter, then re-send the
+   * subscribe messages for every active subscription so the SAME EventStream
+   * keeps delivering mutations after a transient drop.
+   */
+  private scheduleReconnect(): void {
+    if (this.reconnecting) return;
+    this.reconnecting = true;
+
+    const attempt = async (): Promise<void> => {
+      // Bail if the client was closed, or if every subscription was torn down
+      // (e.g. unsubscribed) while a reconnect was in-flight — reconnect was only
+      // opted into because subscriptions existed, so there's nothing to restore.
+      if (this.closed || this.subscriptionParams.size === 0) {
+        this.reconnecting = false;
+        return;
+      }
+      if (
+        this.reconnectMaxAttempts > 0 &&
+        this.reconnectAttempts >= this.reconnectMaxAttempts
+      ) {
+        // Give up: tear down subscriptions and notify consumers.
+        this.reconnecting = false;
+        for (const [, stream] of this.subscriptions) {
+          stream.emit("error", "WebSocket reconnect failed");
+          stream.close();
+        }
+        this.subscriptions.clear();
+        this.subscriptionParams.clear();
+        return;
+      }
+
+      const delay = this.computeBackoff(this.reconnectAttempts);
+      this.reconnectAttempts++;
+      await new Promise((r) => setTimeout(r, delay));
+
+      // Re-check after the backoff delay: close() or a full unsubscribe may have
+      // happened while we were waiting, in which case skip reopening the socket.
+      if (this.closed || this.subscriptionParams.size === 0) {
+        this.reconnecting = false;
+        return;
+      }
+
+      try {
+        // Route through ensureConnected() so a request-driven connect and this
+        // reconnect share one in-flight connectPromise/socket — opening two live
+        // sockets would misroute responses.
+        await this.ensureConnected();
+        // close() may have been called while the connect was in-flight; if so,
+        // tear down the freshly-opened socket instead of leaving it orphaned.
+        if (this.closed) {
+          try {
+            this.ws?.close?.();
+          } catch {
+            /* already closing */
+          }
+          this.ws = null;
+          this.dispatcherRunning = false;
+          this.reconnecting = false;
+          return;
+        }
+        // Success — reset backoff and replay subscriptions.
+        this.reconnectAttempts = 0;
+        this.reconnecting = false;
+        await this.resubscribeAll();
+      } catch {
+        // Connect failed — schedule the next attempt WITHOUT recursive await so
+        // a prolonged outage can't build an unbounded promise chain.
+        setTimeout(() => void attempt(), 0);
+      }
+    };
+
+    void attempt();
+  }
+
+  /** Re-send Subscribe frames for every tracked subscription after a reconnect. */
+  private async resubscribeAll(): Promise<void> {
+    for (const [collection, options] of this.subscriptionParams) {
+      const stream = this.subscriptions.get(collection);
+      if (!stream || stream.closed) continue;
+      const messageId = this.genMessageId();
+      const request: any = {
+        type: "Subscribe",
+        messageId,
+        payload: {
+          collection,
+          ...(options?.filterField && { filter_field: options.filterField }),
+          ...(options?.filterValue && { filter_value: options.filterValue }),
+        },
+      };
+      try {
+        await this.sendRequest(request);
+      } catch {
+        // If the re-subscribe ack fails, leave it tracked; the next
+        // disconnect/reconnect cycle will attempt it again.
+      }
+    }
   }
 
   private routeMessage(msg: any): void {
@@ -3423,14 +3996,7 @@ export class WebSocketClient {
           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;
+          matched = this.settlePending(messageId, msg.type === "Error", msg);
         }
         if (!matched && this.registerToolsAck) {
           const ack = this.registerToolsAck;
@@ -3442,18 +4008,14 @@ export class WebSocketClient {
           }
           matched = true;
         }
-        // Server doesn't echo messageId — if there's exactly one pending
+        // Server doesn't echo messageId at all — 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);
-          }
+        // Only when messageId is absent: a present-but-unmatched id means a late
+        // response for an already-settled/timed-out request, which must NOT be
+        // misrouted to whatever request happens to still be pending.
+        if (!matched && !messageId && this.pendingRequests.size === 1) {
+          const key = this.pendingRequests.keys().next().value!;
+          this.settlePending(key, msg.type === "Error", msg);
         }
         break;
       }
@@ -3555,16 +4117,52 @@ export class WebSocketClient {
     const messageId = request.messageId || request.message_id;
 
     return new Promise((resolve, reject) => {
-      this.pendingRequests.set(messageId, { resolve, reject });
+      // Per-request timeout: reject if no response arrives in the window so a
+      // dropped/never-answered response can't leave the promise pending forever.
+      let timer: ReturnType<typeof setTimeout> | undefined;
+      if (this.requestTimeoutMs > 0) {
+        timer = setTimeout(() => {
+          if (this.pendingRequests.delete(messageId)) {
+            reject(
+              new Error(
+                `WebSocket request "${request.type}" timed out after ${this.requestTimeoutMs}ms`,
+              ),
+            );
+          }
+        }, this.requestTimeoutMs);
+        // Don't keep the process alive just for this timer.
+        (timer as any)?.unref?.();
+      }
+
+      this.pendingRequests.set(messageId, { resolve, reject, timer });
       try {
-        this.ws.send(JSON.stringify(request));
+        this.sendFrame(request);
       } catch (err) {
         this.pendingRequests.delete(messageId);
+        if (timer) clearTimeout(timer);
         reject(err);
       }
     });
   }
 
+  /** Resolve/reject a pending request, clearing its timeout timer. */
+  private settlePending(
+    messageId: string,
+    isError: boolean,
+    msg: any,
+  ): boolean {
+    const pending = this.pendingRequests.get(messageId);
+    if (!pending) return false;
+    this.pendingRequests.delete(messageId);
+    if (pending.timer) clearTimeout(pending.timer);
+    if (isError) {
+      pending.reject(new Error(msg.message || "Unknown error"));
+    } else {
+      pending.resolve(msg.payload);
+    }
+    return true;
+  }
+
   /**
    * Find all records in a collection via WebSocket.
    */
@@ -3595,6 +4193,8 @@ export class WebSocketClient {
     const messageId = this.genMessageId();
     const stream = new EventStream<MutationNotification>();
     this.subscriptions.set(collection, stream);
+    // Track params so the subscription can be replayed on reconnect.
+    this.subscriptionParams.set(collection, options);
 
     const request: any = {
       type: "Subscribe",
@@ -3611,11 +4211,45 @@ export class WebSocketClient {
       await this.sendRequest(request);
     } catch (err) {
       this.subscriptions.delete(collection);
+      this.subscriptionParams.delete(collection);
       throw err;
     }
     return stream;
   }
 
+  /**
+   * Unsubscribe from a collection's mutation notifications. This is an
+   * intentional teardown, so the subscription is NOT replayed on reconnect.
+   */
+  unsubscribe(collection: string): void {
+    const stream = this.subscriptions.get(collection);
+    this.subscriptions.delete(collection);
+    this.subscriptionParams.delete(collection);
+    if (stream && !stream.closed) {
+      stream.close();
+    }
+    // Best-effort: tell the server to stop streaming this collection (the
+    // server already handles an Unsubscribe frame). If the socket isn't open
+    // the local teardown above suffices, since the server drops subscriptions
+    // when the connection closes. A unique messageId is attached so the
+    // server's Success ack carries a correlation id: it has no pending request
+    // to match, so it is simply ignored — and because the id is present, the
+    // single-pending fallback can't misroute it to an unrelated request.
+    if (this.ws && this.ws.readyState === 1 /* WebSocket.OPEN */) {
+      try {
+        this.sendFrame({
+          type: "Unsubscribe",
+          messageId: this.genMessageId(),
+          payload: { collection },
+        });
+      } catch {
+        // Best-effort: the socket can close between the readyState check and the
+        // send. Local teardown already happened, so swallow the failure rather
+        // than throw out of a void teardown call.
+      }
+    }
+  }
+
   /**
    * Send a chat message and receive a streaming response.
    * Returns an EventStream that emits "event" with ChatStreamEvent objects.
@@ -3651,7 +4285,7 @@ export class WebSocketClient {
       },
     };
 
-    this.ws.send(JSON.stringify(request));
+    this.sendFrame(request);
     return stream;
   }
 
@@ -3677,7 +4311,7 @@ export class WebSocketClient {
         resolve: () => resolve(),
         reject: (err) => reject(err),
       };
-      this.ws.send(JSON.stringify(request));
+      this.sendFrame(request);
     });
   }
 
@@ -3704,7 +4338,27 @@ export class WebSocketClient {
       },
     };
 
-    this.ws.send(JSON.stringify(request));
+    this.sendFrame(request);
+  }
+
+  /**
+   * Cancel an in-flight streaming chat. Fire-and-forget: tells the server to
+   * stop generating tokens for the given chat.
+   */
+  async cancelChat(chatId: string): Promise<void> {
+    await this.ensureConnected();
+
+    // Attach a unique messageId (same generator as unsubscribe). Any Success ack
+    // from the server then carries a correlation id: it has no pending request to
+    // match, so it is ignored — and because the id is present, the dispatcher's
+    // single-pending fallback can't misroute the ack to an unrelated request.
+    const request = {
+      type: "CancelChat",
+      messageId: this.genMessageId(),
+      payload: { chat_id: chatId },
+    };
+
+    this.sendFrame(request);
   }
 
   /**
@@ -3930,8 +4584,45 @@ export class WebSocketClient {
 
   /**
    * Close the WebSocket connection.
+   *
+   * This is an INTENTIONAL close: it disables auto-reconnect, rejects any
+   * in-flight requests, and tears down all subscriptions/chat streams so
+   * nothing is replayed afterward.
    */
   close(): void {
+    // Mark intentional so the close handler doesn't trigger a reconnect.
+    this.closed = true;
+    this.reconnecting = false;
+
+    // Reject any in-flight requests and clear their timers.
+    for (const [, pending] of this.pendingRequests) {
+      if (pending.timer) clearTimeout(pending.timer);
+      pending.reject(new Error("WebSocket connection closed"));
+    }
+    this.pendingRequests.clear();
+
+    // Tear down subscriptions + their replay bookkeeping.
+    for (const [, stream] of this.subscriptions) {
+      if (!stream.closed) stream.close();
+    }
+    this.subscriptions.clear();
+    this.subscriptionParams.clear();
+
+    // Reject any in-flight tool registration ack. Done here (not just in the
+    // ws "close" handler) so it's cleaned up even when this.ws is already null.
+    if (this.registerToolsAck) {
+      this.registerToolsAck.reject(new Error("WebSocket connection closed"));
+      this.registerToolsAck = null;
+    }
+
+    // Tear down chat streams immediately; they are one-shot and not replayed,
+    // and we can't rely on the underlying ws "close" event having fired.
+    for (const [, stream] of this.chatStreams) {
+      stream.emit("event", { type: "error", error: "Connection closed" });
+      stream.close();
+    }
+    this.chatStreams.clear();
+
     if (this.ws) {
       this.ws.close();
       this.ws = null;