@ekodb/ekodb-client 0.13.0 → 0.15.0

package/src/client.ts

@@ -203,6 +203,7 @@ export interface ChatMessageRequest {
   force_summarize?: boolean;
   max_iterations?: number;
   tool_config?: ToolConfig;
+  llm_model?: string;
 }
 
 export interface TokenUsage {
@@ -371,6 +372,7 @@ 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 format: SerializationFormat;
@@ -441,22 +443,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 the current authentication token.
-   * Returns null if not yet authenticated. Call refreshToken() first.
-   */
-  getToken(): string | null {
+   * 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.
+   * 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;
+    }
   }
 
   /**
@@ -1499,6 +1553,54 @@ export class EkoDBClient {
 
   // ========== Chat Methods ==========
 
+  /**
+   * Execute a tool via ekoDB's server-side tool pipeline.
+   *
+   * Calls POST /api/chat/tools/execute which goes through the same
+   * execute_tool function as the LLM tool-calling loop — with all
+   * collection filtering, permission enforcement, and internal collection
+   * blocking. No LLM round-trip.
+   *
+   * @returns The tool result if executed, or null if the server doesn't
+   * support the endpoint (older ekoDB versions).
+   */
+  async executeTool(
+    toolName: string,
+    params: { [key: string]: any },
+    chatId?: string,
+  ): Promise<any | null> {
+    const body: { [key: string]: any } = { tool: toolName, params };
+    if (chatId) {
+      body.chat_id = chatId;
+    }
+
+    try {
+      const result = await this.makeRequest<{ [key: string]: any }>(
+        "POST",
+        "/api/chat/tools/execute",
+        body,
+        0,
+        true, // Force JSON for chat operations
+      );
+
+      if (result.success) {
+        return result.result;
+      } else {
+        throw new Error(result.error || "tool execution failed");
+      }
+    } catch (err: any) {
+      // Server doesn't have the endpoint (404) or route mismatch (405)
+      // Parse status from makeRequest error format: "Request failed with status NNN: ..."
+      const message = String(err?.message ?? "");
+      const match = message.match(/Request failed with status (\d+):/);
+      const status = match ? parseInt(match[1], 10) : undefined;
+      if (status === 404 || status === 405) {
+        return null;
+      }
+      throw err;
+    }
+  }
+
   /**
    * Create a new chat session
    */
@@ -1542,6 +1644,115 @@ export class EkoDBClient {
     );
   }
 
+  /**
+   * 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
    */
@@ -1558,6 +1769,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
    */
@@ -1942,6 +2243,517 @@ export class EkoDBClient {
     );
   }
 
+  // ========================================================================
+  // GOAL API
+  // ========================================================================
+
+  /** Create a new goal */
+  async goalCreate(data: Record): Promise<Record> {
+    return this.makeRequest<Record>("POST", "/api/chat/goals", data, 0, true);
+  }
+
+  /** List all goals */
+  async goalList(): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      "/api/chat/goals",
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** Get a goal by ID */
+  async goalGet(id: string): Promise<Record> {
+    return this.makeRequest<Record>(
+      "GET",
+      `/api/chat/goals/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  /** 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,
+    );
+  }
+
+  /** Delete a goal by ID */
+  async goalDelete(id: string): Promise<void> {
+    await this.makeRequest<void>(
+      "DELETE",
+      `/api/chat/goals/${encodeURIComponent(id)}`,
+      undefined,
+      0,
+      true,
+    );
+  }
+
+  // ── 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
   // ========================================================================
@@ -2143,6 +2955,8 @@ export type ChatStreamEvent =
       tokenUsage?: any;
       toolCallHistory?: any;
       executionTimeMs: number;
+      /** Model's context window size in tokens. */
+      contextWindow?: number;
     }
   | {
       type: "toolCall";
@@ -2309,7 +3123,13 @@ export class WebSocketClient {
     switch (msg.type) {
       case "Success":
       case "Error": {
-        const messageId = msg.payload?.message_id || msg.payload?.messageId;
+        // 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);
@@ -2318,7 +3138,9 @@ export class WebSocketClient {
           } else {
             pending.resolve(msg.payload);
           }
-        } else if (this.registerToolsAck) {
+          matched = true;
+        }
+        if (!matched && this.registerToolsAck) {
           const ack = this.registerToolsAck;
           this.registerToolsAck = null;
           if (msg.type === "Error") {
@@ -2326,6 +3148,20 @@ export class WebSocketClient {
           } 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;
       }
@@ -2371,6 +3207,8 @@ export class WebSocketClient {
               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();
@@ -2566,6 +3404,33 @@ export class WebSocketClient {
     this.ws.send(JSON.stringify(request));
   }
 
+  /**
+   * 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.
    */