@dudousxd/nestjs-agent-store-drizzle 0.32.0 → 0.33.0

package/dist/index.cjs

@@ -94,7 +94,9 @@ var agentThread = (0, import_sqlite_core.sqliteTable)("agent_thread", {
     mode: "timestamp_ms"
   })
 }, (table) => [
-  (0, import_sqlite_core.index)("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt)
+  (0, import_sqlite_core.index)("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt),
+  // `ownerOfActiveStream` reads by run id on every stream attach and reconnect.
+  (0, import_sqlite_core.index)("agent_thread_active_stream_idx").on(table.activeStreamId)
 ]);
 var agentMessage = (0, import_sqlite_core.sqliteTable)("agent_message", {
   id: (0, import_sqlite_core.text)("id").primaryKey(),
@@ -144,7 +146,9 @@ var agentMessage = (0, import_sqlite_core.sqliteTable)("agent_message", {
     mode: "timestamp_ms"
   }).notNull()
 }, (table) => [
-  (0, import_sqlite_core.index)("agent_message_thread_created_idx").on(table.threadId, table.createdAt)
+  (0, import_sqlite_core.index)("agent_message_thread_created_idx").on(table.threadId, table.createdAt),
+  // Transcript reads order a thread's messages by `seq`.
+  (0, import_sqlite_core.index)("agent_message_thread_seq_idx").on(table.threadId, table.seq)
 ]);
 var agentToolCall = (0, import_sqlite_core.sqliteTable)("agent_tool_call", {
   proposalId: (0, import_sqlite_core.text)("proposal_id"),
@@ -189,7 +193,11 @@ var agentToolCall = (0, import_sqlite_core.sqliteTable)("agent_tool_call", {
   /** The surface the decision came through (`web`, `slack`, `remembered`, …). */
   decidedVia: (0, import_sqlite_core.text)("decided_via")
 }, (table) => [
-  (0, import_sqlite_core.index)("agent_tool_call_message_idx").on(table.messageId)
+  (0, import_sqlite_core.index)("agent_tool_call_message_idx").on(table.messageId),
+  // The approvals inbox: `pending_approval`, oldest first.
+  (0, import_sqlite_core.index)("agent_tool_call_status_created_idx").on(table.status, table.createdAt),
+  (0, import_sqlite_core.index)("agent_tool_call_run_idx").on(table.runId),
+  (0, import_sqlite_core.index)("agent_tool_call_created_idx").on(table.createdAt)
 ]);
 var agentQueuedMessage = (0, import_sqlite_core.sqliteTable)("agent_queued_message", {
   uiCapabilities: (0, import_sqlite_core.text)("ui_capabilities", {
@@ -253,7 +261,9 @@ var agentTokenUsage = (0, import_sqlite_core.sqliteTable)("agent_token_usage", {
     mode: "timestamp_ms"
   }).notNull()
 }, (table) => [
-  (0, import_sqlite_core.index)("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt)
+  (0, import_sqlite_core.index)("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt),
+  (0, import_sqlite_core.index)("agent_token_usage_created_idx").on(table.createdAt),
+  (0, import_sqlite_core.index)("agent_token_usage_thread_idx").on(table.threadId)
 ]);
 var agentGenuiCall = (0, import_sqlite_core.sqliteTable)("agent_genui_call", {
   id: (0, import_sqlite_core.text)("id").primaryKey(),
@@ -596,7 +606,9 @@ var agentThread2 = (0, import_mysql_core.mysqlTable)("agent_thread", {
   updatedAt: ts("updated_at").notNull(),
   deletedAt: ts("deleted_at")
 }, (table) => [
-  (0, import_mysql_core.index)("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt)
+  (0, import_mysql_core.index)("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt),
+  // `ownerOfActiveStream` reads by run id on every stream attach and reconnect.
+  (0, import_mysql_core.index)("agent_thread_active_stream_idx").on(table.activeStreamId)
 ]);
 var agentMessage2 = (0, import_mysql_core.mysqlTable)("agent_message", {
   id: key("id").primaryKey(),
@@ -622,7 +634,9 @@ var agentMessage2 = (0, import_mysql_core.mysqlTable)("agent_message", {
   seq: (0, import_mysql_core.int)("seq"),
   createdAt: ts("created_at").notNull()
 }, (table) => [
-  (0, import_mysql_core.index)("agent_message_thread_created_idx").on(table.threadId, table.createdAt)
+  (0, import_mysql_core.index)("agent_message_thread_created_idx").on(table.threadId, table.createdAt),
+  // Transcript reads order a thread's messages by `seq`.
+  (0, import_mysql_core.index)("agent_message_thread_seq_idx").on(table.threadId, table.seq)
 ]);
 var agentToolCall2 = (0, import_mysql_core.mysqlTable)("agent_tool_call", {
   proposalId: key("proposal_id"),
@@ -647,7 +661,11 @@ var agentToolCall2 = (0, import_mysql_core.mysqlTable)("agent_tool_call", {
   remember: (0, import_mysql_core.boolean)("remember"),
   decidedVia: key("decided_via")
 }, (table) => [
-  (0, import_mysql_core.index)("agent_tool_call_message_idx").on(table.messageId)
+  (0, import_mysql_core.index)("agent_tool_call_message_idx").on(table.messageId),
+  // The approvals inbox: `pending_approval`, oldest first.
+  (0, import_mysql_core.index)("agent_tool_call_status_created_idx").on(table.status, table.createdAt),
+  (0, import_mysql_core.index)("agent_tool_call_run_idx").on(table.runId),
+  (0, import_mysql_core.index)("agent_tool_call_created_idx").on(table.createdAt)
 ]);
 var agentQueuedMessage2 = (0, import_mysql_core.mysqlTable)("agent_queued_message", {
   uiCapabilities: (0, import_mysql_core.json)("ui_capabilities").$type(),
@@ -689,7 +707,9 @@ var agentTokenUsage2 = (0, import_mysql_core.mysqlTable)("agent_token_usage", {
   costSource: key("cost_source").$type(),
   createdAt: ts("created_at").notNull()
 }, (table) => [
-  (0, import_mysql_core.index)("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt)
+  (0, import_mysql_core.index)("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt),
+  (0, import_mysql_core.index)("agent_token_usage_created_idx").on(table.createdAt),
+  (0, import_mysql_core.index)("agent_token_usage_thread_idx").on(table.threadId)
 ]);
 var agentGenuiCall2 = (0, import_mysql_core.mysqlTable)("agent_genui_call", {
   id: key("id").primaryKey(),
@@ -1048,7 +1068,9 @@ var agentThread3 = (0, import_pg_core.pgTable)("agent_thread", {
   updatedAt: ts2("updated_at").notNull(),
   deletedAt: ts2("deleted_at")
 }, (table) => [
-  (0, import_pg_core.index)("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt)
+  (0, import_pg_core.index)("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt),
+  // `ownerOfActiveStream` reads by run id on every stream attach and reconnect.
+  (0, import_pg_core.index)("agent_thread_active_stream_idx").on(table.activeStreamId)
 ]);
 var agentMessage3 = (0, import_pg_core.pgTable)("agent_message", {
   id: (0, import_pg_core.text)("id").primaryKey(),
@@ -1074,7 +1096,9 @@ var agentMessage3 = (0, import_pg_core.pgTable)("agent_message", {
   seq: (0, import_pg_core.integer)("seq"),
   createdAt: ts2("created_at").notNull()
 }, (table) => [
-  (0, import_pg_core.index)("agent_message_thread_created_idx").on(table.threadId, table.createdAt)
+  (0, import_pg_core.index)("agent_message_thread_created_idx").on(table.threadId, table.createdAt),
+  // Transcript reads order a thread's messages by `seq`.
+  (0, import_pg_core.index)("agent_message_thread_seq_idx").on(table.threadId, table.seq)
 ]);
 var agentToolCall3 = (0, import_pg_core.pgTable)("agent_tool_call", {
   proposalId: (0, import_pg_core.text)("proposal_id"),
@@ -1099,7 +1123,11 @@ var agentToolCall3 = (0, import_pg_core.pgTable)("agent_tool_call", {
   remember: (0, import_pg_core.boolean)("remember"),
   decidedVia: (0, import_pg_core.text)("decided_via")
 }, (table) => [
-  (0, import_pg_core.index)("agent_tool_call_message_idx").on(table.messageId)
+  (0, import_pg_core.index)("agent_tool_call_message_idx").on(table.messageId),
+  // The approvals inbox: `pending_approval`, oldest first.
+  (0, import_pg_core.index)("agent_tool_call_status_created_idx").on(table.status, table.createdAt),
+  (0, import_pg_core.index)("agent_tool_call_run_idx").on(table.runId),
+  (0, import_pg_core.index)("agent_tool_call_created_idx").on(table.createdAt)
 ]);
 var agentQueuedMessage3 = (0, import_pg_core.pgTable)("agent_queued_message", {
   uiCapabilities: (0, import_pg_core.jsonb)("ui_capabilities").$type(),
@@ -1141,7 +1169,9 @@ var agentTokenUsage3 = (0, import_pg_core.pgTable)("agent_token_usage", {
   costSource: (0, import_pg_core.text)("cost_source").$type(),
   createdAt: ts2("created_at").notNull()
 }, (table) => [
-  (0, import_pg_core.index)("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt)
+  (0, import_pg_core.index)("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt),
+  (0, import_pg_core.index)("agent_token_usage_created_idx").on(table.createdAt),
+  (0, import_pg_core.index)("agent_token_usage_thread_idx").on(table.threadId)
 ]);
 var agentGenuiCall3 = (0, import_pg_core.pgTable)("agent_genui_call", {
   id: (0, import_pg_core.text)("id").primaryKey(),
@@ -2420,12 +2450,49 @@ var DrizzleAgentStore = class {
   }
   async listThreads(actorRef, limit = 50) {
     const threads = await this.db.select().from(this.t.agentThread).where((0, import_drizzle_orm4.and)((0, import_drizzle_orm4.eq)(this.t.agentThread.actorRef, actorRef), (0, import_drizzle_orm4.eq)(this.t.agentThread.transient, false), (0, import_drizzle_orm4.isNull)(this.t.agentThread.deletedAt))).orderBy((0, import_drizzle_orm4.desc)(this.t.agentThread.updatedAt)).limit(limit);
-    const summaries = [];
-    for (const thread of threads) {
-      const [last] = await this.db.select().from(this.t.agentMessage).where((0, import_drizzle_orm4.eq)(this.t.agentMessage.threadId, thread.id)).orderBy(...this.messageOrderNewestFirst()).limit(1);
-      summaries.push(this.toSummary(thread, last?.content));
+    const previews = await this.lastMessageContents(threads.map((thread) => thread.id));
+    return threads.map((thread) => this.toSummary(thread, previews.get(thread.id)));
+  }
+  /**
+  * The content of each thread's newest message ({@link messageOrderNewestFirst}), in a fixed two
+  * reads for any number of threads: each thread's highest `seq`, then the rows holding it. Only a
+  * thread whose every message predates `seq` is read on its own, by `created_at`.
+  */
+  async lastMessageContents(threadIds) {
+    const contents = /* @__PURE__ */ new Map();
+    if (threadIds.length === 0) {
+      return contents;
+    }
+    const message = this.t.agentMessage;
+    const highest = await this.db.select({
+      threadId: message.threadId,
+      seq: (0, import_drizzle_orm4.max)(message.seq)
+    }).from(message).where((0, import_drizzle_orm4.inArray)(message.threadId, [
+      ...threadIds
+    ])).groupBy(message.threadId);
+    const numbered = highest.filter((row) => row.seq !== null);
+    if (numbered.length > 0) {
+      const rows = await this.db.select({
+        threadId: message.threadId,
+        content: message.content
+      }).from(message).where((0, import_drizzle_orm4.or)(...numbered.map((row) => (0, import_drizzle_orm4.and)((0, import_drizzle_orm4.eq)(message.threadId, row.threadId), (0, import_drizzle_orm4.eq)(message.seq, Number(row.seq)))))).orderBy(...this.messageOrderNewestFirst());
+      for (const row of rows) {
+        if (!contents.has(row.threadId)) {
+          contents.set(row.threadId, row.content);
+        }
+      }
+    }
+    for (const row of highest) {
+      if (row.seq === null) {
+        const [last] = await this.db.select({
+          content: message.content
+        }).from(message).where((0, import_drizzle_orm4.eq)(message.threadId, row.threadId)).orderBy(...this.messageOrderNewestFirst()).limit(1);
+        if (last !== void 0) {
+          contents.set(row.threadId, last.content);
+        }
+      }
     }
-    return summaries;
+    return contents;
   }
   async softDeleteThread(threadId) {
     await this.db.update(this.t.agentThread).set({
@@ -2802,6 +2869,24 @@ var DrizzleAgentStore = class {
     }).from(agentThread).where((0, import_drizzle_orm4.and)((0, import_drizzle_orm4.eq)(this.t.agentThread.id, threadId), (0, import_drizzle_orm4.isNull)(this.t.agentThread.deletedAt)));
     return thread?.persona ?? null;
   }
+  async threadSettings(threadId) {
+    const thread = this.t.agentThread;
+    const [row] = await this.db.select({
+      actorRef: thread.actorRef,
+      defaultAgent: thread.defaultAgent,
+      persona: thread.persona,
+      model: thread.model
+    }).from(thread).where((0, import_drizzle_orm4.and)((0, import_drizzle_orm4.eq)(thread.id, threadId), (0, import_drizzle_orm4.isNull)(thread.deletedAt)));
+    if (row === void 0) {
+      return null;
+    }
+    return {
+      actorRef: row.actorRef,
+      defaultAgent: row.defaultAgent ?? null,
+      persona: row.persona ?? null,
+      model: row.model ?? null
+    };
+  }
   async defaultAgentForThread(threadId) {
     const [thread] = await this.db.select({
       defaultAgent: this.t.agentThread.defaultAgent
@@ -3117,10 +3202,43 @@ var DrizzleAgentStore = class {
     return this.usageBetween(actorRef, day, day);
   }
   async usageBetween(actorRef, fromDay, toDay) {
-    const start = /* @__PURE__ */ new Date(`${fromDay}T00:00:00.000Z`);
-    const end = /* @__PURE__ */ new Date(`${toDay}T23:59:59.999Z`);
-    const rows = await this.db.select().from(this.t.agentTokenUsage).where((0, import_drizzle_orm4.and)((0, import_drizzle_orm4.eq)(this.t.agentTokenUsage.actorRef, actorRef), (0, import_drizzle_orm4.gte)(this.t.agentTokenUsage.createdAt, start), (0, import_drizzle_orm4.lte)(this.t.agentTokenUsage.createdAt, end)));
-    return (0, import_nestjs_agent_core4.sumUsage)(rows);
+    const [totals] = await this.usageWindows(actorRef, [
+      {
+        fromDay,
+        toDay
+      }
+    ]);
+    return totals ?? (0, import_nestjs_agent_core4.sumUsage)([]);
+  }
+  /**
+  * Every range summed by the database in one statement, each as a conditional aggregate over the
+  * rows of the smallest span covering them all — the actor's `(actor_ref, created_at)` index serves
+  * it, and no usage row crosses the wire.
+  */
+  async usageWindows(actorRef, ranges) {
+    if (ranges.length === 0) {
+      return [];
+    }
+    const usage = this.t.agentTokenUsage;
+    const bounds = ranges.map(({ fromDay, toDay }) => ({
+      start: /* @__PURE__ */ new Date(`${fromDay}T00:00:00.000Z`),
+      end: /* @__PURE__ */ new Date(`${toDay}T23:59:59.999Z`)
+    }));
+    const within = /* @__PURE__ */ __name((start, end) => (0, import_drizzle_orm4.and)((0, import_drizzle_orm4.gte)(usage.createdAt, start), (0, import_drizzle_orm4.lte)(usage.createdAt, end)), "within");
+    const columns = {};
+    bounds.forEach(({ start, end }, index4) => {
+      columns[`tokens${index4}`] = import_drizzle_orm4.sql`sum(case when ${within(start, end)} then ${usage.inputTokens} + ${usage.outputTokens} else 0 end)`;
+      columns[`cost${index4}`] = import_drizzle_orm4.sql`sum(case when ${within(start, end)} then ${usage.costUsd} else 0 end)`;
+      columns[`estimated${index4}`] = import_drizzle_orm4.sql`sum(case when ${within(start, end)} and ${usage.costSource} = 'estimate' then ${usage.costUsd} else 0 end)`;
+    });
+    const from = new Date(Math.min(...bounds.map(({ start }) => start.getTime())));
+    const to = new Date(Math.max(...bounds.map(({ end }) => end.getTime())));
+    const [row] = await this.db.select(columns).from(usage).where((0, import_drizzle_orm4.and)((0, import_drizzle_orm4.eq)(usage.actorRef, actorRef), within(from, to)));
+    return bounds.map((_, index4) => ({
+      usedTokens: Number(row?.[`tokens${index4}`] ?? 0),
+      costUsd: Number(row?.[`cost${index4}`] ?? 0),
+      estimatedCostUsd: Number(row?.[`estimated${index4}`] ?? 0)
+    }));
   }
   toSummary(thread, lastContent) {
     return {
@@ -3416,7 +3534,18 @@ var DrizzleGovernanceQueries = class {
   }
   async usageInRange(range) {
     const { start, end } = (0, import_nestjs_agent_core5.dayBoundsUtc)(range);
-    const rows = await this.db.select().from(this.t.agentTokenUsage).where((0, import_drizzle_orm5.and)((0, import_drizzle_orm5.gte)(this.t.agentTokenUsage.createdAt, start), (0, import_drizzle_orm5.lte)(this.t.agentTokenUsage.createdAt, end)));
+    const usage = this.t.agentTokenUsage;
+    const rows = await this.db.select({
+      modelId: usage.modelId,
+      actorRef: usage.actorRef,
+      threadId: usage.threadId,
+      createdAt: usage.createdAt,
+      inputTokens: usage.inputTokens,
+      outputTokens: usage.outputTokens,
+      cacheWriteTokens: usage.cacheWriteTokens,
+      cacheReadTokens: usage.cacheReadTokens,
+      costUsd: usage.costUsd
+    }).from(usage).where((0, import_drizzle_orm5.and)((0, import_drizzle_orm5.gte)(this.t.agentTokenUsage.createdAt, start), (0, import_drizzle_orm5.lte)(this.t.agentTokenUsage.createdAt, end)));
     return rows.map(toUsageInput);
   }
   async spendByModel(range) {
@@ -3441,7 +3570,11 @@ var DrizzleGovernanceQueries = class {
     if (threadIds.length === 0) {
       return [];
     }
-    const threads = await this.db.select().from(this.t.agentThread).where((0, import_drizzle_orm5.and)((0, import_drizzle_orm5.inArray)(this.t.agentThread.id, threadIds), (0, import_drizzle_orm5.isNull)(this.t.agentThread.deletedAt)));
+    const threads = await this.db.select({
+      id: this.t.agentThread.id,
+      title: this.t.agentThread.title,
+      actorRef: this.t.agentThread.actorRef
+    }).from(this.t.agentThread).where((0, import_drizzle_orm5.and)((0, import_drizzle_orm5.inArray)(this.t.agentThread.id, threadIds), (0, import_drizzle_orm5.isNull)(this.t.agentThread.deletedAt)));
     const threadsById = new Map(threads.map((thread) => [
       thread.id,
       {
@@ -3536,9 +3669,19 @@ var DrizzleGovernanceQueries = class {
     const threads = await this.db.select().from(this.t.agentThread).where((0, import_drizzle_orm5.isNull)(this.t.agentThread.deletedAt)).orderBy((0, import_drizzle_orm5.desc)(this.t.agentThread.updatedAt), (0, import_drizzle_orm5.desc)(this.t.agentThread.id)).limit(limit);
     return this.toThreadActivityRows(threads);
   }
+  /** The range's runs, with only the columns the rollups read. */
   async runsInRange(range) {
     const { start, end } = (0, import_nestjs_agent_core5.dayBoundsUtc)(range);
-    return this.db.select().from(this.t.agentRun).where((0, import_drizzle_orm5.and)((0, import_drizzle_orm5.gte)(this.t.agentRun.startedAt, start), (0, import_drizzle_orm5.lte)(this.t.agentRun.startedAt, end)));
+    const run = this.t.agentRun;
+    return this.db.select({
+      agentName: run.agentName,
+      status: run.status,
+      retries: run.retries,
+      durationMs: run.durationMs,
+      ttftMs: run.ttftMs,
+      parentRunId: run.parentRunId,
+      startedAt: run.startedAt
+    }).from(run).where((0, import_drizzle_orm5.and)((0, import_drizzle_orm5.gte)(run.startedAt, start), (0, import_drizzle_orm5.lte)(run.startedAt, end)));
   }
   async runMetrics(range) {
     const runs = await this.runsInRange(range);
@@ -3602,7 +3745,9 @@ var DrizzleGovernanceQueries = class {
   }
   async runErrors(range) {
     const { start, end } = (0, import_nestjs_agent_core5.dayBoundsUtc)(range);
-    const runs = await this.db.select().from(this.t.agentRun).where((0, import_drizzle_orm5.and)((0, import_drizzle_orm5.gte)(this.t.agentRun.startedAt, start), (0, import_drizzle_orm5.lte)(this.t.agentRun.startedAt, end), (0, import_drizzle_orm5.eq)(this.t.agentRun.status, "failed")));
+    const runs = await this.db.select({
+      errorCode: this.t.agentRun.errorCode
+    }).from(this.t.agentRun).where((0, import_drizzle_orm5.and)((0, import_drizzle_orm5.gte)(this.t.agentRun.startedAt, start), (0, import_drizzle_orm5.lte)(this.t.agentRun.startedAt, end), (0, import_drizzle_orm5.eq)(this.t.agentRun.status, "failed")));
     const byError = /* @__PURE__ */ new Map();
     for (const run of runs) {
       const errorCode = run.errorCode ?? UNCLASSIFIED_ERROR_CODE;
@@ -3683,7 +3828,12 @@ var DrizzleGovernanceQueries = class {
   */
   async toolStats(range) {
     const { start, end } = (0, import_nestjs_agent_core5.dayBoundsUtc)(range);
-    const calls = await this.db.select().from(this.t.agentToolCall).where((0, import_drizzle_orm5.and)((0, import_drizzle_orm5.gte)(this.t.agentToolCall.createdAt, start), (0, import_drizzle_orm5.lte)(this.t.agentToolCall.createdAt, end)));
+    const calls = await this.db.select({
+      toolName: this.t.agentToolCall.toolName,
+      toolType: this.t.agentToolCall.toolType,
+      status: this.t.agentToolCall.status,
+      executionMs: this.t.agentToolCall.executionMs
+    }).from(this.t.agentToolCall).where((0, import_drizzle_orm5.and)((0, import_drizzle_orm5.gte)(this.t.agentToolCall.createdAt, start), (0, import_drizzle_orm5.lte)(this.t.agentToolCall.createdAt, end)));
     const byTool = /* @__PURE__ */ new Map();
     for (const call of calls) {
       const key2 = `${call.toolName} ${call.toolType}`;