@deepagents/context 0.16.0 → 0.17.1

package/dist/index.js

@@ -1,7 +1,5 @@
 // packages/context/src/lib/agent.ts
-import { groq } from "@ai-sdk/groq";
 import {
-  NoSuchToolError,
   Output,
   convertToModelMessages,
   createUIMessageStream,
@@ -13,7 +11,7 @@ import {
 } from "ai";
 import chalk from "chalk";
 import "zod";
-import "@deepagents/agent";
+import { createRepairToolCall } from "@deepagents/agent";
 
 // packages/context/src/lib/fragments.ts
 import { generateId } from "ai";
@@ -169,7 +167,7 @@ var Agent = class _Agent {
       stopWhen: stepCountIs(25),
       tools: this.#options.tools,
       experimental_context: contextVariables,
-      experimental_repairToolCall: repairToolCall,
+      experimental_repairToolCall: createRepairToolCall(this.#options.model),
       toolChoice: this.#options.toolChoice,
       onStepFinish: (step) => {
         const toolCall = step.toolCalls.at(-1);
@@ -225,7 +223,7 @@ var Agent = class _Agent {
       model: this.#options.model,
       system: systemPrompt,
       messages: await convertToModelMessages(messages),
-      experimental_repairToolCall: repairToolCall,
+      experimental_repairToolCall: createRepairToolCall(this.#options.model),
       stopWhen: stepCountIs(50),
       experimental_transform: config?.transform ?? smoothStream(),
       tools: this.#options.tools,
@@ -365,7 +363,7 @@ function structuredOutput(options) {
         system: systemPrompt,
         messages: await convertToModelMessages(messages),
         stopWhen: stepCountIs(25),
-        experimental_repairToolCall: repairToolCall,
+        experimental_repairToolCall: createRepairToolCall(options.model),
         experimental_context: contextVariables,
         output: Output.object({ schema: options.schema }),
         tools: options.tools
@@ -387,7 +385,7 @@ function structuredOutput(options) {
         providerOptions: options.providerOptions,
         model: options.model,
         system: systemPrompt,
-        experimental_repairToolCall: repairToolCall,
+        experimental_repairToolCall: createRepairToolCall(options.model),
         messages: await convertToModelMessages(messages),
         stopWhen: stepCountIs(50),
         experimental_transform: config?.transform ?? smoothStream(),
@@ -398,34 +396,6 @@ function structuredOutput(options) {
     }
   };
 }
-var repairToolCall = async ({
-  toolCall,
-  tools,
-  inputSchema,
-  error
-}) => {
-  console.log(
-    `Debug: ${chalk.yellow("RepairingToolCall")}: ${chalk.bgYellow(toolCall.toolName)}`,
-    error.name,
-    JSON.stringify(toolCall)
-  );
-  if (NoSuchToolError.isInstance(error)) {
-    return null;
-  }
-  const tool = tools[toolCall.toolName];
-  const { output } = await generateText({
-    model: groq("openai/gpt-oss-20b"),
-    output: Output.object({ schema: tool.inputSchema }),
-    prompt: [
-      `The model tried to call the tool "${toolCall.toolName}" with the following inputs:`,
-      JSON.stringify(toolCall.input),
-      `The tool accepts the following schema:`,
-      JSON.stringify(inputSchema(toolCall)),
-      "Please fix the inputs."
-    ].join("\n")
-  });
-  return { ...toolCall, input: JSON.stringify(output) };
-};
 function writeText(writer, text) {
   const feedbackPartId = generateId2();
   writer.write({
@@ -1593,6 +1563,13 @@ var ContextEngine = class {
   get branch() {
     return this.#branchName;
   }
+  /**
+   * Get the current branch head message ID.
+   * Returns undefined if no messages have been saved yet.
+   */
+  get headMessageId() {
+    return this.#branch?.headMessageId ?? void 0;
+  }
   /**
    * Get metadata for the current chat.
    * Returns null if the chat hasn't been initialized yet.
@@ -1699,7 +1676,7 @@ var ContextEngine = class {
   async save(options) {
     await this.#ensureInitialized();
     if (this.#pendingMessages.length === 0) {
-      return;
+      return { headMessageId: this.#branch?.headMessageId ?? void 0 };
     }
     const shouldBranch = options?.branch ?? true;
     for (let i = 0; i < this.#pendingMessages.length; i++) {
@@ -1751,6 +1728,7 @@ var ContextEngine = class {
     await this.#store.updateBranchHead(this.#branch.id, parentId);
     this.#branch.headMessageId = parentId;
     this.#pendingMessages = [];
+    return { headMessageId: this.#branch.headMessageId ?? void 0 };
   }
   /**
    * Resolve a lazy fragment by finding the appropriate ID.
@@ -3916,18 +3894,113 @@ var InMemoryContextStore = class extends SqliteContextStore {
 // packages/context/src/lib/store/postgres.store.ts
 import { createRequire } from "node:module";
 
-// packages/context/src/lib/store/ddl.postgres.sql
-var ddl_postgres_default = "-- Context Store DDL for PostgreSQL\n-- This schema implements a DAG-based message history with branching and checkpoints.\n\n-- Chats table\n-- createdAt/updatedAt: DEFAULT for insert, inline SET for updates\nCREATE TABLE IF NOT EXISTS chats (\n  id TEXT PRIMARY KEY,\n  userId TEXT NOT NULL,\n  title TEXT,\n  metadata JSONB,\n  createdAt BIGINT NOT NULL DEFAULT (EXTRACT(EPOCH FROM NOW()) * 1000)::BIGINT,\n  updatedAt BIGINT NOT NULL DEFAULT (EXTRACT(EPOCH FROM NOW()) * 1000)::BIGINT\n);\n\nCREATE INDEX IF NOT EXISTS idx_chats_updatedAt ON chats(updatedAt);\nCREATE INDEX IF NOT EXISTS idx_chats_userId ON chats(userId);\nCREATE INDEX IF NOT EXISTS idx_chats_metadata ON chats USING GIN (metadata);\n\n-- Messages table (nodes in the DAG)\nCREATE TABLE IF NOT EXISTS messages (\n  id TEXT PRIMARY KEY,\n  chatId TEXT NOT NULL,\n  parentId TEXT,\n  name TEXT NOT NULL,\n  type TEXT,\n  data JSONB NOT NULL,\n  createdAt BIGINT NOT NULL,\n  FOREIGN KEY (chatId) REFERENCES chats(id) ON DELETE CASCADE,\n  FOREIGN KEY (parentId) REFERENCES messages(id)\n);\n\nCREATE INDEX IF NOT EXISTS idx_messages_chatId ON messages(chatId);\nCREATE INDEX IF NOT EXISTS idx_messages_parentId ON messages(parentId);\n\n-- Branches table (pointers to head messages)\nCREATE TABLE IF NOT EXISTS branches (\n  id TEXT PRIMARY KEY,\n  chatId TEXT NOT NULL,\n  name TEXT NOT NULL,\n  headMessageId TEXT,\n  isActive BOOLEAN NOT NULL DEFAULT FALSE,\n  createdAt BIGINT NOT NULL,\n  FOREIGN KEY (chatId) REFERENCES chats(id) ON DELETE CASCADE,\n  FOREIGN KEY (headMessageId) REFERENCES messages(id),\n  UNIQUE(chatId, name)\n);\n\nCREATE INDEX IF NOT EXISTS idx_branches_chatId ON branches(chatId);\n\n-- Checkpoints table (pointers to message nodes)\nCREATE TABLE IF NOT EXISTS checkpoints (\n  id TEXT PRIMARY KEY,\n  chatId TEXT NOT NULL,\n  name TEXT NOT NULL,\n  messageId TEXT NOT NULL,\n  createdAt BIGINT NOT NULL,\n  FOREIGN KEY (chatId) REFERENCES chats(id) ON DELETE CASCADE,\n  FOREIGN KEY (messageId) REFERENCES messages(id),\n  UNIQUE(chatId, name)\n);\n\nCREATE INDEX IF NOT EXISTS idx_checkpoints_chatId ON checkpoints(chatId);\n\n-- Full-text search using tsvector + GIN index\nCREATE TABLE IF NOT EXISTS messages_fts (\n  messageId TEXT PRIMARY KEY REFERENCES messages(id) ON DELETE CASCADE,\n  chatId TEXT NOT NULL,\n  name TEXT NOT NULL,\n  content TEXT NOT NULL,\n  content_vector TSVECTOR\n);\n\nCREATE INDEX IF NOT EXISTS idx_messages_fts_vector ON messages_fts USING GIN(content_vector);\nCREATE INDEX IF NOT EXISTS idx_messages_fts_chatId ON messages_fts(chatId);\n\n-- Trigger to automatically update tsvector on insert/update\nCREATE OR REPLACE FUNCTION messages_fts_update_vector() RETURNS TRIGGER AS $$\nBEGIN\n  NEW.content_vector := to_tsvector('english', NEW.content);\n  RETURN NEW;\nEND;\n$$ LANGUAGE plpgsql;\n\nDROP TRIGGER IF EXISTS messages_fts_vector_update ON messages_fts;\nCREATE TRIGGER messages_fts_vector_update\n  BEFORE INSERT OR UPDATE ON messages_fts\n  FOR EACH ROW\n  EXECUTE FUNCTION messages_fts_update_vector();\n";
+// packages/context/src/lib/store/ddl.postgres.ts
+function storeDDL(schema) {
+  return `
+CREATE SCHEMA IF NOT EXISTS "${schema}";
+
+CREATE TABLE IF NOT EXISTS "${schema}"."chats" (
+  id TEXT PRIMARY KEY,
+  userId TEXT NOT NULL,
+  title TEXT,
+  metadata JSONB,
+  createdAt BIGINT NOT NULL DEFAULT (EXTRACT(EPOCH FROM NOW()) * 1000)::BIGINT,
+  updatedAt BIGINT NOT NULL DEFAULT (EXTRACT(EPOCH FROM NOW()) * 1000)::BIGINT
+);
+
+CREATE INDEX IF NOT EXISTS "idx_${schema}_chats_updatedAt" ON "${schema}"."chats"(updatedAt);
+CREATE INDEX IF NOT EXISTS "idx_${schema}_chats_userId" ON "${schema}"."chats"(userId);
+CREATE INDEX IF NOT EXISTS "idx_${schema}_chats_metadata" ON "${schema}"."chats" USING GIN (metadata);
+
+CREATE TABLE IF NOT EXISTS "${schema}"."messages" (
+  id TEXT PRIMARY KEY,
+  chatId TEXT NOT NULL,
+  parentId TEXT,
+  name TEXT NOT NULL,
+  type TEXT,
+  data JSONB NOT NULL,
+  createdAt BIGINT NOT NULL,
+  FOREIGN KEY (chatId) REFERENCES "${schema}"."chats"(id) ON DELETE CASCADE,
+  FOREIGN KEY (parentId) REFERENCES "${schema}"."messages"(id)
+);
+
+CREATE INDEX IF NOT EXISTS "idx_${schema}_messages_chatId" ON "${schema}"."messages"(chatId);
+CREATE INDEX IF NOT EXISTS "idx_${schema}_messages_parentId" ON "${schema}"."messages"(parentId);
+
+CREATE TABLE IF NOT EXISTS "${schema}"."branches" (
+  id TEXT PRIMARY KEY,
+  chatId TEXT NOT NULL,
+  name TEXT NOT NULL,
+  headMessageId TEXT,
+  isActive BOOLEAN NOT NULL DEFAULT FALSE,
+  createdAt BIGINT NOT NULL,
+  FOREIGN KEY (chatId) REFERENCES "${schema}"."chats"(id) ON DELETE CASCADE,
+  FOREIGN KEY (headMessageId) REFERENCES "${schema}"."messages"(id),
+  UNIQUE(chatId, name)
+);
+
+CREATE INDEX IF NOT EXISTS "idx_${schema}_branches_chatId" ON "${schema}"."branches"(chatId);
+
+CREATE TABLE IF NOT EXISTS "${schema}"."checkpoints" (
+  id TEXT PRIMARY KEY,
+  chatId TEXT NOT NULL,
+  name TEXT NOT NULL,
+  messageId TEXT NOT NULL,
+  createdAt BIGINT NOT NULL,
+  FOREIGN KEY (chatId) REFERENCES "${schema}"."chats"(id) ON DELETE CASCADE,
+  FOREIGN KEY (messageId) REFERENCES "${schema}"."messages"(id),
+  UNIQUE(chatId, name)
+);
+
+CREATE INDEX IF NOT EXISTS "idx_${schema}_checkpoints_chatId" ON "${schema}"."checkpoints"(chatId);
+
+CREATE TABLE IF NOT EXISTS "${schema}"."messages_fts" (
+  messageId TEXT PRIMARY KEY REFERENCES "${schema}"."messages"(id) ON DELETE CASCADE,
+  chatId TEXT NOT NULL,
+  name TEXT NOT NULL,
+  content TEXT NOT NULL,
+  content_vector TSVECTOR
+);
+
+CREATE INDEX IF NOT EXISTS "idx_${schema}_messages_fts_vector" ON "${schema}"."messages_fts" USING GIN(content_vector);
+CREATE INDEX IF NOT EXISTS "idx_${schema}_messages_fts_chatId" ON "${schema}"."messages_fts"(chatId);
+
+CREATE OR REPLACE FUNCTION "${schema}"."messages_fts_update_vector"() RETURNS TRIGGER AS $$
+BEGIN
+  NEW.content_vector := to_tsvector('english', NEW.content);
+  RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
+
+DROP TRIGGER IF EXISTS "${schema}_messages_fts_vector_update" ON "${schema}"."messages_fts";
+CREATE TRIGGER "${schema}_messages_fts_vector_update"
+  BEFORE INSERT OR UPDATE ON "${schema}"."messages_fts"
+  FOR EACH ROW
+  EXECUTE FUNCTION "${schema}"."messages_fts_update_vector"();
+`;
+}
 
 // packages/context/src/lib/store/postgres.store.ts
 var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   #pool;
-  #initialized;
+  #schema;
+  #ownsPool;
+  #isInitialized = false;
   constructor(options) {
     super();
+    const schema = options.schema ?? "public";
+    if (!/^[a-zA-Z_]\w*$/.test(schema)) {
+      throw new Error(`Invalid schema name: "${schema}"`);
+    }
+    this.#schema = schema;
     const pg = _PostgresContextStore.#requirePg();
-    this.#pool = typeof options.pool === "string" ? new pg.Pool({ connectionString: options.pool }) : new pg.Pool(options.pool);
-    this.#initialized = this.#initialize();
+    if (options.pool instanceof pg.Pool) {
+      this.#pool = options.pool;
+      this.#ownsPool = false;
+    } else {
+      this.#pool = typeof options.pool === "string" ? new pg.Pool({ connectionString: options.pool }) : new pg.Pool(options.pool);
+      this.#ownsPool = true;
+    }
   }
   static #requirePg() {
     try {
@@ -3939,21 +4012,27 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
       );
     }
   }
-  async #initialize() {
-    await this.#pool.query(ddl_postgres_default);
+  #t(name) {
+    return `"${this.#schema}"."${name}"`;
   }
-  /**
-   * Ensure initialization is complete before any operation.
-   */
-  async #ensureInitialized() {
-    await this.#initialized;
+  async initialize() {
+    const ddl = storeDDL(this.#schema);
+    await this.#pool.query(ddl);
+    this.#isInitialized = true;
+  }
+  #ensureInitialized() {
+    if (!this.#isInitialized) {
+      throw new Error(
+        "PostgresContextStore not initialized. Call await store.initialize() after construction."
+      );
+    }
   }
   /**
    * Execute a function within a transaction.
    * Automatically commits on success or rolls back on error.
    */
   async #useTransaction(fn) {
-    await this.#ensureInitialized();
+    this.#ensureInitialized();
     const client = await this.#pool.connect();
     try {
       await client.query("BEGIN");
@@ -3971,7 +4050,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
    * Execute a query using the pool (no transaction).
    */
   async #query(sql, params) {
-    await this.#ensureInitialized();
+    this.#ensureInitialized();
     const result = await this.#pool.query(sql, params);
     return result.rows;
   }
@@ -3980,7 +4059,9 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
    * Call this when done with the store.
    */
   async close() {
-    await this.#pool.end();
+    if (this.#ownsPool) {
+      await this.#pool.end();
+    }
   }
   // ==========================================================================
   // Chat Operations
@@ -3988,7 +4069,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   async createChat(chat) {
     return this.#useTransaction(async (client) => {
       const result = await client.query(
-        `INSERT INTO chats (id, userId, title, metadata)
+        `INSERT INTO ${this.#t("chats")} (id, userId, title, metadata)
          VALUES ($1, $2, $3, $4)
          RETURNING *`,
         [
@@ -4000,7 +4081,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
       );
       const row = result.rows[0];
       await client.query(
-        `INSERT INTO branches (id, chatId, name, headMessageId, isActive, createdAt)
+        `INSERT INTO ${this.#t("branches")} (id, chatId, name, headMessageId, isActive, createdAt)
          VALUES ($1, $2, 'main', NULL, TRUE, $3)`,
         [crypto.randomUUID(), chat.id, Date.now()]
       );
@@ -4017,7 +4098,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   async upsertChat(chat) {
     return this.#useTransaction(async (client) => {
       const result = await client.query(
-        `INSERT INTO chats (id, userId, title, metadata)
+        `INSERT INTO ${this.#t("chats")} (id, userId, title, metadata)
          VALUES ($1, $2, $3, $4)
          ON CONFLICT(id) DO UPDATE SET id = EXCLUDED.id
          RETURNING *`,
@@ -4030,7 +4111,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
       );
       const row = result.rows[0];
       await client.query(
-        `INSERT INTO branches (id, chatId, name, headMessageId, isActive, createdAt)
+        `INSERT INTO ${this.#t("branches")} (id, chatId, name, headMessageId, isActive, createdAt)
          VALUES ($1, $2, 'main', NULL, TRUE, $3)
          ON CONFLICT(chatId, name) DO NOTHING`,
         [crypto.randomUUID(), chat.id, Date.now()]
@@ -4046,7 +4127,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
     });
   }
   async getChat(chatId) {
-    const rows = await this.#query("SELECT * FROM chats WHERE id = $1", [chatId]);
+    const rows = await this.#query(`SELECT * FROM ${this.#t("chats")} WHERE id = $1`, [chatId]);
     if (rows.length === 0) {
       return void 0;
     }
@@ -4076,7 +4157,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
     }
     params.push(chatId);
     const rows = await this.#query(
-      `UPDATE chats SET ${setClauses.join(", ")} WHERE id = $${paramIndex} RETURNING *`,
+      `UPDATE ${this.#t("chats")} SET ${setClauses.join(", ")} WHERE id = $${paramIndex} RETURNING *`,
       params
     );
     const row = rows[0];
@@ -4124,9 +4205,9 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
         c.updatedAt,
         COUNT(DISTINCT m.id) as messageCount,
         COUNT(DISTINCT b.id) as branchCount
-      FROM chats c
-      LEFT JOIN messages m ON m.chatId = c.id
-      LEFT JOIN branches b ON b.chatId = c.id
+      FROM ${this.#t("chats")} c
+      LEFT JOIN ${this.#t("messages")} m ON m.chatId = c.id
+      LEFT JOIN ${this.#t("branches")} b ON b.chatId = c.id
       ${whereClause}
       GROUP BY c.id
       ORDER BY c.updatedAt DESC${limitClause}`,
@@ -4145,7 +4226,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   }
   async deleteChat(chatId, options) {
     return this.#useTransaction(async (client) => {
-      let sql = "DELETE FROM chats WHERE id = $1";
+      let sql = `DELETE FROM ${this.#t("chats")} WHERE id = $1`;
       const params = [chatId];
       if (options?.userId !== void 0) {
         sql += " AND userId = $2";
@@ -4164,7 +4245,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
     }
     await this.#useTransaction(async (client) => {
       await client.query(
-        `INSERT INTO messages (id, chatId, parentId, name, type, data, createdAt)
+        `INSERT INTO ${this.#t("messages")} (id, chatId, parentId, name, type, data, createdAt)
          VALUES ($1, $2, $3, $4, $5, $6, $7)
          ON CONFLICT(id) DO UPDATE SET
            name = EXCLUDED.name,
@@ -4182,7 +4263,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
       );
       const content = typeof message2.data === "string" ? message2.data : JSON.stringify(message2.data);
       await client.query(
-        `INSERT INTO messages_fts (messageId, chatId, name, content)
+        `INSERT INTO ${this.#t("messages_fts")} (messageId, chatId, name, content)
          VALUES ($1, $2, $3, $4)
          ON CONFLICT(messageId) DO UPDATE SET
            chatId = EXCLUDED.chatId,
@@ -4193,7 +4274,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
     });
   }
   async getMessage(messageId) {
-    const rows = await this.#query("SELECT * FROM messages WHERE id = $1", [messageId]);
+    const rows = await this.#query(`SELECT * FROM ${this.#t("messages")} WHERE id = $1`, [messageId]);
     if (rows.length === 0) {
       return void 0;
     }
@@ -4211,9 +4292,9 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   async getMessageChain(headId) {
     const rows = await this.#query(
       `WITH RECURSIVE chain AS (
-        SELECT *, 0 as depth FROM messages WHERE id = $1
+        SELECT *, 0 as depth FROM ${this.#t("messages")} WHERE id = $1
         UNION ALL
-        SELECT m.*, c.depth + 1 FROM messages m
+        SELECT m.*, c.depth + 1 FROM ${this.#t("messages")} m
         INNER JOIN chain c ON m.id = c.parentId
         WHERE c.depth < 10000
       )
@@ -4233,7 +4314,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   }
   async hasChildren(messageId) {
     const rows = await this.#query(
-      "SELECT EXISTS(SELECT 1 FROM messages WHERE parentId = $1) as exists",
+      `SELECT EXISTS(SELECT 1 FROM ${this.#t("messages")} WHERE parentId = $1) as exists`,
       [messageId]
     );
     return rows[0].exists;
@@ -4254,7 +4335,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   // ==========================================================================
   async createBranch(branch) {
     await this.#query(
-      `INSERT INTO branches (id, chatId, name, headMessageId, isActive, createdAt)
+      `INSERT INTO ${this.#t("branches")} (id, chatId, name, headMessageId, isActive, createdAt)
        VALUES ($1, $2, $3, $4, $5, $6)`,
       [
         branch.id,
@@ -4267,7 +4348,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
     );
   }
   async getBranch(chatId, name) {
-    const rows = await this.#query("SELECT * FROM branches WHERE chatId = $1 AND name = $2", [
+    const rows = await this.#query(`SELECT * FROM ${this.#t("branches")} WHERE chatId = $1 AND name = $2`, [
       chatId,
       name
     ]);
@@ -4285,9 +4366,10 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
     };
   }
   async getActiveBranch(chatId) {
-    const rows = await this.#query("SELECT * FROM branches WHERE chatId = $1 AND isActive = TRUE", [
-      chatId
-    ]);
+    const rows = await this.#query(
+      `SELECT * FROM ${this.#t("branches")} WHERE chatId = $1 AND isActive = TRUE`,
+      [chatId]
+    );
     if (rows.length === 0) {
       return void 0;
     }
@@ -4304,19 +4386,20 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   async setActiveBranch(chatId, branchId) {
     await this.#useTransaction(async (client) => {
       await client.query(
-        "UPDATE branches SET isActive = FALSE WHERE chatId = $1",
+        `UPDATE ${this.#t("branches")} SET isActive = FALSE WHERE chatId = $1`,
         [chatId]
       );
-      await client.query("UPDATE branches SET isActive = TRUE WHERE id = $1", [
-        branchId
-      ]);
+      await client.query(
+        `UPDATE ${this.#t("branches")} SET isActive = TRUE WHERE id = $1`,
+        [branchId]
+      );
     });
   }
   async updateBranchHead(branchId, messageId) {
-    await this.#query("UPDATE branches SET headMessageId = $1 WHERE id = $2", [
-      messageId,
-      branchId
-    ]);
+    await this.#query(
+      `UPDATE ${this.#t("branches")} SET headMessageId = $1 WHERE id = $2`,
+      [messageId, branchId]
+    );
   }
   async listBranches(chatId) {
     const branches = await this.#query(
@@ -4326,7 +4409,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
         headMessageId,
         isActive,
         createdAt
-      FROM branches
+      FROM ${this.#t("branches")}
       WHERE chatId = $1
       ORDER BY createdAt ASC`,
       [chatId]
@@ -4337,9 +4420,9 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
       if (branch.headmessageid) {
         const countRows = await this.#query(
           `WITH RECURSIVE chain AS (
-            SELECT id, parentId FROM messages WHERE id = $1
+            SELECT id, parentId FROM ${this.#t("messages")} WHERE id = $1
             UNION ALL
-            SELECT m.id, m.parentId FROM messages m
+            SELECT m.id, m.parentId FROM ${this.#t("messages")} m
             INNER JOIN chain c ON m.id = c.parentId
           )
           SELECT COUNT(*) as count FROM chain`,
@@ -4363,7 +4446,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   // ==========================================================================
   async createCheckpoint(checkpoint) {
     await this.#query(
-      `INSERT INTO checkpoints (id, chatId, name, messageId, createdAt)
+      `INSERT INTO ${this.#t("checkpoints")} (id, chatId, name, messageId, createdAt)
        VALUES ($1, $2, $3, $4, $5)
        ON CONFLICT(chatId, name) DO UPDATE SET
          messageId = EXCLUDED.messageId,
@@ -4378,10 +4461,10 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
     );
   }
   async getCheckpoint(chatId, name) {
-    const rows = await this.#query("SELECT * FROM checkpoints WHERE chatId = $1 AND name = $2", [
-      chatId,
-      name
-    ]);
+    const rows = await this.#query(
+      `SELECT * FROM ${this.#t("checkpoints")} WHERE chatId = $1 AND name = $2`,
+      [chatId, name]
+    );
     if (rows.length === 0) {
       return void 0;
     }
@@ -4397,7 +4480,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   async listCheckpoints(chatId) {
     const rows = await this.#query(
       `SELECT id, name, messageId, createdAt
-       FROM checkpoints
+       FROM ${this.#t("checkpoints")}
        WHERE chatId = $1
        ORDER BY createdAt DESC`,
       [chatId]
@@ -4411,7 +4494,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   }
   async deleteCheckpoint(chatId, name) {
     await this.#query(
-      "DELETE FROM checkpoints WHERE chatId = $1 AND name = $2",
+      `DELETE FROM ${this.#t("checkpoints")} WHERE chatId = $1 AND name = $2`,
       [chatId, name]
     );
   }
@@ -4433,8 +4516,8 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
         ts_rank(fts.content_vector, plainto_tsquery('english', $2)) as rank,
         ts_headline('english', fts.content, plainto_tsquery('english', $2),
           'StartSel=<mark>, StopSel=</mark>, MaxWords=32, MinWords=5, MaxFragments=1') as snippet
-      FROM messages_fts fts
-      JOIN messages m ON m.id = fts.messageId
+      FROM ${this.#t("messages_fts")} fts
+      JOIN ${this.#t("messages")} m ON m.id = fts.messageId
       WHERE fts.content_vector @@ plainto_tsquery('english', $2)
         AND fts.chatId = $1
     `;
@@ -4468,7 +4551,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
   async getGraph(chatId) {
     const messageRows = await this.#query(
       `SELECT id, parentId, name, data, createdAt
-       FROM messages
+       FROM ${this.#t("messages")}
        WHERE chatId = $1
        ORDER BY createdAt ASC`,
       [chatId]
@@ -4486,7 +4569,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
     });
     const branchRows = await this.#query(
       `SELECT name, headMessageId, isActive
-       FROM branches
+       FROM ${this.#t("branches")}
        WHERE chatId = $1
        ORDER BY createdAt ASC`,
       [chatId]
@@ -4498,7 +4581,7 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
     }));
     const checkpointRows = await this.#query(
       `SELECT name, messageId
-       FROM checkpoints
+       FROM ${this.#t("checkpoints")}
        WHERE chatId = $1
        ORDER BY createdAt ASC`,
       [chatId]
@@ -4520,12 +4603,11 @@ var PostgresContextStore = class _PostgresContextStore extends ContextStore {
 import { createRequire as createRequire2 } from "node:module";
 
 // packages/context/src/lib/store/ddl.sqlserver.ts
-function storeDDL(schema) {
-  const s = schema;
+function storeDDL2(schema) {
   return `
-IF OBJECT_ID('[${s}].[chats]', 'U') IS NULL
+IF OBJECT_ID('[${schema}].[chats]', 'U') IS NULL
 BEGIN
-  CREATE TABLE [${s}].[chats] (
+  CREATE TABLE [${schema}].[chats] (
     id NVARCHAR(255) PRIMARY KEY,
     userId NVARCHAR(255) NOT NULL,
     title NVARCHAR(MAX),
@@ -4535,15 +4617,15 @@ BEGIN
   );
 END;
 
-IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${s}_chats_updatedAt' AND object_id = OBJECT_ID('[${s}].[chats]'))
-  CREATE INDEX [idx_${s}_chats_updatedAt] ON [${s}].[chats](updatedAt);
+IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${schema}_chats_updatedAt' AND object_id = OBJECT_ID('[${schema}].[chats]'))
+  CREATE INDEX [idx_${schema}_chats_updatedAt] ON [${schema}].[chats](updatedAt);
 
-IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${s}_chats_userId' AND object_id = OBJECT_ID('[${s}].[chats]'))
-  CREATE INDEX [idx_${s}_chats_userId] ON [${s}].[chats](userId);
+IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${schema}_chats_userId' AND object_id = OBJECT_ID('[${schema}].[chats]'))
+  CREATE INDEX [idx_${schema}_chats_userId] ON [${schema}].[chats](userId);
 
-IF OBJECT_ID('[${s}].[messages]', 'U') IS NULL
+IF OBJECT_ID('[${schema}].[messages]', 'U') IS NULL
 BEGIN
-  CREATE TABLE [${s}].[messages] (
+  CREATE TABLE [${schema}].[messages] (
     id NVARCHAR(255) PRIMARY KEY,
     chatId NVARCHAR(255) NOT NULL,
     parentId NVARCHAR(255),
@@ -4551,85 +4633,85 @@ BEGIN
     type NVARCHAR(255),
     data NVARCHAR(MAX) NOT NULL,
     createdAt BIGINT NOT NULL,
-    FOREIGN KEY (chatId) REFERENCES [${s}].[chats](id) ON DELETE CASCADE,
-    FOREIGN KEY (parentId) REFERENCES [${s}].[messages](id)
+    FOREIGN KEY (chatId) REFERENCES [${schema}].[chats](id) ON DELETE CASCADE,
+    FOREIGN KEY (parentId) REFERENCES [${schema}].[messages](id)
   );
 END;
 
-IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${s}_messages_chatId' AND object_id = OBJECT_ID('[${s}].[messages]'))
-  CREATE INDEX [idx_${s}_messages_chatId] ON [${s}].[messages](chatId);
+IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${schema}_messages_chatId' AND object_id = OBJECT_ID('[${schema}].[messages]'))
+  CREATE INDEX [idx_${schema}_messages_chatId] ON [${schema}].[messages](chatId);
 
-IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${s}_messages_parentId' AND object_id = OBJECT_ID('[${s}].[messages]'))
-  CREATE INDEX [idx_${s}_messages_parentId] ON [${s}].[messages](parentId);
+IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${schema}_messages_parentId' AND object_id = OBJECT_ID('[${schema}].[messages]'))
+  CREATE INDEX [idx_${schema}_messages_parentId] ON [${schema}].[messages](parentId);
 
-IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${s}_messages_chatId_parentId' AND object_id = OBJECT_ID('[${s}].[messages]'))
-  CREATE INDEX [idx_${s}_messages_chatId_parentId] ON [${s}].[messages](chatId, parentId);
+IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${schema}_messages_chatId_parentId' AND object_id = OBJECT_ID('[${schema}].[messages]'))
+  CREATE INDEX [idx_${schema}_messages_chatId_parentId] ON [${schema}].[messages](chatId, parentId);
 
-IF OBJECT_ID('[${s}].[branches]', 'U') IS NULL
+IF OBJECT_ID('[${schema}].[branches]', 'U') IS NULL
 BEGIN
-  CREATE TABLE [${s}].[branches] (
+  CREATE TABLE [${schema}].[branches] (
     id NVARCHAR(255) PRIMARY KEY,
     chatId NVARCHAR(255) NOT NULL,
     name NVARCHAR(255) NOT NULL,
     headMessageId NVARCHAR(255),
     isActive BIT NOT NULL DEFAULT 0,
     createdAt BIGINT NOT NULL,
-    FOREIGN KEY (chatId) REFERENCES [${s}].[chats](id) ON DELETE CASCADE,
-    FOREIGN KEY (headMessageId) REFERENCES [${s}].[messages](id),
-    CONSTRAINT [UQ_${s}_branches_chatId_name] UNIQUE(chatId, name)
+    FOREIGN KEY (chatId) REFERENCES [${schema}].[chats](id) ON DELETE CASCADE,
+    FOREIGN KEY (headMessageId) REFERENCES [${schema}].[messages](id),
+    CONSTRAINT [UQ_${schema}_branches_chatId_name] UNIQUE(chatId, name)
   );
 END;
 
-IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${s}_branches_chatId' AND object_id = OBJECT_ID('[${s}].[branches]'))
-  CREATE INDEX [idx_${s}_branches_chatId] ON [${s}].[branches](chatId);
+IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${schema}_branches_chatId' AND object_id = OBJECT_ID('[${schema}].[branches]'))
+  CREATE INDEX [idx_${schema}_branches_chatId] ON [${schema}].[branches](chatId);
 
-IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${s}_branches_chatId_isActive' AND object_id = OBJECT_ID('[${s}].[branches]'))
-  CREATE INDEX [idx_${s}_branches_chatId_isActive] ON [${s}].[branches](chatId, isActive);
+IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${schema}_branches_chatId_isActive' AND object_id = OBJECT_ID('[${schema}].[branches]'))
+  CREATE INDEX [idx_${schema}_branches_chatId_isActive] ON [${schema}].[branches](chatId, isActive);
 
-IF OBJECT_ID('[${s}].[checkpoints]', 'U') IS NULL
+IF OBJECT_ID('[${schema}].[checkpoints]', 'U') IS NULL
 BEGIN
-  CREATE TABLE [${s}].[checkpoints] (
+  CREATE TABLE [${schema}].[checkpoints] (
     id NVARCHAR(255) PRIMARY KEY,
     chatId NVARCHAR(255) NOT NULL,
     name NVARCHAR(255) NOT NULL,
     messageId NVARCHAR(255) NOT NULL,
     createdAt BIGINT NOT NULL,
-    FOREIGN KEY (chatId) REFERENCES [${s}].[chats](id) ON DELETE CASCADE,
-    FOREIGN KEY (messageId) REFERENCES [${s}].[messages](id),
-    CONSTRAINT [UQ_${s}_checkpoints_chatId_name] UNIQUE(chatId, name)
+    FOREIGN KEY (chatId) REFERENCES [${schema}].[chats](id) ON DELETE CASCADE,
+    FOREIGN KEY (messageId) REFERENCES [${schema}].[messages](id),
+    CONSTRAINT [UQ_${schema}_checkpoints_chatId_name] UNIQUE(chatId, name)
   );
 END;
 
-IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${s}_checkpoints_chatId' AND object_id = OBJECT_ID('[${s}].[checkpoints]'))
-  CREATE INDEX [idx_${s}_checkpoints_chatId] ON [${s}].[checkpoints](chatId);
+IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${schema}_checkpoints_chatId' AND object_id = OBJECT_ID('[${schema}].[checkpoints]'))
+  CREATE INDEX [idx_${schema}_checkpoints_chatId] ON [${schema}].[checkpoints](chatId);
 
-IF OBJECT_ID('[${s}].[messages_fts]', 'U') IS NULL
+IF OBJECT_ID('[${schema}].[messages_fts]', 'U') IS NULL
 BEGIN
-  CREATE TABLE [${s}].[messages_fts] (
+  CREATE TABLE [${schema}].[messages_fts] (
     messageId NVARCHAR(255) NOT NULL,
     chatId NVARCHAR(255) NOT NULL,
     name NVARCHAR(255) NOT NULL,
     content NVARCHAR(MAX) NOT NULL,
-    CONSTRAINT [PK_${s}_messages_fts] PRIMARY KEY (messageId),
-    FOREIGN KEY (messageId) REFERENCES [${s}].[messages](id) ON DELETE CASCADE
+    CONSTRAINT [PK_${schema}_messages_fts] PRIMARY KEY (messageId),
+    FOREIGN KEY (messageId) REFERENCES [${schema}].[messages](id) ON DELETE CASCADE
   );
 END;
 
-IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${s}_messages_fts_chatId' AND object_id = OBJECT_ID('[${s}].[messages_fts]'))
-  CREATE INDEX [idx_${s}_messages_fts_chatId] ON [${s}].[messages_fts](chatId);
+IF NOT EXISTS (SELECT * FROM sys.indexes WHERE name = 'idx_${schema}_messages_fts_chatId' AND object_id = OBJECT_ID('[${schema}].[messages_fts]'))
+  CREATE INDEX [idx_${schema}_messages_fts_chatId] ON [${schema}].[messages_fts](chatId);
 
 GO
 
 IF SERVERPROPERTY('IsFullTextInstalled') = 1
 BEGIN
-  IF NOT EXISTS (SELECT * FROM sys.fulltext_catalogs WHERE name = '${s}_context_store_catalog')
-    CREATE FULLTEXT CATALOG [${s}_context_store_catalog];
+  IF NOT EXISTS (SELECT * FROM sys.fulltext_catalogs WHERE name = '${schema}_context_store_catalog')
+    CREATE FULLTEXT CATALOG [${schema}_context_store_catalog];
 
-  IF NOT EXISTS (SELECT * FROM sys.fulltext_indexes WHERE object_id = OBJECT_ID('[${s}].[messages_fts]'))
+  IF NOT EXISTS (SELECT * FROM sys.fulltext_indexes WHERE object_id = OBJECT_ID('[${schema}].[messages_fts]'))
   BEGIN
-    CREATE FULLTEXT INDEX ON [${s}].[messages_fts](content)
-      KEY INDEX [PK_${s}_messages_fts]
-      ON [${s}_context_store_catalog]
+    CREATE FULLTEXT INDEX ON [${schema}].[messages_fts](content)
+      KEY INDEX [PK_${schema}_messages_fts]
+      ON [${schema}_context_store_catalog]
       WITH STOPLIST = SYSTEM;
   END;
 END;
@@ -4684,7 +4766,7 @@ var SqlServerContextStore = class _SqlServerContextStore extends ContextStore {
         EXEC sp_executesql @sql;
       END
     `);
-    const ddl = storeDDL(this.#schema);
+    const ddl = storeDDL2(this.#schema);
     const batches = ddl.split(/\bGO\b/i).filter((b) => b.trim());
     for (const batch of batches) {
       if (batch.trim()) {
@@ -5469,6 +5551,43 @@ var SqliteStreamStore = class extends StreamStore {
       stream.error
     );
   }
+  async upsertStream(stream) {
+    const row = this.#stmt(
+      `INSERT INTO streams (id, status, createdAt, startedAt, finishedAt, cancelRequestedAt, error)
+       VALUES (?, ?, ?, ?, ?, ?, ?)
+       ON CONFLICT(id) DO NOTHING
+       RETURNING *`
+    ).get(
+      stream.id,
+      stream.status,
+      stream.createdAt,
+      stream.startedAt,
+      stream.finishedAt,
+      stream.cancelRequestedAt,
+      stream.error
+    );
+    if (row) {
+      return {
+        stream: {
+          id: row.id,
+          status: row.status,
+          createdAt: row.createdAt,
+          startedAt: row.startedAt,
+          finishedAt: row.finishedAt,
+          cancelRequestedAt: row.cancelRequestedAt,
+          error: row.error
+        },
+        created: true
+      };
+    }
+    const existing = await this.getStream(stream.id);
+    if (!existing) {
+      throw new Error(
+        `Stream "${stream.id}" disappeared between upsert and fetch`
+      );
+    }
+    return { stream: existing, created: false };
+  }
   async getStream(streamId) {
     const row = this.#stmt("SELECT * FROM streams WHERE id = ?").get(
       streamId
@@ -5575,7 +5694,7 @@ var StreamManager = class {
     return this.#store;
   }
   async register(streamId) {
-    await this.#store.createStream({
+    return this.#store.upsertStream({
       id: streamId,
       status: "queued",
       createdAt: Date.now(),