@open-mercato/ai-assistant 0.6.1-develop.3291.1.6fad645fd0 → 0.6.1

package/src/modules/ai_assistant/migrations/Migration20260508160000_ai_agent_loop_overrides.ts

@@ -0,0 +1,25 @@
+import { Migration } from '@mikro-orm/migrations';
+
+export class Migration20260508160000_ai_agent_loop_overrides extends Migration {
+
+  override async up(): Promise<void> {
+    this.addSql(`alter table "ai_agent_runtime_overrides" add column "loop_disabled" boolean null;`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" add column "loop_max_steps" int null;`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" add column "loop_max_tool_calls" int null;`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" add column "loop_max_wall_clock_ms" int null;`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" add column "loop_max_tokens" int null;`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" add column "loop_stop_when_json" jsonb null;`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" add column "loop_active_tools_json" jsonb null;`);
+  }
+
+  override async down(): Promise<void> {
+    this.addSql(`alter table "ai_agent_runtime_overrides" drop column "loop_disabled";`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" drop column "loop_max_steps";`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" drop column "loop_max_tool_calls";`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" drop column "loop_max_wall_clock_ms";`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" drop column "loop_max_tokens";`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" drop column "loop_stop_when_json";`);
+    this.addSql(`alter table "ai_agent_runtime_overrides" drop column "loop_active_tools_json";`);
+  }
+
+}

package/src/modules/ai_assistant/migrations/Migration20260508170000_ai_token_usage.ts

@@ -0,0 +1,89 @@
+import { Migration } from '@mikro-orm/migrations'
+
+export class Migration20260508170000_ai_token_usage extends Migration {
+
+  override async up(): Promise<void> {
+    this.addSql(`
+      create table "ai_token_usage_events" (
+        "id" uuid not null default gen_random_uuid(),
+        "tenant_id" uuid not null,
+        "organization_id" uuid null,
+        "user_id" uuid not null,
+        "agent_id" text not null,
+        "module_id" text not null,
+        "session_id" uuid not null,
+        "turn_id" uuid not null,
+        "step_index" int not null,
+        "provider_id" text not null,
+        "model_id" text not null,
+        "input_tokens" int not null,
+        "output_tokens" int not null,
+        "cached_input_tokens" int null,
+        "reasoning_tokens" int null,
+        "finish_reason" text null,
+        "loop_abort_reason" text null,
+        "created_at" timestamptz not null,
+        "updated_at" timestamptz not null,
+        constraint "ai_token_usage_events_pkey" primary key ("id")
+      );
+    `)
+    this.addSql(`
+      create index "ai_token_usage_events_tenant_created_idx"
+        on "ai_token_usage_events" ("tenant_id", "created_at" desc);
+    `)
+    this.addSql(`
+      create index "ai_token_usage_events_tenant_agent_created_idx"
+        on "ai_token_usage_events" ("tenant_id", "agent_id", "created_at" desc);
+    `)
+    this.addSql(`
+      create index "ai_token_usage_events_tenant_model_created_idx"
+        on "ai_token_usage_events" ("tenant_id", "model_id", "created_at" desc);
+    `)
+    this.addSql(`
+      create index "ai_token_usage_events_tenant_session_turn_step_idx"
+        on "ai_token_usage_events" ("tenant_id", "session_id", "turn_id", "step_index");
+    `)
+
+    this.addSql(`
+      create table "ai_token_usage_daily" (
+        "id" uuid not null default gen_random_uuid(),
+        "tenant_id" uuid not null,
+        "organization_id" uuid null,
+        "day" date not null,
+        "agent_id" text not null,
+        "model_id" text not null,
+        "provider_id" text not null,
+        "input_tokens" bigint not null default 0,
+        "output_tokens" bigint not null default 0,
+        "cached_input_tokens" bigint not null default 0,
+        "reasoning_tokens" bigint not null default 0,
+        "step_count" bigint not null default 0,
+        "turn_count" bigint not null default 0,
+        "session_count" bigint not null default 0,
+        "created_at" timestamptz not null,
+        "updated_at" timestamptz not null,
+        constraint "ai_token_usage_daily_pkey" primary key ("id")
+      );
+    `)
+    this.addSql(`
+      create unique index "ai_token_usage_daily_tenant_day_agent_model_org_uq"
+        on "ai_token_usage_daily" ("tenant_id", "day", "agent_id", "model_id", "organization_id")
+        where "organization_id" is not null;
+    `)
+    this.addSql(`
+      create unique index "ai_token_usage_daily_tenant_day_agent_model_null_org_uq"
+        on "ai_token_usage_daily" ("tenant_id", "day", "agent_id", "model_id")
+        where "organization_id" is null;
+    `)
+    this.addSql(`
+      create index "ai_token_usage_daily_tenant_day_idx"
+        on "ai_token_usage_daily" ("tenant_id", "day");
+    `)
+  }
+
+  override async down(): Promise<void> {
+    this.addSql(`drop table if exists "ai_token_usage_events" cascade;`)
+    this.addSql(`drop table if exists "ai_token_usage_daily" cascade;`)
+  }
+
+}

package/src/modules/ai_assistant/setup.ts

@@ -1,6 +1,7 @@
 import type { ModuleSetupConfig } from '@open-mercato/shared/modules/setup'
 
 const PENDING_ACTION_CLEANUP_SCHEDULE_ID = 'ai_assistant:pending-action-cleanup'
+const TOKEN_USAGE_PRUNE_SCHEDULE_ID = 'ai_assistant:token-usage-prune'
 
 /**
  * System-scoped recurring schedule: every 5 minutes, enqueue a job to the
@@ -49,6 +50,53 @@ async function ensurePendingActionCleanupSchedule(
   }
 }
 
+/**
+ * System-scoped daily schedule: enqueue a job to the `ai-token-usage-prune`
+ * queue to prune events older than the retention window and reconcile the
+ * daily rollup session counts.
+ *
+ * Phase 6.4 of spec `2026-04-28-ai-agents-agentic-loop-controls`.
+ */
+async function ensureTokenUsagePruneSchedule(
+  container: import('awilix').AwilixContainer | undefined,
+): Promise<void> {
+  if (!container) return
+  let schedulerService:
+    | {
+        register: (registration: Record<string, unknown>) => Promise<void>
+      }
+    | undefined
+  try {
+    schedulerService = container.resolve('schedulerService')
+  } catch {
+    schedulerService = undefined
+  }
+  if (!schedulerService) return
+  try {
+    await schedulerService.register({
+      id: TOKEN_USAGE_PRUNE_SCHEDULE_ID,
+      name: 'AI token-usage prune',
+      description:
+        'Delete ai_token_usage_events rows older than AI_TOKEN_USAGE_EVENTS_RETENTION_DAYS (default 90) and reconcile session_count on the daily rollup.',
+      scopeType: 'system',
+      scheduleType: 'interval',
+      scheduleValue: '24h',
+      timezone: 'UTC',
+      targetType: 'queue',
+      targetQueue: 'ai-token-usage-prune',
+      targetPayload: {},
+      sourceType: 'module',
+      sourceModule: 'ai_assistant',
+      isEnabled: true,
+    })
+  } catch (error) {
+    console.warn(
+      '[ai_assistant] Failed to register token-usage prune schedule:',
+      error instanceof Error ? error.message : error,
+    )
+  }
+}
+
 export const setup: ModuleSetupConfig = {
   defaultRoleFeatures: {
     admin: [
@@ -64,6 +112,7 @@ export const setup: ModuleSetupConfig = {
 
   async seedDefaults({ container }) {
     await ensurePendingActionCleanupSchedule(container)
+    await ensureTokenUsagePruneSchedule(container)
   },
 }
 

package/src/modules/ai_assistant/workers/__tests__/ai-token-usage-prune.test.ts

@@ -0,0 +1,144 @@
+import { runTokenUsagePrune } from '../ai-token-usage-prune'
+import type { EntityManager } from '@mikro-orm/postgresql'
+
+type ExecuteResult = { affectedRows?: number; rowCount?: number }
+type ConnectionExecuteSpy = jest.Mock<Promise<unknown[]> | Promise<ExecuteResult>>
+
+interface ConnectionStub {
+  execute: ConnectionExecuteSpy
+}
+
+function makeEmStub(connectionStub: ConnectionStub) {
+  return {
+    getConnection: () => connectionStub,
+  } as unknown as EntityManager
+}
+
+function makeConnectionStub(options: {
+  deleteBatches?: number[]
+  dailyRows?: Array<Record<string, unknown>>
+}): ConnectionStub {
+  const { deleteBatches = [], dailyRows = [] } = options
+  let deleteCallIndex = 0
+
+  const execute: ConnectionExecuteSpy = jest.fn(async (sql: string) => {
+    const normalized = sql.replace(/\s+/g, ' ').trim()
+
+    if (normalized.includes('delete from ai_token_usage_events')) {
+      const affected = deleteBatches[deleteCallIndex] ?? 0
+      deleteCallIndex++
+      return { affectedRows: affected } as unknown as ExecuteResult
+    }
+
+    if (normalized.includes('from ai_token_usage_daily')) {
+      return dailyRows
+    }
+
+    if (normalized.includes('update ai_token_usage_daily')) {
+      return [] as unknown[]
+    }
+
+    return [] as unknown[]
+  })
+
+  return { execute }
+}
+
+describe('runTokenUsagePrune', () => {
+  beforeEach(() => {
+    jest.clearAllMocks()
+  })
+
+  it('returns zero counts when there are no events or daily rows', async () => {
+    const connection = makeConnectionStub({ deleteBatches: [0], dailyRows: [] })
+    const em = makeEmStub(connection)
+    const summary = await runTokenUsagePrune({ em, now: new Date('2026-06-01T00:00:00Z'), retentionDays: 90 })
+    expect(summary.eventsDeleted).toBe(0)
+    expect(summary.dailyRowsReconciled).toBe(0)
+  })
+
+  it('deletes a single partial batch and stops looping', async () => {
+    const connection = makeConnectionStub({ deleteBatches: [42], dailyRows: [] })
+    const em = makeEmStub(connection)
+    const summary = await runTokenUsagePrune({ em, now: new Date('2026-06-01'), retentionDays: 90, batchSize: 5000 })
+    expect(summary.eventsDeleted).toBe(42)
+    const deleteCalls = (connection.execute as jest.Mock).mock.calls.filter(([sql]: [string]) =>
+      sql.includes('delete from ai_token_usage_events'),
+    )
+    expect(deleteCalls).toHaveLength(1)
+  })
+
+  it('loops until a partial batch completes multi-batch pruning', async () => {
+    const connection = makeConnectionStub({ deleteBatches: [100, 100, 50], dailyRows: [] })
+    const em = makeEmStub(connection)
+    const summary = await runTokenUsagePrune({ em, now: new Date('2026-06-01'), retentionDays: 90, batchSize: 100 })
+    expect(summary.eventsDeleted).toBe(250)
+    const deleteCalls = (connection.execute as jest.Mock).mock.calls.filter(([sql]: [string]) =>
+      sql.includes('delete from ai_token_usage_events'),
+    )
+    expect(deleteCalls).toHaveLength(3)
+  })
+
+  it('reconciles daily rows that are returned by the select', async () => {
+    const dailyRows = [
+      { id: 'row-1', computed_session_count: '3' },
+      { id: 'row-2', computed_session_count: 7 },
+    ]
+    const connection = makeConnectionStub({ deleteBatches: [0], dailyRows })
+    const em = makeEmStub(connection)
+    const summary = await runTokenUsagePrune({ em, now: new Date('2026-06-01'), retentionDays: 90 })
+    expect(summary.dailyRowsReconciled).toBe(2)
+    const updateCalls = (connection.execute as jest.Mock).mock.calls.filter(([sql]: [string]) =>
+      sql.includes('update ai_token_usage_daily'),
+    )
+    expect(updateCalls).toHaveLength(2)
+  })
+
+  it('does not throw when the delete query fails — returns zero deleted', async () => {
+    const badConnection: ConnectionStub = {
+      execute: jest.fn(async (sql: string) => {
+        if (sql.includes('delete from ai_token_usage_events')) {
+          throw new Error('DB connection error')
+        }
+        return []
+      }) as unknown as ConnectionExecuteSpy,
+    }
+    const em = makeEmStub(badConnection)
+    const consoleSpy = jest.spyOn(console, 'error').mockImplementation(() => {})
+    const summary = await runTokenUsagePrune({ em, now: new Date('2026-06-01'), retentionDays: 90 })
+    expect(summary.eventsDeleted).toBe(0)
+    consoleSpy.mockRestore()
+  })
+
+  it('does not throw when the reconcile query fails — returns zero reconciled', async () => {
+    const badConnection: ConnectionStub = {
+      execute: jest.fn(async (sql: string) => {
+        if (sql.includes('delete from ai_token_usage_events')) {
+          return [{ affectedRows: 0 }]
+        }
+        throw new Error('Reconcile error')
+      }) as unknown as ConnectionExecuteSpy,
+    }
+    const em = makeEmStub(badConnection)
+    const consoleSpy = jest.spyOn(console, 'error').mockImplementation(() => {})
+    const summary = await runTokenUsagePrune({ em, now: new Date('2026-06-01'), retentionDays: 90 })
+    expect(summary.dailyRowsReconciled).toBe(0)
+    consoleSpy.mockRestore()
+  })
+
+  it('uses AI_TOKEN_USAGE_EVENTS_RETENTION_DAYS env var when retentionDays is not supplied', async () => {
+    const original = process.env.AI_TOKEN_USAGE_EVENTS_RETENTION_DAYS
+    process.env.AI_TOKEN_USAGE_EVENTS_RETENTION_DAYS = '30'
+    const connection = makeConnectionStub({ deleteBatches: [0], dailyRows: [] })
+    const em = makeEmStub(connection)
+    await runTokenUsagePrune({ em, now: new Date('2026-06-01') })
+    const deleteCall = (connection.execute as jest.Mock).mock.calls.find(([sql]: [string]) =>
+      sql.includes('delete from ai_token_usage_events'),
+    )
+    expect(deleteCall).toBeDefined()
+    const cutoffParam = deleteCall![1][0] as Date
+    const cutoffIso = cutoffParam.toISOString().slice(0, 10)
+    expect(cutoffIso).toBe('2026-05-02')
+    process.env.AI_TOKEN_USAGE_EVENTS_RETENTION_DAYS = original
+  })
+})

package/src/modules/ai_assistant/workers/ai-token-usage-prune.ts

@@ -0,0 +1,188 @@
+import type { EntityManager } from '@mikro-orm/postgresql'
+import type { JobContext, QueuedJob, WorkerMeta } from '@open-mercato/queue'
+
+/** Default retention window for the ai_token_usage_events table (days). */
+const DEFAULT_RETENTION_DAYS = 90
+/** Batch size for delete operations to avoid long locks. */
+const DELETE_BATCH_SIZE = 5_000
+/** Number of trailing days to reconcile session_count against events table. */
+const RECONCILE_TRAILING_DAYS = 7
+
+export const metadata: WorkerMeta = {
+  queue: 'ai-token-usage-prune',
+  id: 'ai_assistant:token-usage-prune',
+  concurrency: 1,
+}
+
+export interface TokenUsagePruneRunOptions {
+  em: EntityManager
+  /** Override for deterministic tests. Defaults to `new Date()`. */
+  now?: Date
+  /** Override retention days (default from env). */
+  retentionDays?: number
+  /** Override batch size (default 5000). */
+  batchSize?: number
+}
+
+export interface TokenUsagePruneSummary {
+  eventsDeleted: number
+  dailyRowsReconciled: number
+}
+
+/**
+ * Reads `AI_TOKEN_USAGE_EVENTS_RETENTION_DAYS` env var (default 90).
+ */
+function resolveRetentionDays(): number {
+  const raw = process.env.AI_TOKEN_USAGE_EVENTS_RETENTION_DAYS
+  if (!raw) return DEFAULT_RETENTION_DAYS
+  const parsed = parseInt(raw.trim(), 10)
+  return !isNaN(parsed) && parsed > 0 ? parsed : DEFAULT_RETENTION_DAYS
+}
+
+/**
+ * Deletes events older than the retention cutoff in batches of `batchSize` to
+ * avoid long table locks. Returns total rows deleted.
+ */
+async function pruneOldEvents(
+  connection: ReturnType<EntityManager['getConnection']>,
+  cutoff: Date,
+  batchSize: number,
+): Promise<number> {
+  let totalDeleted = 0
+  for (;;) {
+    const result = await connection.execute(
+      `
+      delete from ai_token_usage_events
+      where id in (
+        select id from ai_token_usage_events
+        where created_at < ?
+        limit ?
+      )
+      `,
+      [cutoff, batchSize],
+      'run',
+    ) as { affectedRows?: number; rowCount?: number } | undefined
+    const deleted = result?.affectedRows ?? result?.rowCount ?? 0
+    totalDeleted += deleted
+    if (deleted < batchSize) break
+  }
+  return totalDeleted
+}
+
+/**
+ * Reconciles `session_count` on daily rollup rows for the trailing N days by
+ * recomputing it directly from the events table. This corrects any drift caused
+ * by out-of-order event delivery, retention pruning, or failed incremental
+ * writes.
+ */
+async function reconcileSessionCounts(
+  connection: ReturnType<EntityManager['getConnection']>,
+  now: Date,
+  trailingDays: number,
+): Promise<number> {
+  const trailingStart = new Date(now)
+  trailingStart.setUTCDate(trailingStart.getUTCDate() - trailingDays)
+  const from = trailingStart.toISOString().slice(0, 10)
+
+  // Recompute session_count for each (tenant_id, day, agent_id, model_id, org)
+  // combination by counting distinct session_ids from the events table.
+  const rows = await connection.execute(
+    `
+    select
+      d.id,
+      count(distinct e.session_id)::bigint as computed_session_count
+    from ai_token_usage_daily d
+    left join ai_token_usage_events e
+      on e.tenant_id = d.tenant_id
+      and e.agent_id = d.agent_id
+      and e.model_id = d.model_id
+      and date_trunc('day', e.created_at)::date = d.day
+      and (
+        (d.organization_id is null and e.organization_id is null)
+        or (d.organization_id is not null and e.organization_id = d.organization_id)
+      )
+    where d.day >= ?::date
+    group by d.id
+    `,
+    [from],
+    'all',
+  )
+
+  if (!Array.isArray(rows) || rows.length === 0) return 0
+
+  let reconciled = 0
+  for (const row of rows as Array<Record<string, unknown>>) {
+    const rowId = row.id as string
+    const computed = typeof row.computed_session_count === 'string'
+      ? parseInt(row.computed_session_count, 10)
+      : (row.computed_session_count as number) ?? 0
+    await connection.execute(
+      `update ai_token_usage_daily set session_count = ?, updated_at = now() where id = ?`,
+      [computed, rowId],
+      'run',
+    )
+    reconciled += 1
+  }
+
+  return reconciled
+}
+
+/**
+ * Core logic for the token-usage prune worker. Exported for unit testing.
+ *
+ * 1. Resolves the retention cutoff from `AI_TOKEN_USAGE_EVENTS_RETENTION_DAYS`.
+ * 2. Deletes events older than the cutoff in batches of 5_000.
+ * 3. Reconciles `session_count` on the daily rollup for trailing 7 days.
+ *
+ * Phase 6.4 of spec `2026-04-28-ai-agents-agentic-loop-controls`.
+ */
+export async function runTokenUsagePrune(
+  options: TokenUsagePruneRunOptions,
+): Promise<TokenUsagePruneSummary> {
+  const now = options.now ?? new Date()
+  const retentionDays = options.retentionDays ?? resolveRetentionDays()
+  const batchSize = options.batchSize ?? DELETE_BATCH_SIZE
+
+  const cutoff = new Date(now)
+  cutoff.setUTCDate(cutoff.getUTCDate() - retentionDays)
+
+  const connection = options.em.getConnection()
+
+  let eventsDeleted = 0
+  try {
+    eventsDeleted = await pruneOldEvents(connection, cutoff, batchSize)
+  } catch (error) {
+    console.error(
+      '[ai-token-usage-prune] Failed to prune old events:',
+      error instanceof Error ? error.message : error,
+    )
+  }
+
+  let dailyRowsReconciled = 0
+  try {
+    dailyRowsReconciled = await reconcileSessionCounts(connection, now, RECONCILE_TRAILING_DAYS)
+  } catch (error) {
+    console.error(
+      '[ai-token-usage-prune] Failed to reconcile session counts:',
+      error instanceof Error ? error.message : error,
+    )
+  }
+
+  console.info(
+    `[ai-token-usage-prune] Done. eventsDeleted=${eventsDeleted}, dailyRowsReconciled=${dailyRowsReconciled}, retentionDays=${retentionDays}.`,
+  )
+
+  return { eventsDeleted, dailyRowsReconciled }
+}
+
+type HandlerContext = JobContext & {
+  resolve: <T = unknown>(name: string) => T
+}
+
+export default async function handle(
+  _job: QueuedJob,
+  ctx: HandlerContext,
+): Promise<void> {
+  const em = ctx.resolve<EntityManager>('em')
+  await runTokenUsagePrune({ em })
+}