@dudousxd/nestjs-agent-store-drizzle 0.8.1 → 0.9.0

package/README.md

@@ -36,9 +36,58 @@ export class AppModule {}
 ```
 
 The package ships the `agentSchema` (Drizzle tables), `ensureAgentSchema` (a non-destructive
-`CREATE TABLE IF NOT EXISTS` helper for a quick start), `DrizzleAgentStore`, and
-`DrizzleGovernanceQueries`. For production, prefer your normal drizzle-kit migrations over the
-`ensureAgentSchema` helper.
+DDL helper for a quick start), `DrizzleAgentStore`, and `DrizzleGovernanceQueries`. For production,
+prefer your normal drizzle-kit migrations over the `ensureAgentSchema` helper.
+
+## The read a turn makes
+
+The agent loop does not call `getThread` to build a prompt. This store implements the core SPI's
+`ThreadTurnReader`, so the loop asks it for `loadThreadForTurn({ threadId, messageLimit })` instead:
+the thread's newest `messageLimit` messages oldest-first, bounded by the database rather than in
+memory, projected to the columns a model turn reads (`usage`, `follow_ups` and `run_id` stay in the
+table), plus the title, the default agent, and whether the thread has EVER been answered.
+
+`messageLimit` is the configured `HistoryPolicy.maxMessages`. A policy that summarizes, or whose
+ceiling is only a token budget, names no row bound and the whole thread is read — see the core
+README. Nothing to wire: the loop probes for the method, and `getThread` remains the right read for
+a client rendering a transcript.
+
+## Upgrading a database that already exists
+
+`ensureAgentSchema` is additive only — `CREATE TABLE IF NOT EXISTS` plus `CREATE INDEX IF NOT
+EXISTS`, and an explicit add-column pass for columns this package introduced after a table shipped,
+since the `CREATE TABLE` guard is inert against a table that already exists. Calling it on boot is
+enough; it never drops a column or changes a type.
+
+On your own drizzle-kit migrations, those additive statements are yours to write:
+
+```sql
+ALTER TABLE agent_thread ADD COLUMN default_agent TEXT;
+ALTER TABLE agent_message ADD COLUMN run_id TEXT;
+ALTER TABLE agent_message ADD COLUMN attachments TEXT;
+ALTER TABLE agent_run ADD COLUMN parent_run_id TEXT;
+CREATE INDEX agent_tool_call_message_idx ON agent_tool_call (message_id);
+```
+
+The index is not optional on Postgres: `agent_tool_call.message_id` carries a foreign key, and
+Postgres — unlike MySQL — does not index one for you, while both message-scoped reads (the thread
+reader's `IN (…)` and `truncateFrom`'s delete) filter on it.
+
+## Attachment housekeeping
+
+`referencedMediaIds(actorRef, mediaIds)` answers which of a set of media ids a message that still
+exists carries, for one actor — the inverse of the host's own staged-media inventory, and the half a
+sweep cannot work out for itself. It needs **no schema change**: it reads the `attachments` JSON
+column messages have carried since attachments shipped, so an existing database answers correctly
+the moment you upgrade, with nothing to backfill.
+
+The match happens in memory rather than in SQL. The column holds an array of objects, every dialect
+spells that query differently, and none of them can use an index for it; the scan is bounded to one
+actor's attachment-bearing messages, and a message with no attachments never leaves the database.
+
+A thread that was soft-deleted still counts as holding its references: the message rows survive, so
+the bytes are still reachable from stored state. Delete the thread for real and the cascade makes
+them collectable on the next sweep.
 
 ## Cost accounting
 

package/dist/index.cjs

@@ -47,6 +47,8 @@ var agentThread = (0, import_sqlite_core.sqliteTable)("agent_thread", {
     mode: "boolean"
   }).notNull().default(false),
   activeStreamId: (0, import_sqlite_core.text)("active_stream_id"),
+  /** The agent a new turn on this thread defaults to when the caller names none. */
+  defaultAgent: (0, import_sqlite_core.text)("default_agent"),
   createdAt: (0, import_sqlite_core.integer)("created_at", {
     mode: "timestamp_ms"
   }).notNull(),
@@ -72,6 +74,9 @@ var agentMessage = (0, import_sqlite_core.sqliteTable)("agent_message", {
   toolResults: (0, import_sqlite_core.text)("tool_results", {
     mode: "json"
   }).$type(),
+  attachments: (0, import_sqlite_core.text)("attachments", {
+    mode: "json"
+  }).$type(),
   followUps: (0, import_sqlite_core.text)("follow_ups", {
     mode: "json"
   }).$type(),
@@ -79,6 +84,7 @@ var agentMessage = (0, import_sqlite_core.sqliteTable)("agent_message", {
     mode: "json"
   }).$type(),
   agentName: (0, import_sqlite_core.text)("agent_name"),
+  runId: (0, import_sqlite_core.text)("run_id"),
   createdAt: (0, import_sqlite_core.integer)("created_at", {
     mode: "timestamp_ms"
   }).notNull()
@@ -110,7 +116,9 @@ var agentToolCall = (0, import_sqlite_core.sqliteTable)("agent_tool_call", {
   }),
   /** The run (turn) this call belongs to, for a trace deep-link; null for a pre-rollout row. */
   runId: (0, import_sqlite_core.text)("run_id")
-});
+}, (table) => [
+  (0, import_sqlite_core.index)("agent_tool_call_message_idx").on(table.messageId)
+]);
 var agentTokenUsage = (0, import_sqlite_core.sqliteTable)("agent_token_usage", {
   id: (0, import_sqlite_core.text)("id").primaryKey(),
   threadId: (0, import_sqlite_core.text)("thread_id").notNull().references(() => agentThread.id, {
@@ -169,7 +177,15 @@ var agentRun = (0, import_sqlite_core.sqliteTable)("agent_run", {
     mode: "timestamp_ms"
   }),
   /** sha256 hex of the run's resolved (pre-RAG) system prompt; null for a pre-existing run. */
-  promptHash: (0, import_sqlite_core.text)("prompt_hash")
+  promptHash: (0, import_sqlite_core.text)("prompt_hash"),
+  /**
+   * The run that delegated this one; null for a turn nobody delegated. The durable runtime's
+   * journal holds the same edge, but only there — this column is what lets a reader of run ROWS
+   * build the delegation tree and roll a child's cost up to the turn that asked for it, including
+   * for a detached child, which outlives its parent's turn and so is paired by nothing in the
+   * transcript.
+   */
+  parentRunId: (0, import_sqlite_core.text)("parent_run_id")
 }, (table) => [
   (0, import_sqlite_core.index)("agent_run_started_idx").on(table.startedAt)
 ]);
@@ -192,6 +208,7 @@ var statements = [
     title TEXT NOT NULL,
     transient INTEGER NOT NULL DEFAULT 0,
     active_stream_id TEXT,
+    default_agent TEXT,
     created_at INTEGER NOT NULL,
     updated_at INTEGER NOT NULL,
     deleted_at INTEGER
@@ -205,9 +222,11 @@ var statements = [
     content TEXT NOT NULL,
     tool_calls TEXT,
     tool_results TEXT,
+    attachments TEXT,
     follow_ups TEXT,
     usage TEXT,
     agent_name TEXT,
+    run_id TEXT,
     created_at INTEGER NOT NULL
   )`,
   `CREATE INDEX IF NOT EXISTS agent_message_thread_created_idx
@@ -227,6 +246,8 @@ var statements = [
     executed_at INTEGER,
     run_id TEXT
   )`,
+  `CREATE INDEX IF NOT EXISTS agent_tool_call_message_idx
+    ON agent_tool_call (message_id)`,
   `CREATE TABLE IF NOT EXISTS agent_token_usage (
     id TEXT PRIMARY KEY NOT NULL,
     thread_id TEXT NOT NULL REFERENCES agent_thread(id) ON DELETE CASCADE,
@@ -265,20 +286,59 @@ var statements = [
     retries INTEGER NOT NULL DEFAULT 0,
     started_at INTEGER NOT NULL,
     settled_at INTEGER,
-    prompt_hash TEXT
+    prompt_hash TEXT,
+    parent_run_id TEXT
   )`,
   `CREATE INDEX IF NOT EXISTS agent_run_started_idx
     ON agent_run (started_at)`
 ];
+var additiveColumns = [
+  {
+    table: "agent_message",
+    column: "run_id",
+    ddl: "ALTER TABLE agent_message ADD COLUMN run_id TEXT"
+  },
+  {
+    table: "agent_message",
+    column: "attachments",
+    ddl: "ALTER TABLE agent_message ADD COLUMN attachments TEXT"
+  },
+  {
+    table: "agent_thread",
+    column: "default_agent",
+    ddl: "ALTER TABLE agent_thread ADD COLUMN default_agent TEXT"
+  },
+  {
+    table: "agent_run",
+    column: "parent_run_id",
+    ddl: "ALTER TABLE agent_run ADD COLUMN parent_run_id TEXT"
+  }
+];
 async function ensureAgentSchema(db) {
   for (const statement of statements) {
     await db.run(import_drizzle_orm.sql.raw(statement));
   }
+  for (const { table, column, ddl } of additiveColumns) {
+    const existing = await db.all(import_drizzle_orm.sql.raw(`PRAGMA table_info(${table})`));
+    if (!existing.some((row) => row.name === column)) {
+      await db.run(import_drizzle_orm.sql.raw(ddl));
+    }
+  }
 }
 __name(ensureAgentSchema, "ensureAgentSchema");
 
 // src/drizzle-agent-store.ts
 var import_drizzle_orm2 = require("drizzle-orm");
+var TURN_MESSAGE_COLUMNS = {
+  id: agentMessage.id,
+  role: agentMessage.role,
+  content: agentMessage.content,
+  agentName: agentMessage.agentName,
+  toolCalls: agentMessage.toolCalls,
+  toolResults: agentMessage.toolResults,
+  attachments: agentMessage.attachments,
+  createdAt: agentMessage.createdAt
+};
 var DrizzleAgentStore = class {
   static {
     __name(this, "DrizzleAgentStore");
@@ -296,6 +356,7 @@ var DrizzleAgentStore = class {
       title: input.title ?? "New chat",
       transient: input.transient ?? false,
       activeStreamId: null,
+      defaultAgent: null,
       createdAt: now,
       updatedAt: now,
       deletedAt: null
@@ -318,6 +379,50 @@ var DrizzleAgentStore = class {
       } : {}
     };
   }
+  /**
+  * The window a turn reads off a thread: its newest `messageLimit` messages (oldest first), the
+  * title, the default agent, and whether the thread has EVER been answered.
+  *
+  * {@link getThread} is the wrong read for a turn. It materializes the transcript — every message,
+  * every attachment, every tool output the thread ever recorded — and the run then journals what it
+  * loaded, so a long thread pays for its whole history on every turn and again on every replay.
+  *
+  * `hasAssistantMessage` is answered over the WHOLE thread, not the returned page: it answers "has
+  * this conversation been answered before?", and a thread whose window happens to hold only the
+  * user's last questions has still been answered. `null` when the thread is unknown or
+  * soft-deleted, same as {@link getThread}.
+  */
+  async loadThreadForTurn(query) {
+    const [thread] = await this.db.select({
+      title: agentThread.title,
+      defaultAgent: agentThread.defaultAgent
+    }).from(agentThread).where((0, import_drizzle_orm2.and)((0, import_drizzle_orm2.eq)(agentThread.id, query.threadId), (0, import_drizzle_orm2.isNull)(agentThread.deletedAt)));
+    if (thread === void 0) {
+      return null;
+    }
+    const messages = await this.loadTurnWindow(query.threadId, query.messageLimit);
+    const [answered] = await this.db.select({
+      id: agentMessage.id
+    }).from(agentMessage).where((0, import_drizzle_orm2.and)((0, import_drizzle_orm2.eq)(agentMessage.threadId, query.threadId), (0, import_drizzle_orm2.eq)(agentMessage.role, "assistant"))).limit(1);
+    return {
+      title: thread.title,
+      defaultAgent: thread.defaultAgent,
+      hasAssistantMessage: answered !== void 0,
+      messages
+    };
+  }
+  /**
+  * The newest `limit` messages of a thread, oldest first, projected to the columns a model turn
+  * reads. Read newest-first so the LIMIT is the database's job, then reversed for the prompt.
+  */
+  async loadTurnWindow(threadId, limit) {
+    if (limit !== void 0 && limit <= 0) {
+      return [];
+    }
+    const window = this.db.select(TURN_MESSAGE_COLUMNS).from(agentMessage).where((0, import_drizzle_orm2.eq)(agentMessage.threadId, threadId)).orderBy((0, import_drizzle_orm2.desc)(agentMessage.createdAt), (0, import_drizzle_orm2.desc)(agentMessage.id));
+    const rows = limit === void 0 ? await window : await window.limit(limit);
+    return rows.reverse().map((row) => this.toStoredMessage(row));
+  }
   async listThreads(actorRef, limit = 50) {
     const threads = await this.db.select().from(agentThread).where((0, import_drizzle_orm2.and)((0, import_drizzle_orm2.eq)(agentThread.actorRef, actorRef), (0, import_drizzle_orm2.eq)(agentThread.transient, false), (0, import_drizzle_orm2.isNull)(agentThread.deletedAt))).orderBy((0, import_drizzle_orm2.desc)(agentThread.updatedAt)).limit(limit);
     const summaries = [];
@@ -348,6 +453,7 @@ var DrizzleAgentStore = class {
       title: source.title,
       transient: false,
       activeStreamId: null,
+      defaultAgent: source.defaultAgent,
       createdAt: now,
       updatedAt: now,
       deletedAt: null
@@ -361,9 +467,14 @@ var DrizzleAgentStore = class {
         content: message.content,
         toolCalls: message.toolCalls,
         toolResults: message.toolResults,
+        attachments: message.attachments,
         followUps: message.followUps,
         usage: message.usage,
         agentName: message.agentName,
+        // The copy is the same message, so it keeps the run that wrote it. A reader asking which
+        // turn produced this text gets the truthful answer; the run's own thread is still the
+        // original, so a run-scoped read never picks the fork's rows up.
+        runId: message.runId,
         createdAt: message.createdAt
       })));
     }
@@ -383,9 +494,10 @@ var DrizzleAgentStore = class {
   }
   async runForToolCall(toolCallId) {
     const [row] = await this.db.select({
+      runId: agentToolCall.runId,
       activeStreamId: agentThread.activeStreamId
     }).from(agentToolCall).innerJoin(agentMessage, (0, import_drizzle_orm2.eq)(agentToolCall.messageId, agentMessage.id)).innerJoin(agentThread, (0, import_drizzle_orm2.eq)(agentMessage.threadId, agentThread.id)).where((0, import_drizzle_orm2.eq)(agentToolCall.id, toolCallId));
-    return row?.activeStreamId ?? null;
+    return row?.runId ?? row?.activeStreamId ?? null;
   }
   async ownerOfActiveStream(runId) {
     const [thread] = await this.db.select({
@@ -410,7 +522,43 @@ var DrizzleAgentStore = class {
       activeStreamId: runId
     }).where((0, import_drizzle_orm2.eq)(agentThread.id, threadId));
   }
-  /** Persist the start of a run (turn). Replay-safe: called under a durable localStep. */
+  /**
+  * Patch a thread's `title` and/or `defaultAgent`. Each field is only touched when PRESENT in
+  * `patch` — `defaultAgent: null` clears it (falls back to the app default), while an omitted
+  * `defaultAgent` leaves whatever is stored untouched. A no-op (including `updatedAt`) when the
+  * patch is empty or the thread doesn't exist.
+  */
+  async updateThread(threadId, patch) {
+    const updates = {};
+    if (patch.title !== void 0) {
+      updates.title = patch.title;
+    }
+    if (patch.defaultAgent !== void 0) {
+      updates.defaultAgent = patch.defaultAgent;
+    }
+    if (Object.keys(updates).length === 0) {
+      return;
+    }
+    updates.updatedAt = /* @__PURE__ */ new Date();
+    await this.db.update(agentThread).set(updates).where((0, import_drizzle_orm2.eq)(agentThread.id, threadId));
+  }
+  /**
+  * The thread's default agent, projected — the caller wants one nullable scalar to decide which
+  * agent answers the next turn, and {@link getThread} would materialize the whole transcript to
+  * hand it over. `null` when the thread is unknown, soft-deleted, or has no default set.
+  */
+  async defaultAgentForThread(threadId) {
+    const [thread] = await this.db.select({
+      defaultAgent: agentThread.defaultAgent
+    }).from(agentThread).where((0, import_drizzle_orm2.and)((0, import_drizzle_orm2.eq)(agentThread.id, threadId), (0, import_drizzle_orm2.isNull)(agentThread.deletedAt)));
+    return thread?.defaultAgent ?? null;
+  }
+  /**
+  * Persist the start of a run (turn). Replay-safe: called under a durable localStep.
+  *
+  * Takes the SPI's own {@link RecordRunStartInput} rather than a hand-copied shape: a field added
+  * to the input is otherwise accepted and dropped, silently, by every adapter that re-declares it.
+  */
   async recordRunStart(run) {
     const runRow = {
       id: run.runId,
@@ -424,7 +572,8 @@ var DrizzleAgentStore = class {
       retries: 0,
       startedAt: /* @__PURE__ */ new Date(),
       settledAt: null,
-      promptHash: run.promptHash ?? null
+      promptHash: run.promptHash ?? null,
+      parentRunId: run.parentRunId ?? null
     };
     await this.db.insert(agentRun).values(runRow);
   }
@@ -464,9 +613,11 @@ var DrizzleAgentStore = class {
       content: input.content,
       toolCalls: input.toolCalls ?? null,
       toolResults: input.toolResults ?? null,
+      attachments: input.attachments ?? null,
       followUps: input.followUps ?? null,
       usage: input.usage ?? null,
       agentName: input.agentName ?? null,
+      runId: input.runId ?? null,
       createdAt: now
     };
     await this.db.insert(agentMessage).values(message);
@@ -475,6 +626,11 @@ var DrizzleAgentStore = class {
     }).where((0, import_drizzle_orm2.eq)(agentThread.id, input.threadId));
     return this.toStoredMessage(message);
   }
+  async setMessageToolResults(messageId, results) {
+    await this.db.update(agentMessage).set({
+      toolResults: results
+    }).where((0, import_drizzle_orm2.eq)(agentMessage.id, messageId));
+  }
   async truncateFrom(threadId, messageId) {
     const messages = await this.db.select().from(agentMessage).where((0, import_drizzle_orm2.eq)(agentMessage.threadId, threadId)).orderBy((0, import_drizzle_orm2.asc)(agentMessage.createdAt), (0, import_drizzle_orm2.asc)(agentMessage.id));
     const cutoff = messages.findIndex((message) => message.id === messageId);
@@ -523,6 +679,39 @@ var DrizzleAgentStore = class {
     }
     await this.db.update(agentToolCall).set(updates).where((0, import_drizzle_orm2.eq)(agentToolCall.id, input.toolCallId));
   }
+  /**
+  * Of `mediaIds`, the ones a surviving message in one of this actor's threads still carries.
+  *
+  * Reads the `attachments` JSON back out and matches in memory rather than pushing the match into
+  * SQL: the column holds an array of objects, and every dialect this adapter's sibling targets
+  * spells that query differently (`jsonb` containment, `json_table`, `json_each`), while none of
+  * them can use an index for it anyway. The scan is bounded by ONE actor's attachment-bearing
+  * messages, which is a small set — attachments are rare, and a message without one is excluded
+  * by the `is not null` before any JSON is parsed.
+  *
+  * Soft-deleted threads are deliberately included: their message rows survive, so the bytes they
+  * point at are still reachable from stored state and are not garbage.
+  */
+  async referencedMediaIds(actorRef, mediaIds) {
+    if (mediaIds.length === 0) {
+      return [];
+    }
+    const rows = await this.db.select({
+      attachments: agentMessage.attachments
+    }).from(agentMessage).innerJoin(agentThread, (0, import_drizzle_orm2.eq)(agentMessage.threadId, agentThread.id)).where((0, import_drizzle_orm2.and)((0, import_drizzle_orm2.eq)(agentThread.actorRef, actorRef), (0, import_drizzle_orm2.isNotNull)(agentMessage.attachments)));
+    const wanted = new Set(mediaIds);
+    const found = /* @__PURE__ */ new Set();
+    for (const row of rows) {
+      for (const attachment of row.attachments ?? []) {
+        if (wanted.has(attachment.mediaId)) {
+          found.add(attachment.mediaId);
+        }
+      }
+    }
+    return [
+      ...wanted
+    ].filter((mediaId) => found.has(mediaId));
+  }
   async recordUsage(input) {
     await this.db.insert(agentTokenUsage).values({
       id: crypto.randomUUID(),
@@ -557,6 +746,7 @@ var DrizzleAgentStore = class {
       transient: thread.transient,
       createdAt: thread.createdAt.toISOString(),
       updatedAt: thread.updatedAt.toISOString(),
+      defaultAgent: thread.defaultAgent,
       ...lastContent !== void 0 ? {
         lastMessagePreview: lastContent.slice(0, 120)
       } : {}
@@ -574,6 +764,9 @@ var DrizzleAgentStore = class {
       ...message.toolResults != null ? {
         toolResults: message.toolResults
       } : {},
+      ...message.attachments != null ? {
+        attachments: message.attachments
+      } : {},
       ...message.followUps != null ? {
         followUps: message.followUps
       } : {},
@@ -582,6 +775,9 @@ var DrizzleAgentStore = class {
       } : {},
       ...message.agentName != null ? {
         agentName: message.agentName
+      } : {},
+      ...message.runId != null ? {
+        runId: message.runId
       } : {}
     };
   }
@@ -615,7 +811,8 @@ __name(isToolKind, "isToolKind");
 var RUN_STATUSES = [
   "running",
   "completed",
-  "failed"
+  "failed",
+  "cancelled"
 ];
 function isRunStatus(value) {
   return RUN_STATUSES.includes(value);
@@ -656,7 +853,8 @@ function toRecentRunRow(run) {
     errorMessage: run.errorMessage ?? null,
     retries: run.retries,
     startedAt: run.startedAt.toISOString(),
-    promptHash: run.promptHash ?? null
+    promptHash: run.promptHash ?? null,
+    parentRunId: run.parentRunId ?? null
   };
 }
 __name(toRecentRunRow, "toRecentRunRow");