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

package/dist/index.cjs

@@ -240,7 +240,7 @@ var agentQueuedMessage = (0, import_sqlite_core.sqliteTable)("agent_queued_messa
 ]);
 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, {
+  threadId: (0, import_sqlite_core.text)("thread_id").references(() => agentThread.id, {
     onDelete: "cascade"
   }),
   actorRef: (0, import_sqlite_core.text)("actor_ref").notNull(),
@@ -257,13 +257,28 @@ var agentTokenUsage = (0, import_sqlite_core.sqliteTable)("agent_token_usage", {
   costUsd: (0, import_sqlite_core.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: (0, import_sqlite_core.text)("cost_source").$type(),
+  /** The run the row belongs to. */
+  runId: (0, import_sqlite_core.text)("run_id"),
+  /** What the spend was for: `agent:<name>` for an agent turn, the app's own otherwise; null → `agent`. */
+  feature: (0, import_sqlite_core.text)("feature"),
+  provider: (0, import_sqlite_core.text)("provider"),
+  tenantRef: (0, import_sqlite_core.text)("tenant_ref"),
+  /** Units a row bills by when they are not tokens. */
+  audioSeconds: (0, import_sqlite_core.real)("audio_seconds"),
+  characters: (0, import_sqlite_core.integer)("characters"),
+  requests: (0, import_sqlite_core.integer)("requests"),
+  /** The caller's own JSON, never read by the library. */
+  metadata: (0, import_sqlite_core.text)("metadata", {
+    mode: "json"
+  }).$type(),
   createdAt: (0, import_sqlite_core.integer)("created_at", {
     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_created_idx").on(table.createdAt),
-  (0, import_sqlite_core.index)("agent_token_usage_thread_idx").on(table.threadId)
+  (0, import_sqlite_core.index)("agent_token_usage_thread_idx").on(table.threadId),
+  (0, import_sqlite_core.index)("agent_token_usage_feature_created_idx").on(table.feature, table.createdAt)
 ]);
 var agentGenuiCall = (0, import_sqlite_core.sqliteTable)("agent_genui_call", {
   id: (0, import_sqlite_core.text)("id").primaryKey(),
@@ -345,7 +360,17 @@ var agentRun = (0, import_sqlite_core.sqliteTable)("agent_run", {
    * 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")
+  parentRunId: (0, import_sqlite_core.text)("parent_run_id"),
+  /** Enqueue → body start (ms); null when not read. */
+  queueMs: (0, import_sqlite_core.integer)("queue_ms"),
+  /** Where the completed run's time went: model and tool totals, counts, and the per-step list. */
+  modelMs: (0, import_sqlite_core.integer)("model_ms"),
+  toolsMs: (0, import_sqlite_core.integer)("tools_ms"),
+  modelCalls: (0, import_sqlite_core.integer)("model_calls"),
+  toolCalls: (0, import_sqlite_core.integer)("tool_calls"),
+  timings: (0, import_sqlite_core.text)("timings", {
+    mode: "json"
+  }).$type()
 }, (table) => [
   (0, import_sqlite_core.index)("agent_run_started_idx").on(table.startedAt)
 ]);
@@ -691,7 +716,7 @@ var agentQueuedMessage2 = (0, import_mysql_core.mysqlTable)("agent_queued_messag
 ]);
 var agentTokenUsage2 = (0, import_mysql_core.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(),
@@ -705,11 +730,26 @@ var agentTokenUsage2 = (0, import_mysql_core.mysqlTable)("agent_token_usage", {
   costUsd: (0, import_mysql_core.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: (0, import_mysql_core.double)("audio_seconds"),
+  characters: (0, import_mysql_core.bigint)("characters", {
+    mode: "number"
+  }),
+  requests: (0, import_mysql_core.int)("requests"),
+  /** The caller's own JSON, never read by the library. */
+  metadata: (0, import_mysql_core.json)("metadata").$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_created_idx").on(table.createdAt),
-  (0, import_mysql_core.index)("agent_token_usage_thread_idx").on(table.threadId)
+  (0, import_mysql_core.index)("agent_token_usage_thread_idx").on(table.threadId),
+  (0, import_mysql_core.index)("agent_token_usage_feature_created_idx").on(table.feature, table.createdAt)
 ]);
 var agentGenuiCall2 = (0, import_mysql_core.mysqlTable)("agent_genui_call", {
   id: key("id").primaryKey(),
@@ -781,7 +821,15 @@ var agentRun2 = (0, import_mysql_core.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: (0, import_mysql_core.int)("queue_ms"),
+  /** Where the completed run's time went: model and tool totals, counts, and the per-step list. */
+  modelMs: (0, import_mysql_core.int)("model_ms"),
+  toolsMs: (0, import_mysql_core.int)("tools_ms"),
+  modelCalls: (0, import_mysql_core.int)("model_calls"),
+  toolCalls: (0, import_mysql_core.int)("tool_calls"),
+  timings: (0, import_mysql_core.json)("timings").$type()
 }, (table) => [
   (0, import_mysql_core.index)("agent_run_started_idx").on(table.startedAt)
 ]);
@@ -1153,7 +1201,7 @@ var agentQueuedMessage3 = (0, import_pg_core.pgTable)("agent_queued_message", {
 ]);
 var agentTokenUsage3 = (0, import_pg_core.pgTable)("agent_token_usage", {
   id: (0, import_pg_core.text)("id").primaryKey(),
-  threadId: (0, import_pg_core.text)("thread_id").notNull().references(() => agentThread3.id, {
+  threadId: (0, import_pg_core.text)("thread_id").references(() => agentThread3.id, {
     onDelete: "cascade"
   }),
   actorRef: (0, import_pg_core.text)("actor_ref").notNull(),
@@ -1167,11 +1215,26 @@ var agentTokenUsage3 = (0, import_pg_core.pgTable)("agent_token_usage", {
   costUsd: (0, import_pg_core.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: (0, import_pg_core.text)("cost_source").$type(),
+  /** The run the row belongs to. */
+  runId: (0, import_pg_core.text)("run_id"),
+  /** What the spend was for: `agent:<name>` for an agent turn, the app's own otherwise; null → `agent`. */
+  feature: (0, import_pg_core.text)("feature"),
+  provider: (0, import_pg_core.text)("provider"),
+  tenantRef: (0, import_pg_core.text)("tenant_ref"),
+  /** Units a row bills by when they are not tokens. */
+  audioSeconds: (0, import_pg_core.doublePrecision)("audio_seconds"),
+  characters: (0, import_pg_core.bigint)("characters", {
+    mode: "number"
+  }),
+  requests: (0, import_pg_core.integer)("requests"),
+  /** The caller's own JSON, never read by the library. */
+  metadata: (0, import_pg_core.jsonb)("metadata").$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_created_idx").on(table.createdAt),
-  (0, import_pg_core.index)("agent_token_usage_thread_idx").on(table.threadId)
+  (0, import_pg_core.index)("agent_token_usage_thread_idx").on(table.threadId),
+  (0, import_pg_core.index)("agent_token_usage_feature_created_idx").on(table.feature, table.createdAt)
 ]);
 var agentGenuiCall3 = (0, import_pg_core.pgTable)("agent_genui_call", {
   id: (0, import_pg_core.text)("id").primaryKey(),
@@ -1243,7 +1306,15 @@ var agentRun3 = (0, import_pg_core.pgTable)("agent_run", {
   startedAt: ts2("started_at").notNull(),
   settledAt: ts2("settled_at"),
   promptHash: (0, import_pg_core.text)("prompt_hash"),
-  parentRunId: (0, import_pg_core.text)("parent_run_id")
+  parentRunId: (0, import_pg_core.text)("parent_run_id"),
+  /** Enqueue → body start (ms); null when not read. */
+  queueMs: (0, import_pg_core.integer)("queue_ms"),
+  /** Where the completed run's time went: model and tool totals, counts, and the per-step list. */
+  modelMs: (0, import_pg_core.integer)("model_ms"),
+  toolsMs: (0, import_pg_core.integer)("tools_ms"),
+  modelCalls: (0, import_pg_core.integer)("model_calls"),
+  toolCalls: (0, import_pg_core.integer)("tool_calls"),
+  timings: (0, import_pg_core.jsonb)("timings").$type()
 }, (table) => [
   (0, import_pg_core.index)("agent_run_started_idx").on(table.startedAt)
 ]);
@@ -1701,10 +1772,60 @@ function specs(dialect) {
 }
 __name(specs, "specs");
 async function existingColumns(db, dialect, table) {
-  const rows = dialect === "sqlite" ? await querySql(db, dialect, import_drizzle_orm2.sql.raw(`PRAGMA table_info("${table}")`)) : await querySql(db, dialect, import_drizzle_orm2.sql`select column_name as name from information_schema.columns where table_schema = ${import_drizzle_orm2.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, import_drizzle_orm2.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, import_drizzle_orm2.sql`select column_name as name, is_nullable as nullable from information_schema.columns where table_schema = ${import_drizzle_orm2.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, import_drizzle_orm2.sql.raw(`ALTER TABLE ${q(spec.name)} ALTER COLUMN ${q(column.name)} DROP NOT NULL`));
+    return;
+  }
+  if (dialect === "mysql") {
+    await runSql(db, dialect, import_drizzle_orm2.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, import_drizzle_orm2.sql.raw("SAVEPOINT agent_schema_rebuild"));
+  try {
+    await runSql(db, dialect, import_drizzle_orm2.sql.raw(`DROP TABLE IF EXISTS ${q(scratch)}`));
+    await runSql(db, dialect, import_drizzle_orm2.sql.raw(createTableSql(dialect, {
+      ...spec,
+      name: scratch
+    })));
+    await runSql(db, dialect, import_drizzle_orm2.sql.raw(`INSERT INTO ${q(scratch)} (${list}) SELECT ${list} FROM ${q(spec.name)}`));
+    await runSql(db, dialect, import_drizzle_orm2.sql.raw(`DROP TABLE ${q(spec.name)}`));
+    await runSql(db, dialect, import_drizzle_orm2.sql.raw(`ALTER TABLE ${q(scratch)} RENAME TO ${q(spec.name)}`));
+    await runSql(db, dialect, import_drizzle_orm2.sql.raw("RELEASE SAVEPOINT agent_schema_rebuild"));
+  } catch (error) {
+    await runSql(db, dialect, import_drizzle_orm2.sql.raw("ROLLBACK TO SAVEPOINT agent_schema_rebuild")).catch(() => void 0);
+    await runSql(db, dialect, import_drizzle_orm2.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") {
@@ -1718,7 +1839,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, import_drizzle_orm2.sql.raw(createTableSql(dialect, spec)));
     const present = await existingColumns(db, dialect, spec.name);
     for (const column of spec.columns) {
@@ -1726,6 +1848,13 @@ async function heal(db, dialect) {
         await runSql(db, dialect, import_drizzle_orm2.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, import_drizzle_orm2.sql.raw(createIndexSql(dialect, spec.name, index4)));
@@ -2916,7 +3045,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);
   }
@@ -2944,6 +3079,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((0, import_drizzle_orm4.eq)(this.t.agentRun.id, end.runId));
   }
   /** Bump the run's llm-step retry counter. Atomic `retries = retries + 1`, no read-modify-write. */
@@ -3146,9 +3288,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,
@@ -3466,10 +3616,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,
@@ -3485,6 +3651,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,
@@ -3532,7 +3704,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 } = (0, import_nestjs_agent_core5.dayBoundsUtc)(range);
     const usage = this.t.agentTokenUsage;
     const rows = await this.db.select({
@@ -3544,7 +3721,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((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);
   }
@@ -3552,6 +3737,14 @@ var DrizzleGovernanceQueries = class {
     const pricing = await this.loadPricing();
     return (0, import_nestjs_agent_core5.bucketByModel)(await this.usageInRange(range), pricing);
   }
+  async spendByFeature(range) {
+    const pricing = await this.loadPricing();
+    return (0, import_nestjs_agent_core5.bucketByFeature)(await this.usageInRange(range, true), pricing);
+  }
+  async usageBreakdown(range, filter) {
+    const pricing = await this.loadPricing();
+    return (0, import_nestjs_agent_core5.bucketUsageBreakdown)(await this.usageInRange(range, true), pricing, filter);
+  }
   async spendByActor(range) {
     const pricing = await this.loadPricing();
     return (0, import_nestjs_agent_core5.bucketByActor)(await this.usageInRange(range), pricing);
@@ -3565,7 +3758,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 [];
@@ -3643,6 +3838,7 @@ var DrizzleGovernanceQueries = class {
       outputTokens: (0, import_drizzle_orm5.sum)(this.t.agentTokenUsage.outputTokens)
     }).from(this.t.agentTokenUsage).where((0, import_drizzle_orm5.inArray)(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);
@@ -4019,7 +4215,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)
     };
   }
   /**