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

package/dist/index.js

@@ -32,7 +32,9 @@ var agentThread = sqliteTable("agent_thread", {
     mode: "timestamp_ms"
   })
 }, (table) => [
-  index("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt)
+  index("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt),
+  // `ownerOfActiveStream` reads by run id on every stream attach and reconnect.
+  index("agent_thread_active_stream_idx").on(table.activeStreamId)
 ]);
 var agentMessage = sqliteTable("agent_message", {
   id: text("id").primaryKey(),
@@ -82,7 +84,9 @@ var agentMessage = sqliteTable("agent_message", {
     mode: "timestamp_ms"
   }).notNull()
 }, (table) => [
-  index("agent_message_thread_created_idx").on(table.threadId, table.createdAt)
+  index("agent_message_thread_created_idx").on(table.threadId, table.createdAt),
+  // Transcript reads order a thread's messages by `seq`.
+  index("agent_message_thread_seq_idx").on(table.threadId, table.seq)
 ]);
 var agentToolCall = sqliteTable("agent_tool_call", {
   proposalId: text("proposal_id"),
@@ -127,7 +131,11 @@ var agentToolCall = sqliteTable("agent_tool_call", {
   /** The surface the decision came through (`web`, `slack`, `remembered`, …). */
   decidedVia: text("decided_via")
 }, (table) => [
-  index("agent_tool_call_message_idx").on(table.messageId)
+  index("agent_tool_call_message_idx").on(table.messageId),
+  // The approvals inbox: `pending_approval`, oldest first.
+  index("agent_tool_call_status_created_idx").on(table.status, table.createdAt),
+  index("agent_tool_call_run_idx").on(table.runId),
+  index("agent_tool_call_created_idx").on(table.createdAt)
 ]);
 var agentQueuedMessage = sqliteTable("agent_queued_message", {
   uiCapabilities: text("ui_capabilities", {
@@ -191,7 +199,9 @@ var agentTokenUsage = sqliteTable("agent_token_usage", {
     mode: "timestamp_ms"
   }).notNull()
 }, (table) => [
-  index("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt)
+  index("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt),
+  index("agent_token_usage_created_idx").on(table.createdAt),
+  index("agent_token_usage_thread_idx").on(table.threadId)
 ]);
 var agentGenuiCall = sqliteTable("agent_genui_call", {
   id: text("id").primaryKey(),
@@ -249,6 +259,12 @@ var agentRun = sqliteTable("agent_run", {
   agentName: text("agent_name"),
   status: text("status").$type().notNull(),
   durationMs: integer("duration_ms"),
+  /** Turn start to the first `text` frame; null for a run that streamed none or did not complete. */
+  ttftMs: integer("ttft_ms"),
+  /** The run's known spend (the sum of its model calls with a known cost); null when none had one. */
+  costUsd: real("cost_usd"),
+  /** The model the run's last model call resolved to. */
+  modelId: text("model_id"),
   errorCode: text("error_code"),
   errorMessage: text("error_message"),
   retries: integer("retries").notNull().default(0),
@@ -528,7 +544,9 @@ var agentThread2 = mysqlTable("agent_thread", {
   updatedAt: ts("updated_at").notNull(),
   deletedAt: ts("deleted_at")
 }, (table) => [
-  index2("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt)
+  index2("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt),
+  // `ownerOfActiveStream` reads by run id on every stream attach and reconnect.
+  index2("agent_thread_active_stream_idx").on(table.activeStreamId)
 ]);
 var agentMessage2 = mysqlTable("agent_message", {
   id: key("id").primaryKey(),
@@ -554,7 +572,9 @@ var agentMessage2 = mysqlTable("agent_message", {
   seq: int("seq"),
   createdAt: ts("created_at").notNull()
 }, (table) => [
-  index2("agent_message_thread_created_idx").on(table.threadId, table.createdAt)
+  index2("agent_message_thread_created_idx").on(table.threadId, table.createdAt),
+  // Transcript reads order a thread's messages by `seq`.
+  index2("agent_message_thread_seq_idx").on(table.threadId, table.seq)
 ]);
 var agentToolCall2 = mysqlTable("agent_tool_call", {
   proposalId: key("proposal_id"),
@@ -579,7 +599,11 @@ var agentToolCall2 = mysqlTable("agent_tool_call", {
   remember: boolean("remember"),
   decidedVia: key("decided_via")
 }, (table) => [
-  index2("agent_tool_call_message_idx").on(table.messageId)
+  index2("agent_tool_call_message_idx").on(table.messageId),
+  // The approvals inbox: `pending_approval`, oldest first.
+  index2("agent_tool_call_status_created_idx").on(table.status, table.createdAt),
+  index2("agent_tool_call_run_idx").on(table.runId),
+  index2("agent_tool_call_created_idx").on(table.createdAt)
 ]);
 var agentQueuedMessage2 = mysqlTable("agent_queued_message", {
   uiCapabilities: json("ui_capabilities").$type(),
@@ -621,7 +645,9 @@ var agentTokenUsage2 = mysqlTable("agent_token_usage", {
   costSource: key("cost_source").$type(),
   createdAt: ts("created_at").notNull()
 }, (table) => [
-  index2("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt)
+  index2("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt),
+  index2("agent_token_usage_created_idx").on(table.createdAt),
+  index2("agent_token_usage_thread_idx").on(table.threadId)
 ]);
 var agentGenuiCall2 = mysqlTable("agent_genui_call", {
   id: key("id").primaryKey(),
@@ -681,6 +707,12 @@ var agentRun2 = mysqlTable("agent_run", {
   agentName: key("agent_name"),
   status: key("status").$type().notNull(),
   durationMs: int("duration_ms"),
+  /** Turn start to the first `text` frame; null for a run that streamed none or did not complete. */
+  ttftMs: int("ttft_ms"),
+  /** The run's known spend (the sum of its model calls with a known cost); null when none had one. */
+  costUsd: double("cost_usd"),
+  /** The model the run's last model call resolved to. */
+  modelId: key("model_id"),
   errorCode: key("error_code"),
   errorMessage: longtext("error_message"),
   retries: int("retries").notNull().default(0),
@@ -974,7 +1006,9 @@ var agentThread3 = pgTable("agent_thread", {
   updatedAt: ts2("updated_at").notNull(),
   deletedAt: ts2("deleted_at")
 }, (table) => [
-  index3("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt)
+  index3("agent_thread_actor_updated_idx").on(table.actorRef, table.updatedAt),
+  // `ownerOfActiveStream` reads by run id on every stream attach and reconnect.
+  index3("agent_thread_active_stream_idx").on(table.activeStreamId)
 ]);
 var agentMessage3 = pgTable("agent_message", {
   id: text2("id").primaryKey(),
@@ -1000,7 +1034,9 @@ var agentMessage3 = pgTable("agent_message", {
   seq: integer2("seq"),
   createdAt: ts2("created_at").notNull()
 }, (table) => [
-  index3("agent_message_thread_created_idx").on(table.threadId, table.createdAt)
+  index3("agent_message_thread_created_idx").on(table.threadId, table.createdAt),
+  // Transcript reads order a thread's messages by `seq`.
+  index3("agent_message_thread_seq_idx").on(table.threadId, table.seq)
 ]);
 var agentToolCall3 = pgTable("agent_tool_call", {
   proposalId: text2("proposal_id"),
@@ -1025,7 +1061,11 @@ var agentToolCall3 = pgTable("agent_tool_call", {
   remember: boolean2("remember"),
   decidedVia: text2("decided_via")
 }, (table) => [
-  index3("agent_tool_call_message_idx").on(table.messageId)
+  index3("agent_tool_call_message_idx").on(table.messageId),
+  // The approvals inbox: `pending_approval`, oldest first.
+  index3("agent_tool_call_status_created_idx").on(table.status, table.createdAt),
+  index3("agent_tool_call_run_idx").on(table.runId),
+  index3("agent_tool_call_created_idx").on(table.createdAt)
 ]);
 var agentQueuedMessage3 = pgTable("agent_queued_message", {
   uiCapabilities: jsonb("ui_capabilities").$type(),
@@ -1067,7 +1107,9 @@ var agentTokenUsage3 = pgTable("agent_token_usage", {
   costSource: text2("cost_source").$type(),
   createdAt: ts2("created_at").notNull()
 }, (table) => [
-  index3("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt)
+  index3("agent_token_usage_actor_created_idx").on(table.actorRef, table.createdAt),
+  index3("agent_token_usage_created_idx").on(table.createdAt),
+  index3("agent_token_usage_thread_idx").on(table.threadId)
 ]);
 var agentGenuiCall3 = pgTable("agent_genui_call", {
   id: text2("id").primaryKey(),
@@ -1127,6 +1169,12 @@ var agentRun3 = pgTable("agent_run", {
   agentName: text2("agent_name"),
   status: text2("status").$type().notNull(),
   durationMs: integer2("duration_ms"),
+  /** Turn start to the first `text` frame; null for a run that streamed none or did not complete. */
+  ttftMs: integer2("ttft_ms"),
+  /** The run's known spend (the sum of its model calls with a known cost); null when none had one. */
+  costUsd: doublePrecision("cost_usd"),
+  /** The model the run's last model call resolved to. */
+  modelId: text2("model_id"),
   errorCode: text2("error_code"),
   errorMessage: text2("error_message"),
   retries: integer2("retries").notNull().default(0),
@@ -2340,12 +2388,49 @@ var DrizzleAgentStore = class {
   }
   async listThreads(actorRef, limit = 50) {
     const threads = await this.db.select().from(this.t.agentThread).where(and2(eq2(this.t.agentThread.actorRef, actorRef), eq2(this.t.agentThread.transient, false), isNull2(this.t.agentThread.deletedAt))).orderBy(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(eq2(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: max2(message.seq)
+    }).from(message).where(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(or2(...numbered.map((row) => and2(eq2(message.threadId, row.threadId), eq2(message.seq, Number(row.seq)))))).orderBy(...this.messageOrderNewestFirst());
+      for (const row of rows) {
+        if (!contents.has(row.threadId)) {
+          contents.set(row.threadId, row.content);
+        }
+      }
     }
-    return summaries;
+    for (const row of highest) {
+      if (row.seq === null) {
+        const [last] = await this.db.select({
+          content: message.content
+        }).from(message).where(eq2(message.threadId, row.threadId)).orderBy(...this.messageOrderNewestFirst()).limit(1);
+        if (last !== void 0) {
+          contents.set(row.threadId, last.content);
+        }
+      }
+    }
+    return contents;
   }
   async softDeleteThread(threadId) {
     await this.db.update(this.t.agentThread).set({
@@ -2722,6 +2807,24 @@ var DrizzleAgentStore = class {
     }).from(agentThread).where(and2(eq2(this.t.agentThread.id, threadId), isNull2(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(and2(eq2(thread.id, threadId), isNull2(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
@@ -2742,6 +2845,9 @@ var DrizzleAgentStore = class {
       agentName: run.agentName ?? null,
       status: "running",
       durationMs: null,
+      ttftMs: null,
+      costUsd: null,
+      modelId: null,
       errorCode: null,
       errorMessage: null,
       retries: 0,
@@ -2761,6 +2867,15 @@ var DrizzleAgentStore = class {
     if (end.durationMs !== void 0) {
       updates.durationMs = end.durationMs;
     }
+    if (end.ttftMs !== void 0) {
+      updates.ttftMs = end.ttftMs;
+    }
+    if (end.costUsd !== void 0) {
+      updates.costUsd = end.costUsd;
+    }
+    if (end.modelId !== void 0) {
+      updates.modelId = end.modelId;
+    }
     if (end.errorCode !== void 0) {
       updates.errorCode = end.errorCode;
     }
@@ -3025,10 +3140,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(and2(eq2(this.t.agentTokenUsage.actorRef, actorRef), gte(this.t.agentTokenUsage.createdAt, start), lte2(this.t.agentTokenUsage.createdAt, end)));
-    return sumUsage(rows);
+    const [totals] = await this.usageWindows(actorRef, [
+      {
+        fromDay,
+        toDay
+      }
+    ]);
+    return totals ?? 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) => and2(gte(usage.createdAt, start), lte2(usage.createdAt, end)), "within");
+    const columns = {};
+    bounds.forEach(({ start, end }, index4) => {
+      columns[`tokens${index4}`] = sql3`sum(case when ${within(start, end)} then ${usage.inputTokens} + ${usage.outputTokens} else 0 end)`;
+      columns[`cost${index4}`] = sql3`sum(case when ${within(start, end)} then ${usage.costUsd} else 0 end)`;
+      columns[`estimated${index4}`] = sql3`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(and2(eq2(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 {
@@ -3246,6 +3394,11 @@ function toRecentRunRow(run) {
     agentName: run.agentName ?? null,
     status: run.status,
     durationMs: run.durationMs ?? null,
+    ttftMs: run.ttftMs ?? null,
+    // MySQL's DOUBLE and Postgres' double precision come back as numbers; a driver that hands a
+    // decimal back as a string is normalized here so the SPI row is always a number.
+    costUsd: run.costUsd == null ? null : Number(run.costUsd),
+    modelId: run.modelId ?? null,
     errorCode: run.errorCode ?? null,
     errorMessage: run.errorMessage ?? null,
     retries: run.retries,
@@ -3319,7 +3472,18 @@ var DrizzleGovernanceQueries = class {
   }
   async usageInRange(range) {
     const { start, end } = dayBoundsUtc(range);
-    const rows = await this.db.select().from(this.t.agentTokenUsage).where(and3(gte2(this.t.agentTokenUsage.createdAt, start), lte3(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(and3(gte2(this.t.agentTokenUsage.createdAt, start), lte3(this.t.agentTokenUsage.createdAt, end)));
     return rows.map(toUsageInput);
   }
   async spendByModel(range) {
@@ -3344,7 +3508,11 @@ var DrizzleGovernanceQueries = class {
     if (threadIds.length === 0) {
       return [];
     }
-    const threads = await this.db.select().from(this.t.agentThread).where(and3(inArray2(this.t.agentThread.id, threadIds), isNull3(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(and3(inArray2(this.t.agentThread.id, threadIds), isNull3(this.t.agentThread.deletedAt)));
     const threadsById = new Map(threads.map((thread) => [
       thread.id,
       {
@@ -3439,9 +3607,19 @@ var DrizzleGovernanceQueries = class {
     const threads = await this.db.select().from(this.t.agentThread).where(isNull3(this.t.agentThread.deletedAt)).orderBy(desc2(this.t.agentThread.updatedAt), desc2(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 } = dayBoundsUtc(range);
-    return this.db.select().from(this.t.agentRun).where(and3(gte2(this.t.agentRun.startedAt, start), lte3(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(and3(gte2(run.startedAt, start), lte3(run.startedAt, end)));
   }
   async runMetrics(range) {
     const runs = await this.runsInRange(range);
@@ -3449,6 +3627,7 @@ var DrizzleGovernanceQueries = class {
     let failed = 0;
     let retries = 0;
     const settledDurationsMs = [];
+    const ttftsMs = [];
     for (const run of runs) {
       if (run.status === "completed") {
         completed += 1;
@@ -3459,8 +3638,12 @@ var DrizzleGovernanceQueries = class {
       if ((run.status === "completed" || run.status === "failed") && run.durationMs != null) {
         settledDurationsMs.push(run.durationMs);
       }
+      if (run.status === "completed" && run.parentRunId == null && run.ttftMs != null) {
+        ttftsMs.push(run.ttftMs);
+      }
     }
     settledDurationsMs.sort((left, right) => left - right);
+    ttftsMs.sort((left, right) => left - right);
     return {
       runs: runs.length,
       completed,
@@ -3468,7 +3651,9 @@ var DrizzleGovernanceQueries = class {
       successRate: runs.length === 0 ? 0 : completed / runs.length,
       retries,
       durationP50Ms: percentileMs(settledDurationsMs, 0.5),
-      durationP95Ms: percentileMs(settledDurationsMs, 0.95)
+      durationP95Ms: percentileMs(settledDurationsMs, 0.95),
+      ttftP50Ms: percentileMs(ttftsMs, 0.5),
+      ttftP95Ms: percentileMs(ttftsMs, 0.95)
     };
   }
   async runsByAgent(range) {
@@ -3498,7 +3683,9 @@ var DrizzleGovernanceQueries = class {
   }
   async runErrors(range) {
     const { start, end } = dayBoundsUtc(range);
-    const runs = await this.db.select().from(this.t.agentRun).where(and3(gte2(this.t.agentRun.startedAt, start), lte3(this.t.agentRun.startedAt, end), eq3(this.t.agentRun.status, "failed")));
+    const runs = await this.db.select({
+      errorCode: this.t.agentRun.errorCode
+    }).from(this.t.agentRun).where(and3(gte2(this.t.agentRun.startedAt, start), lte3(this.t.agentRun.startedAt, end), eq3(this.t.agentRun.status, "failed")));
     const byError = /* @__PURE__ */ new Map();
     for (const run of runs) {
       const errorCode = run.errorCode ?? UNCLASSIFIED_ERROR_CODE;
@@ -3579,7 +3766,12 @@ var DrizzleGovernanceQueries = class {
   */
   async toolStats(range) {
     const { start, end } = dayBoundsUtc(range);
-    const calls = await this.db.select().from(this.t.agentToolCall).where(and3(gte2(this.t.agentToolCall.createdAt, start), lte3(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(and3(gte2(this.t.agentToolCall.createdAt, start), lte3(this.t.agentToolCall.createdAt, end)));
     const byTool = /* @__PURE__ */ new Map();
     for (const call of calls) {
       const key2 = `${call.toolName} ${call.toolType}`;