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

package/dist/index.js

@@ -178,7 +178,7 @@ var agentQueuedMessage = sqliteTable("agent_queued_message", {
 ]);
 var agentTokenUsage = sqliteTable("agent_token_usage", {
   id: text("id").primaryKey(),
-  threadId: text("thread_id").notNull().references(() => agentThread.id, {
+  threadId: text("thread_id").references(() => agentThread.id, {
     onDelete: "cascade"
   }),
   actorRef: text("actor_ref").notNull(),
@@ -195,13 +195,28 @@ var agentTokenUsage = sqliteTable("agent_token_usage", {
   costUsd: real("cost_usd"),
   /** Where `cost_usd` came from: `'provider'` (reported) or `'estimate'` (tokens × the price row); null on rows from before it existed (reported). */
   costSource: text("cost_source").$type(),
+  /** The run the row belongs to. */
+  runId: text("run_id"),
+  /** What the spend was for: `agent:<name>` for an agent turn, the app's own otherwise; null → `agent`. */
+  feature: text("feature"),
+  provider: text("provider"),
+  tenantRef: text("tenant_ref"),
+  /** Units a row bills by when they are not tokens. */
+  audioSeconds: real("audio_seconds"),
+  characters: integer("characters"),
+  requests: integer("requests"),
+  /** The caller's own JSON, never read by the library. */
+  metadata: text("metadata", {
+    mode: "json"
+  }).$type(),
   createdAt: integer("created_at", {
     mode: "timestamp_ms"
   }).notNull()
 }, (table) => [
   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)
+  index("agent_token_usage_thread_idx").on(table.threadId),
+  index("agent_token_usage_feature_created_idx").on(table.feature, table.createdAt)
 ]);
 var agentGenuiCall = sqliteTable("agent_genui_call", {
   id: text("id").primaryKey(),
@@ -283,7 +298,17 @@ var agentRun = sqliteTable("agent_run", {
    * for a detached child, which outlives its parent's turn and so is paired by nothing in the
    * transcript.
    */
-  parentRunId: text("parent_run_id")
+  parentRunId: text("parent_run_id"),
+  /** Enqueue → body start (ms); null when not read. */
+  queueMs: integer("queue_ms"),
+  /** Where the completed run's time went: model and tool totals, counts, and the per-step list. */
+  modelMs: integer("model_ms"),
+  toolsMs: integer("tools_ms"),
+  modelCalls: integer("model_calls"),
+  toolCalls: integer("tool_calls"),
+  timings: text("timings", {
+    mode: "json"
+  }).$type()
 }, (table) => [
   index("agent_run_started_idx").on(table.startedAt)
 ]);
@@ -629,7 +654,7 @@ var agentQueuedMessage2 = mysqlTable("agent_queued_message", {
 ]);
 var agentTokenUsage2 = mysqlTable("agent_token_usage", {
   id: key("id").primaryKey(),
-  threadId: key("thread_id").notNull().references(() => agentThread2.id, {
+  threadId: key("thread_id").references(() => agentThread2.id, {
     onDelete: "cascade"
   }),
   actorRef: key("actor_ref").notNull(),
@@ -643,11 +668,26 @@ var agentTokenUsage2 = mysqlTable("agent_token_usage", {
   costUsd: double("cost_usd"),
   /** Where `cost_usd` came from: `'provider'` (reported) or `'estimate'` (tokens × the price row); null on rows from before it existed (reported). */
   costSource: key("cost_source").$type(),
+  /** The run the row belongs to. */
+  runId: key("run_id"),
+  /** What the spend was for: `agent:<name>` for an agent turn, the app's own otherwise; null → `agent`. */
+  feature: key("feature"),
+  provider: key("provider"),
+  tenantRef: key("tenant_ref"),
+  /** Units a row bills by when they are not tokens. */
+  audioSeconds: double("audio_seconds"),
+  characters: bigint("characters", {
+    mode: "number"
+  }),
+  requests: int("requests"),
+  /** The caller's own JSON, never read by the library. */
+  metadata: json("metadata").$type(),
   createdAt: ts("created_at").notNull()
 }, (table) => [
   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)
+  index2("agent_token_usage_thread_idx").on(table.threadId),
+  index2("agent_token_usage_feature_created_idx").on(table.feature, table.createdAt)
 ]);
 var agentGenuiCall2 = mysqlTable("agent_genui_call", {
   id: key("id").primaryKey(),
@@ -719,7 +759,15 @@ var agentRun2 = mysqlTable("agent_run", {
   startedAt: ts("started_at").notNull(),
   settledAt: ts("settled_at"),
   promptHash: key("prompt_hash"),
-  parentRunId: key("parent_run_id")
+  parentRunId: key("parent_run_id"),
+  /** Enqueue → body start (ms); null when not read. */
+  queueMs: int("queue_ms"),
+  /** Where the completed run's time went: model and tool totals, counts, and the per-step list. */
+  modelMs: int("model_ms"),
+  toolsMs: int("tools_ms"),
+  modelCalls: int("model_calls"),
+  toolCalls: int("tool_calls"),
+  timings: json("timings").$type()
 }, (table) => [
   index2("agent_run_started_idx").on(table.startedAt)
 ]);
@@ -1091,7 +1139,7 @@ var agentQueuedMessage3 = pgTable("agent_queued_message", {
 ]);
 var agentTokenUsage3 = pgTable("agent_token_usage", {
   id: text2("id").primaryKey(),
-  threadId: text2("thread_id").notNull().references(() => agentThread3.id, {
+  threadId: text2("thread_id").references(() => agentThread3.id, {
     onDelete: "cascade"
   }),
   actorRef: text2("actor_ref").notNull(),
@@ -1105,11 +1153,26 @@ var agentTokenUsage3 = pgTable("agent_token_usage", {
   costUsd: doublePrecision("cost_usd"),
   /** Where `cost_usd` came from: `'provider'` (reported) or `'estimate'` (tokens × the price row); null on rows from before it existed (reported). */
   costSource: text2("cost_source").$type(),
+  /** The run the row belongs to. */
+  runId: text2("run_id"),
+  /** What the spend was for: `agent:<name>` for an agent turn, the app's own otherwise; null → `agent`. */
+  feature: text2("feature"),
+  provider: text2("provider"),
+  tenantRef: text2("tenant_ref"),
+  /** Units a row bills by when they are not tokens. */
+  audioSeconds: doublePrecision("audio_seconds"),
+  characters: bigint2("characters", {
+    mode: "number"
+  }),
+  requests: integer2("requests"),
+  /** The caller's own JSON, never read by the library. */
+  metadata: jsonb("metadata").$type(),
   createdAt: ts2("created_at").notNull()
 }, (table) => [
   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)
+  index3("agent_token_usage_thread_idx").on(table.threadId),
+  index3("agent_token_usage_feature_created_idx").on(table.feature, table.createdAt)
 ]);
 var agentGenuiCall3 = pgTable("agent_genui_call", {
   id: text2("id").primaryKey(),
@@ -1181,7 +1244,15 @@ var agentRun3 = pgTable("agent_run", {
   startedAt: ts2("started_at").notNull(),
   settledAt: ts2("settled_at"),
   promptHash: text2("prompt_hash"),
-  parentRunId: text2("parent_run_id")
+  parentRunId: text2("parent_run_id"),
+  /** Enqueue → body start (ms); null when not read. */
+  queueMs: integer2("queue_ms"),
+  /** Where the completed run's time went: model and tool totals, counts, and the per-step list. */
+  modelMs: integer2("model_ms"),
+  toolsMs: integer2("tools_ms"),
+  modelCalls: integer2("model_calls"),
+  toolCalls: integer2("tool_calls"),
+  timings: jsonb("timings").$type()
 }, (table) => [
   index3("agent_run_started_idx").on(table.startedAt)
 ]);
@@ -1639,10 +1710,60 @@ function specs(dialect) {
 }
 __name(specs, "specs");
 async function existingColumns(db, dialect, table) {
-  const rows = dialect === "sqlite" ? await querySql(db, dialect, sql.raw(`PRAGMA table_info("${table}")`)) : await querySql(db, dialect, sql`select column_name as name from information_schema.columns where table_schema = ${sql.raw(dialect === "mysql" ? "database()" : "current_schema()")} and table_name = ${table}`);
-  return new Set(rows.map((row) => String(row.name)));
+  return new Set((await columnNullability(db, dialect, table)).keys());
 }
 __name(existingColumns, "existingColumns");
+async function columnNullability(db, dialect, table) {
+  if (dialect === "sqlite") {
+    const rows2 = await querySql(db, dialect, sql.raw(`PRAGMA table_info("${table}")`));
+    return new Map(rows2.map((row) => [
+      String(row.name),
+      Number(row.notnull) === 0
+    ]));
+  }
+  const rows = await querySql(db, dialect, sql`select column_name as name, is_nullable as nullable from information_schema.columns where table_schema = ${sql.raw(dialect === "mysql" ? "database()" : "current_schema()")} and table_name = ${table}`);
+  return new Map(rows.map((row) => [
+    String(row.name ?? row.NAME),
+    String(row.nullable ?? row.NULLABLE).toUpperCase() === "YES"
+  ]));
+}
+__name(columnNullability, "columnNullability");
+async function relaxNotNull(db, dialect, spec, column, all) {
+  const q = /* @__PURE__ */ __name((identifier) => quote(dialect, identifier), "q");
+  if (dialect === "postgres") {
+    await runSql(db, dialect, sql.raw(`ALTER TABLE ${q(spec.name)} ALTER COLUMN ${q(column.name)} DROP NOT NULL`));
+    return;
+  }
+  if (dialect === "mysql") {
+    await runSql(db, dialect, sql.raw(`ALTER TABLE ${q(spec.name)} MODIFY ${columnDefinition(dialect, column)}`));
+    return;
+  }
+  if (all.some((other) => other.foreignKeys.some((key2) => key2.table === spec.name && other !== spec))) {
+    throw new Error(`agent schema: cannot make ${spec.name}.${column.name} nullable on SQLite \u2014 other tables reference ${spec.name}`);
+  }
+  const scratch = `${spec.name}__rebuild`;
+  const present = [
+    ...(await columnNullability(db, dialect, spec.name)).keys()
+  ];
+  const list = present.map(q).join(", ");
+  await runSql(db, dialect, sql.raw("SAVEPOINT agent_schema_rebuild"));
+  try {
+    await runSql(db, dialect, sql.raw(`DROP TABLE IF EXISTS ${q(scratch)}`));
+    await runSql(db, dialect, sql.raw(createTableSql(dialect, {
+      ...spec,
+      name: scratch
+    })));
+    await runSql(db, dialect, sql.raw(`INSERT INTO ${q(scratch)} (${list}) SELECT ${list} FROM ${q(spec.name)}`));
+    await runSql(db, dialect, sql.raw(`DROP TABLE ${q(spec.name)}`));
+    await runSql(db, dialect, sql.raw(`ALTER TABLE ${q(scratch)} RENAME TO ${q(spec.name)}`));
+    await runSql(db, dialect, sql.raw("RELEASE SAVEPOINT agent_schema_rebuild"));
+  } catch (error) {
+    await runSql(db, dialect, sql.raw("ROLLBACK TO SAVEPOINT agent_schema_rebuild")).catch(() => void 0);
+    await runSql(db, dialect, sql.raw("RELEASE SAVEPOINT agent_schema_rebuild")).catch(() => void 0);
+    throw error;
+  }
+}
+__name(relaxNotNull, "relaxNotNull");
 async function indexExists(db, dialect, table, index4) {
   let query;
   if (dialect === "sqlite") {
@@ -1656,7 +1777,8 @@ async function indexExists(db, dialect, table, index4) {
 }
 __name(indexExists, "indexExists");
 async function heal(db, dialect) {
-  for (const spec of specs(dialect)) {
+  const all = specs(dialect);
+  for (const spec of all) {
     await runSql(db, dialect, sql.raw(createTableSql(dialect, spec)));
     const present = await existingColumns(db, dialect, spec.name);
     for (const column of spec.columns) {
@@ -1664,6 +1786,13 @@ async function heal(db, dialect) {
         await runSql(db, dialect, sql.raw(`ALTER TABLE ${quote(dialect, spec.name)} ADD COLUMN ${columnDefinition(dialect, column)}`));
       }
     }
+    const nullable3 = await columnNullability(db, dialect, spec.name);
+    for (const column of spec.columns) {
+      if (column.notNull || spec.primaryKey.includes(column.name)) continue;
+      if (nullable3.get(column.name) === false) {
+        await relaxNotNull(db, dialect, spec, column, all);
+      }
+    }
     for (const index4 of spec.indexes) {
       if (!await indexExists(db, dialect, spec.name, index4.name)) {
         await runSql(db, dialect, sql.raw(createIndexSql(dialect, spec.name, index4)));
@@ -2437,7 +2566,7 @@ var DrizzleAgentStore = class {
       deletedAt: /* @__PURE__ */ new Date()
     }).where(eq2(this.t.agentThread.id, threadId));
   }
-  async forkThread(threadId, fromMessageId) {
+  async forkThread(threadId, fromMessageId, options = {}) {
     const [source] = await this.db.select().from(this.t.agentThread).where(eq2(this.t.agentThread.id, threadId));
     if (source === void 0) {
       throw new Error(`thread ${threadId} not found`);
@@ -2448,8 +2577,8 @@ var DrizzleAgentStore = class {
     const now = /* @__PURE__ */ new Date();
     const fork = {
       id: crypto.randomUUID(),
-      actorRef: source.actorRef,
-      tenantRef: source.tenantRef,
+      actorRef: options.actorRef ?? source.actorRef,
+      tenantRef: options.tenantRef !== void 0 ? options.tenantRef : source.tenantRef,
       title: source.title,
       transient: false,
       activeStreamId: null,
@@ -2854,7 +2983,13 @@ var DrizzleAgentStore = class {
       startedAt: /* @__PURE__ */ new Date(),
       settledAt: null,
       promptHash: run.promptHash ?? null,
-      parentRunId: run.parentRunId ?? null
+      parentRunId: run.parentRunId ?? null,
+      queueMs: run.queueMs !== void 0 ? Math.max(0, Math.round(run.queueMs)) : null,
+      modelMs: null,
+      toolsMs: null,
+      modelCalls: null,
+      toolCalls: null,
+      timings: null
     };
     await this.db.insert(this.t.agentRun).values(runRow);
   }
@@ -2882,6 +3017,13 @@ var DrizzleAgentStore = class {
     if (end.errorMessage !== void 0) {
       updates.errorMessage = end.errorMessage;
     }
+    if (end.timings !== void 0) {
+      updates.modelMs = Math.round(end.timings.modelMs);
+      updates.toolsMs = Math.round(end.timings.toolsMs);
+      updates.modelCalls = end.timings.modelCalls;
+      updates.toolCalls = end.timings.toolCalls;
+      updates.timings = end.timings.steps;
+    }
     await this.db.update(this.t.agentRun).set(updates).where(eq2(this.t.agentRun.id, end.runId));
   }
   /** Bump the run's llm-step retry counter. Atomic `retries = retries + 1`, no read-modify-write. */
@@ -3084,9 +3226,17 @@ var DrizzleAgentStore = class {
   async recordUsage(input) {
     await this.db.insert(this.t.agentTokenUsage).values({
       id: crypto.randomUUID(),
-      threadId: input.threadId,
+      threadId: input.threadId ?? null,
       actorRef: input.actorRef,
       messageId: input.messageId ?? null,
+      runId: input.runId ?? null,
+      feature: input.feature ?? null,
+      provider: input.provider ?? null,
+      tenantRef: input.tenantRef ?? null,
+      audioSeconds: input.audioSeconds ?? null,
+      characters: input.characters ?? null,
+      requests: input.requests ?? null,
+      metadata: input.metadata ?? null,
       modelId: input.modelId,
       purpose: input.purpose,
       inputTokens: input.usage.inputTokens,
@@ -3282,7 +3432,7 @@ function toQueuedMessage(row) {
 __name(toQueuedMessage, "toQueuedMessage");
 
 // src/drizzle-governance-queries.ts
-import { bucketByActor, bucketByModel, bucketByThread, bucketUsageTrend, dayBoundsUtc, rollupThreadUsage, summarizeGenuiCalls, truncateDetailContent } from "@dudousxd/nestjs-agent-core";
+import { bucketByActor, bucketByFeature, bucketByModel, bucketByThread, bucketUsageBreakdown, bucketUsageTrend, dayBoundsUtc, rollupThreadUsage, summarizeGenuiCalls, truncateDetailContent } from "@dudousxd/nestjs-agent-core";
 import { and as and3, asc as asc3, count, desc as desc2, eq as eq3, gte as gte2, inArray as inArray2, isNull as isNull3, like, lte as lte3, sql as sql4, sum } from "drizzle-orm";
 var UNCLASSIFIED_ERROR_CODE = "unknown";
 var DEFAULT_AGENT_BUCKET = "(default)";
@@ -3404,10 +3554,26 @@ function toRecentRunRow(run) {
     retries: run.retries,
     startedAt: run.startedAt.toISOString(),
     promptHash: run.promptHash ?? null,
-    parentRunId: run.parentRunId ?? null
+    parentRunId: run.parentRunId ?? null,
+    queueMs: run.queueMs ?? null,
+    modelMs: run.modelMs ?? null,
+    toolsMs: run.toolsMs ?? null,
+    modelCalls: run.modelCalls ?? null,
+    toolCalls: run.toolCalls ?? null
   };
 }
 __name(toRecentRunRow, "toRecentRunRow");
+function toRunTimings(run) {
+  if (!Array.isArray(run.timings)) return null;
+  return {
+    modelMs: run.modelMs ?? 0,
+    toolsMs: run.toolsMs ?? 0,
+    modelCalls: run.modelCalls ?? 0,
+    toolCalls: run.toolCalls ?? 0,
+    steps: run.timings
+  };
+}
+__name(toRunTimings, "toRunTimings");
 function toRunToolCallRow(toolCall) {
   return {
     toolCallId: toolCall.id,
@@ -3423,6 +3589,12 @@ function toRunToolCallRow(toolCall) {
 __name(toRunToolCallRow, "toRunToolCallRow");
 function toUsageInput(row) {
   return {
+    feature: row.feature ?? null,
+    purpose: row.purpose ?? null,
+    costSource: row.costSource ?? null,
+    audioSeconds: row.audioSeconds == null ? null : Number(row.audioSeconds),
+    characters: row.characters == null ? null : Number(row.characters),
+    requests: row.requests ?? null,
     modelId: row.modelId,
     actorRef: row.actorRef,
     threadId: row.threadId,
@@ -3470,7 +3642,12 @@ var DrizzleGovernanceQueries = class {
     }
     return pricing;
   }
-  async usageInRange(range) {
+  /**
+  * The usage rows in range, with only what the bucketers read: `requests` always (a recorded row
+  * counts as its own requests in every rollup), and the feature/unit columns only for the
+  * per-feature reads that group by them.
+  */
+  async usageInRange(range, withFeature = false) {
     const { start, end } = dayBoundsUtc(range);
     const usage = this.t.agentTokenUsage;
     const rows = await this.db.select({
@@ -3482,7 +3659,15 @@ var DrizzleGovernanceQueries = class {
       outputTokens: usage.outputTokens,
       cacheWriteTokens: usage.cacheWriteTokens,
       cacheReadTokens: usage.cacheReadTokens,
-      costUsd: usage.costUsd
+      costUsd: usage.costUsd,
+      requests: usage.requests,
+      ...withFeature ? {
+        feature: usage.feature,
+        purpose: usage.purpose,
+        costSource: usage.costSource,
+        audioSeconds: usage.audioSeconds,
+        characters: usage.characters
+      } : {}
     }).from(usage).where(and3(gte2(this.t.agentTokenUsage.createdAt, start), lte3(this.t.agentTokenUsage.createdAt, end)));
     return rows.map(toUsageInput);
   }
@@ -3490,6 +3675,14 @@ var DrizzleGovernanceQueries = class {
     const pricing = await this.loadPricing();
     return bucketByModel(await this.usageInRange(range), pricing);
   }
+  async spendByFeature(range) {
+    const pricing = await this.loadPricing();
+    return bucketByFeature(await this.usageInRange(range, true), pricing);
+  }
+  async usageBreakdown(range, filter) {
+    const pricing = await this.loadPricing();
+    return bucketUsageBreakdown(await this.usageInRange(range, true), pricing, filter);
+  }
   async spendByActor(range) {
     const pricing = await this.loadPricing();
     return bucketByActor(await this.usageInRange(range), pricing);
@@ -3503,7 +3696,9 @@ var DrizzleGovernanceQueries = class {
     const pricing = await this.loadPricing();
     const rows = await this.usageInRange(range);
     const threadIds = [
-      ...new Set(rows.map((row) => row.threadId))
+      ...new Set(rows.flatMap((row) => row.threadId === null ? [] : [
+        row.threadId
+      ]))
     ];
     if (threadIds.length === 0) {
       return [];
@@ -3581,6 +3776,7 @@ var DrizzleGovernanceQueries = class {
       outputTokens: sum(this.t.agentTokenUsage.outputTokens)
     }).from(this.t.agentTokenUsage).where(inArray2(this.t.agentTokenUsage.threadId, threadIds)).groupBy(this.t.agentTokenUsage.threadId);
     for (const row of usageTotals) {
+      if (row.threadId === null) continue;
       const bucket = rollups.get(row.threadId);
       if (bucket !== void 0) {
         bucket.totalTokens = Number(row.inputTokens ?? 0) + Number(row.outputTokens ?? 0);
@@ -3957,7 +4153,8 @@ var DrizzleGovernanceQueries = class {
         actorRef: row.thread.actorRef,
         deleted: row.thread.deletedAt != null
       },
-      toolCalls: toolCalls.map(toRunToolCallRow)
+      toolCalls: toolCalls.map(toRunToolCallRow),
+      timings: toRunTimings(row.run)
     };
   }
   /**