feat(memory): implement stability stage

Co-authored-by: multica-agent <github@multica.ai>

apps/web/server/utils/ai-provider.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from 'vitest'
-import { buildJsonCompletion, parseJsonCompletion, resolveAiProvider } from './ai-provider'
+import { buildJsonCompletion, parseAiUsage, parseJsonCompletion, resolveAiProvider } from './ai-provider'
 
 const request = { system: 'Return JSON.', messages: [{ role: 'user' as const, content: 'Create a world' }], schemaName: 'world', jsonSchema: { type: 'object' } }
 
@@ -22,4 +22,11 @@ describe('AI provider adapter', () => {
     expect(() => parseJsonCompletion({ choices: [] })).toThrow()
     expect(parseJsonCompletion({ choices: [{ message: { content: '{"ok":true}' } }] })).toEqual({ ok: true })
   })
+
+  it('normalizes provider usage telemetry', () => {
+    expect(parseAiUsage({
+      choices: [{ message: { content: '{}' } }],
+      usage: { prompt_tokens: 120, completion_tokens: 30, cost: 0.0042 },
+    })).toEqual({ inputTokens: 120, outputTokens: 30, costUsd: 0.0042 })
+  })
 })

apps/web/server/utils/ai-provider.ts

@@ -5,6 +5,14 @@ const CompletionResponseSchema = z.object({
   choices: z.array(z.object({
     message: z.object({ content: z.string().min(1) }),
   })).min(1),
+  usage: z.object({
+    prompt_tokens: z.number().nonnegative().nullish(),
+    completion_tokens: z.number().nonnegative().nullish(),
+    input_tokens: z.number().nonnegative().nullish(),
+    output_tokens: z.number().nonnegative().nullish(),
+    cost: z.number().nonnegative().nullish(),
+    total_cost: z.number().nonnegative().nullish(),
+  }).optional(),
 })
 
 export type AiProvider = z.infer<typeof ProviderSchema>
@@ -26,10 +34,18 @@ export interface JsonCompletionRequest {
   jsonSchema: Record<string, unknown>
 }
 
-interface JsonCompletion {
+export interface JsonCompletion {
   endpoint: string
   headers: Record<string, string>
   body: Record<string, unknown>
+  provider: AiProvider
+  model: string
+}
+
+export interface AiUsageSnapshot {
+  inputTokens: number
+  outputTokens: number
+  costUsd: number
 }
 
 const nonEmptyString = (value: unknown, name: string) => {
@@ -73,6 +89,8 @@ export function buildJsonCompletion(config: AiRuntimeConfig, request: JsonComple
   if (provider.provider === 'deepseek') {
     return {
       endpoint: provider.endpoint,
+      provider: provider.provider,
+      model: provider.model,
       headers: { Authorization: `Bearer ${provider.apiKey}`, 'Content-Type': 'application/json' },
       body: {
         ...common,
@@ -88,6 +106,8 @@ export function buildJsonCompletion(config: AiRuntimeConfig, request: JsonComple
 
   return {
     endpoint: provider.endpoint,
+    provider: provider.provider,
+    model: provider.model,
     headers: { Authorization: `Bearer ${provider.apiKey}`, 'Content-Type': 'application/json', 'X-Title': 'Dungeons & Ground' },
     body: {
       ...common,
@@ -101,3 +121,12 @@ export function parseJsonCompletion(payload: unknown): unknown {
   const completion = CompletionResponseSchema.parse(payload)
   return JSON.parse(completion.choices[0]!.message.content)
 }
+
+export function parseAiUsage(payload: unknown): AiUsageSnapshot {
+  const completion = CompletionResponseSchema.parse(payload)
+  return {
+    inputTokens: Math.max(0, Math.trunc(completion.usage?.prompt_tokens ?? completion.usage?.input_tokens ?? 0)),
+    outputTokens: Math.max(0, Math.trunc(completion.usage?.completion_tokens ?? completion.usage?.output_tokens ?? 0)),
+    costUsd: Math.max(0, completion.usage?.cost ?? completion.usage?.total_cost ?? 0),
+  }
+}

apps/web/server/utils/stage-two-supabase.ts

@@ -1,6 +1,8 @@
 import { z } from 'zod'
+import { randomUUID } from 'node:crypto'
 import type { H3Event } from 'h3'
 import { moderate13Plus } from '@dng/shared'
+import { parseAiUsage, type JsonCompletion } from './ai-provider'
 
 const UuidSchema = z.string().uuid()
 const AuthEmailSchema = z.preprocess(
@@ -87,6 +89,41 @@ export function stageTwoRpc<T>(name: string, body: Record<string, unknown>): Pro
   return stageTwoDatabase<T>(`rpc/${name}`, { method: 'POST', body: JSON.stringify(body) })
 }
 
+export function stageSixConsumeAiQuota(userId: string, campaignId: string | null, requestKind: string) {
+  return stageTwoRpc<{ userRemaining: number; campaignRemaining: number | null; tokenRemaining: number }>('stage_six_consume_ai_quota', {
+    p_user_id: userId,
+    p_campaign_id: campaignId,
+    p_request_kind: requestKind,
+  })
+}
+
+export async function stageSixRecordAiUsage(input: {
+  userId: string
+  campaignId: string | null
+  requestKind: string
+  completion: JsonCompletion
+  payload: unknown
+  latencyMs: number
+}): Promise<void> {
+  const usage = parseAiUsage(input.payload)
+  await stageTwoDatabase('ai_usage', {
+    method: 'POST',
+    prefer: 'return=minimal',
+    body: JSON.stringify({
+      usage_key: `${input.requestKind}:${randomUUID()}`,
+      user_id: input.userId,
+      campaign_id: input.campaignId,
+      provider: input.completion.provider,
+      request_kind: input.requestKind,
+      model: input.completion.model,
+      input_tokens: usage.inputTokens,
+      output_tokens: usage.outputTokens,
+      cost_usd: usage.costUsd,
+      latency_ms: Math.max(0, Math.trunc(input.latencyMs)),
+    }),
+  })
+}
+
 export async function requireStageTwoUser(event: H3Event): Promise<StageTwoUser> {
   const authorization = getHeader(event, 'authorization')
   if (!authorization?.startsWith('Bearer ') || authorization.length <= 7) {
@@ -159,7 +196,8 @@ export function stageTwoApiError(error: unknown): never {
       })
     }
     const conflict = /duplicate key|already exists|not claimable/i.test(error.message)
-    throw createError({ statusCode: conflict ? 409 : error.status >= 500 ? 502 : 400, statusMessage: error.message })
+    const quota = /AI (?:rate limit|.*quota)/i.test(error.message)
+    throw createError({ statusCode: quota ? 429 : conflict ? 409 : error.status >= 500 ? 502 : 400, statusMessage: error.message })
   }
   if (error instanceof z.ZodError) {
     throw createError({ statusCode: 400, statusMessage: error.issues.map(issue => issue.message).join('; ') })