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

package/dist/index.js

@@ -12,6 +12,8 @@ var agentThread = sqliteTable("agent_thread", {
     mode: "boolean"
   }).notNull().default(false),
   activeStreamId: text("active_stream_id"),
+  /** The agent a new turn on this thread defaults to when the caller names none. */
+  defaultAgent: text("default_agent"),
   createdAt: integer("created_at", {
     mode: "timestamp_ms"
   }).notNull(),
@@ -37,6 +39,9 @@ var agentMessage = sqliteTable("agent_message", {
   toolResults: text("tool_results", {
     mode: "json"
   }).$type(),
+  attachments: text("attachments", {
+    mode: "json"
+  }).$type(),
   followUps: text("follow_ups", {
     mode: "json"
   }).$type(),
@@ -44,6 +49,7 @@ var agentMessage = sqliteTable("agent_message", {
     mode: "json"
   }).$type(),
   agentName: text("agent_name"),
+  runId: text("run_id"),
   createdAt: integer("created_at", {
     mode: "timestamp_ms"
   }).notNull()
@@ -75,7 +81,9 @@ var agentToolCall = sqliteTable("agent_tool_call", {
   }),
   /** The run (turn) this call belongs to, for a trace deep-link; null for a pre-rollout row. */
   runId: text("run_id")
-});
+}, (table) => [
+  index("agent_tool_call_message_idx").on(table.messageId)
+]);
 var agentTokenUsage = sqliteTable("agent_token_usage", {
   id: text("id").primaryKey(),
   threadId: text("thread_id").notNull().references(() => agentThread.id, {
@@ -134,7 +142,15 @@ var agentRun = sqliteTable("agent_run", {
     mode: "timestamp_ms"
   }),
   /** sha256 hex of the run's resolved (pre-RAG) system prompt; null for a pre-existing run. */
-  promptHash: text("prompt_hash")
+  promptHash: 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: text("parent_run_id")
 }, (table) => [
   index("agent_run_started_idx").on(table.startedAt)
 ]);
@@ -157,6 +173,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
@@ -170,9 +187,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
@@ -192,6 +211,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,
@@ -230,20 +251,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(sql.raw(statement));
   }
+  for (const { table, column, ddl } of additiveColumns) {
+    const existing = await db.all(sql.raw(`PRAGMA table_info(${table})`));
+    if (!existing.some((row) => row.name === column)) {
+      await db.run(sql.raw(ddl));
+    }
+  }
 }
 __name(ensureAgentSchema, "ensureAgentSchema");
 
 // src/drizzle-agent-store.ts
-import { and, asc, desc, eq, gte, inArray, isNull, lte, sql as sql2 } from "drizzle-orm";
+import { and, asc, desc, eq, gte, inArray, isNotNull, isNull, lte, sql as sql2 } from "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");
@@ -261,6 +321,7 @@ var DrizzleAgentStore = class {
       title: input.title ?? "New chat",
       transient: input.transient ?? false,
       activeStreamId: null,
+      defaultAgent: null,
       createdAt: now,
       updatedAt: now,
       deletedAt: null
@@ -283,6 +344,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(and(eq(agentThread.id, query.threadId), 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(and(eq(agentMessage.threadId, query.threadId), 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(eq(agentMessage.threadId, threadId)).orderBy(desc(agentMessage.createdAt), 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(and(eq(agentThread.actorRef, actorRef), eq(agentThread.transient, false), isNull(agentThread.deletedAt))).orderBy(desc(agentThread.updatedAt)).limit(limit);
     const summaries = [];
@@ -313,6 +418,7 @@ var DrizzleAgentStore = class {
       title: source.title,
       transient: false,
       activeStreamId: null,
+      defaultAgent: source.defaultAgent,
       createdAt: now,
       updatedAt: now,
       deletedAt: null
@@ -326,9 +432,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
       })));
     }
@@ -348,9 +459,10 @@ var DrizzleAgentStore = class {
   }
   async runForToolCall(toolCallId) {
     const [row] = await this.db.select({
+      runId: agentToolCall.runId,
       activeStreamId: agentThread.activeStreamId
     }).from(agentToolCall).innerJoin(agentMessage, eq(agentToolCall.messageId, agentMessage.id)).innerJoin(agentThread, eq(agentMessage.threadId, agentThread.id)).where(eq(agentToolCall.id, toolCallId));
-    return row?.activeStreamId ?? null;
+    return row?.runId ?? row?.activeStreamId ?? null;
   }
   async ownerOfActiveStream(runId) {
     const [thread] = await this.db.select({
@@ -375,7 +487,43 @@ var DrizzleAgentStore = class {
       activeStreamId: runId
     }).where(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(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(and(eq(agentThread.id, threadId), 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,
@@ -389,7 +537,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);
   }
@@ -429,9 +578,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);
@@ -440,6 +591,11 @@ var DrizzleAgentStore = class {
     }).where(eq(agentThread.id, input.threadId));
     return this.toStoredMessage(message);
   }
+  async setMessageToolResults(messageId, results) {
+    await this.db.update(agentMessage).set({
+      toolResults: results
+    }).where(eq(agentMessage.id, messageId));
+  }
   async truncateFrom(threadId, messageId) {
     const messages = await this.db.select().from(agentMessage).where(eq(agentMessage.threadId, threadId)).orderBy(asc(agentMessage.createdAt), asc(agentMessage.id));
     const cutoff = messages.findIndex((message) => message.id === messageId);
@@ -488,6 +644,39 @@ var DrizzleAgentStore = class {
     }
     await this.db.update(agentToolCall).set(updates).where(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, eq(agentMessage.threadId, agentThread.id)).where(and(eq(agentThread.actorRef, actorRef), 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(),
@@ -522,6 +711,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)
       } : {}
@@ -539,6 +729,9 @@ var DrizzleAgentStore = class {
       ...message.toolResults != null ? {
         toolResults: message.toolResults
       } : {},
+      ...message.attachments != null ? {
+        attachments: message.attachments
+      } : {},
       ...message.followUps != null ? {
         followUps: message.followUps
       } : {},
@@ -547,6 +740,9 @@ var DrizzleAgentStore = class {
       } : {},
       ...message.agentName != null ? {
         agentName: message.agentName
+      } : {},
+      ...message.runId != null ? {
+        runId: message.runId
       } : {}
     };
   }
@@ -580,7 +776,8 @@ __name(isToolKind, "isToolKind");
 var RUN_STATUSES = [
   "running",
   "completed",
-  "failed"
+  "failed",
+  "cancelled"
 ];
 function isRunStatus(value) {
   return RUN_STATUSES.includes(value);
@@ -621,7 +818,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");