fix(ai-billing): bill thread title generation and tool-call repair (#20067)

## Summary Two `generateText` call sites were unbilled — identified during the 2026-04-26 incident audit: - **Thread title generation** (`AgentTitleGenerationService.generateThreadTitle`) — fires once per new chat thread; low-volume but completeness matters. - **Tool-call repair** (`repairToolCall` util, used inside `experimental_repairToolCall` callbacks) — can fire `MAX_STEPS` times per agent turn if a model gets stuck producing malformed tool calls. ## What changed - `AgentTitleGenerationService` — inject `AiBillingService`, expand `generateThreadTitle` signature to accept `workspaceId` and `userWorkspaceId`, bill in a `finally` block with `UsageOperationType.AI_CHAT_TOKEN`. Restructured so the no-default-model short-circuit happens before the `try`, avoiding a fake billing call. - `repair-tool-call.util.ts` — added an optional `billingContext` arg containing `aiBillingService`, `modelId`, `workspaceId`, `userWorkspaceId`, `operationType`. Wraps the `generateText` call in `try/finally`; bills in finally with the operation type provided by the caller. Optional so existing callers keep compiling. - `chat-execution.service.ts` — threads `billingContext` (`AI_CHAT_TOKEN`) into the repair callback. - `agent-chat.service.ts` — passes `workspaceId` and `thread.userWorkspaceId` to `generateThreadTitle`. - `agent-async-executor.service.ts` — TODO comment marking that the repair callback should thread billing once `executeAgent` accepts `workspaceId` (depends on the executeAgent-finally PR). ## Follow-up After the executeAgent-finally PR lands, `workspaceId` is in scope at the `experimental_repairToolCall` callback in `agent-async-executor.service.ts:198`. Thread `billingContext` through to fire repair-call billing for workflow agents too, and remove the TODO. Tiny follow-up. ## Test plan - [ ] Create a new chat thread; verify a `usageEvent` row is written for the title generation call. - [ ] Send a chat message that triggers tool-call repair (e.g. force a malformed tool call); verify a `usageEvent` row for the repair sub-call. ## Notes for review - `billingContext` arg made **optional** to allow incremental rollout — the executeAgent-finally PR plus a small follow-up will close the agent-async-executor side. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-27 20:54:40 +02:00
parent 875795cc30
commit 5b3ee3f1a7
4 changed files with 102 additions and 14 deletions
@@ -1,6 +1,17 @@
-import { type LanguageModel, NoSuchToolError, Output, generateText } from 'ai';
+import {
+  type LanguageModel,
+  type LanguageModelUsage,
+  NoSuchToolError,
+  Output,
+  type StepResult,
+  type ToolSet,
+  generateText,
+} from 'ai';
 import { type z } from 'zod';

+import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum';
+import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service';
+import { extractCacheCreationTokensFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util';
 import { AI_TELEMETRY_CONFIG } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-telemetry.const';

 type ToolCall = {
@@ -10,18 +21,28 @@ type ToolCall = {
  input: string;
 };

+type RepairToolCallBillingContext = {
+  aiBillingService: AiBillingService;
+  modelId: string;
+  workspaceId: string;
+  userWorkspaceId: string | null;
+  operationType: UsageOperationType;
+};
+
 export const repairToolCall = async ({
  toolCall,
  tools,
  inputSchema,
  error,
  model,
+  billingContext,
 }: {
  toolCall: ToolCall;
  tools: Record<string, unknown>;
  inputSchema: (toolCall: { toolName: string }) => unknown;
  error: Error;
  model: LanguageModel;
+  billingContext?: RepairToolCallBillingContext;
 }): Promise<ToolCall | null> => {
  // Don't attempt to fix invalid tool names
  if (NoSuchToolError.isInstance(error)) {
@@ -40,8 +61,11 @@ export const repairToolCall = async ({
    return null;
  }

+  let usage: LanguageModelUsage | undefined;
+  let steps: StepResult<ToolSet>[] | undefined;
+
  try {
-    const { output: repairedInput } = await generateText({
+    const result = await generateText({
      model,
      output: Output.object({ schema: schema as z.ZodTypeAny }),
      prompt: [
@@ -62,6 +86,11 @@ export const repairToolCall = async ({
      experimental_telemetry: AI_TELEMETRY_CONFIG,
    });

+    usage = result.usage;
+    steps = result.steps;
+
+    const repairedInput = result.output;
+
    if (repairedInput == null) {
      return null;
    }
@@ -75,5 +104,20 @@ export const repairToolCall = async ({
  } catch {
    // If repair fails, return null to let the error propagate
    return null;
+  } finally {
+    if (billingContext && usage) {
+      const cacheCreationTokens = steps
+        ? extractCacheCreationTokensFromSteps(steps)
+        : 0;
+
+      billingContext.aiBillingService.calculateAndBillUsage(
+        billingContext.modelId,
+        { usage, cacheCreationTokens },
+        billingContext.workspaceId,
+        billingContext.operationType,
+        null,
+        billingContext.userWorkspaceId,
+      );
+    }
  }
 };
@@ -328,8 +328,11 @@ export class AgentChatService {
      return null;
    }

-    const title =
-      await this.titleGenerationService.generateThreadTitle(messageContent);
+    const title = await this.titleGenerationService.generateThreadTitle(
+      messageContent,
+      workspaceId,
+      thread.userWorkspaceId,
+    );

    await this.threadRepository.update(threadId, { title });

@@ -1,7 +1,15 @@
 import { Injectable, Logger } from '@nestjs/common';

-import { generateText } from 'ai';
+import {
+  type LanguageModelUsage,
+  type StepResult,
+  type ToolSet,
+  generateText,
+} from 'ai';

+import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-operation-type.enum';
+import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service';
+import { extractCacheCreationTokensFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util';
 import { AI_TELEMETRY_CONFIG } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-telemetry.const';
 import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service';

@@ -11,29 +19,55 @@ export class AgentTitleGenerationService {

  constructor(
    private readonly aiModelRegistryService: AiModelRegistryService,
+    private readonly aiBillingService: AiBillingService,
  ) {}

-  async generateThreadTitle(messageContent: string): Promise<string> {
+  async generateThreadTitle(
+    messageContent: string,
+    workspaceId: string,
+    userWorkspaceId: string | null,
+  ): Promise<string> {
+    const defaultModel = this.aiModelRegistryService.getDefaultSpeedModel();
+
+    if (!defaultModel) {
+      this.logger.warn('No default AI model available for title generation');
+
+      return this.generateFallbackTitle(messageContent);
+    }
+
+    let usage: LanguageModelUsage | undefined;
+    let steps: StepResult<ToolSet>[] | undefined;
+
    try {
-      const defaultModel = this.aiModelRegistryService.getDefaultSpeedModel();
-
-      if (!defaultModel) {
-        this.logger.warn('No default AI model available for title generation');
-
-        return this.generateFallbackTitle(messageContent);
-      }
-
      const result = await generateText({
        model: defaultModel.model,
        prompt: `Generate a concise, descriptive title (maximum 60 characters) for a chat thread based on the following message. The title should capture the main topic or purpose of the conversation. Return only the title, nothing else. Message: "${messageContent}"`,
        experimental_telemetry: AI_TELEMETRY_CONFIG,
      });

+      usage = result.usage;
+      steps = result.steps;
+
      return this.cleanTitle(result.text);
    } catch (error) {
      this.logger.error('Failed to generate title with AI:', error);

      return this.generateFallbackTitle(messageContent);
+    } finally {
+      if (usage) {
+        const cacheCreationTokens = steps
+          ? extractCacheCreationTokensFromSteps(steps)
+          : 0;
+
+        this.aiBillingService.calculateAndBillUsage(
+          defaultModel.modelId,
+          { usage, cacheCreationTokens },
+          workspaceId,
+          UsageOperationType.AI_CHAT_TOKEN,
+          null,
+          userWorkspaceId,
+        );
+      }
    }
  }

@@ -348,6 +348,13 @@ export class ChatExecutionService {
          inputSchema,
          error,
          model: registeredModel.model,
+          billingContext: {
+            aiBillingService: this.aiBillingService,
+            modelId: registeredModel.modelId,
+            workspaceId: workspace.id,
+            userWorkspaceId,
+            operationType: UsageOperationType.AI_CHAT_TOKEN,
+          },
        });
      },
    });