feat(credits): charge credits per chat turn (atomic wallet debit + audit)

app/lib/workflows/__tests__/runAgentWorkflow.test.ts

@@ -5,6 +5,7 @@ import { clearChatActiveStream } from "@/lib/chat/clearChatActiveStream";
 import { closeChatStream } from "@/app/lib/workflows/closeChatStream";
 import { generateAssistantMessageId } from "@/app/lib/workflows/generateAssistantMessageId";
 import { persistAssistantMessage } from "@/lib/chat/persistAssistantMessage";
+import { handleChatCredits } from "@/lib/credits/handleChatCredits";
 
 vi.mock("@/app/lib/workflows/runAgentStep", () => ({
   runAgentStep: vi.fn(),
@@ -21,6 +22,9 @@ vi.mock("@/app/lib/workflows/generateAssistantMessageId", () => ({
 vi.mock("@/lib/chat/persistAssistantMessage", () => ({
   persistAssistantMessage: vi.fn(),
 }));
+vi.mock("@/lib/credits/handleChatCredits", () => ({
+  handleChatCredits: vi.fn(),
+}));
 // Captured writable stub so tests can assert closeChatStream got the
 // same instance the workflow body holds.
 const writableStub = new WritableStream();
@@ -43,6 +47,7 @@ const baseInput = {
   messages: [{ id: "m1", role: "user", parts: [{ type: "text", text: "hi" }] } as never],
   chatId: "chat-1",
   sessionId: "session-1",
+  accountId: "acc-1",
   modelId: "anthropic/claude-haiku-4.5",
   agentContext: {
     sandbox: { state: { type: "vercel" }, workingDirectory: "/sandbox/mono" },
@@ -170,4 +175,82 @@ describe("runAgentWorkflow", () => {
       expect.objectContaining({ assistantMessageId: "asst-in-progress" }),
     );
   });
+
+  it("calls handleChatCredits with the gateway cost + token usage from responseMessage.metadata", async () => {
+    const responseMessage = {
+      id: "assistant-msg-xyz",
+      role: "assistant",
+      parts: [{ type: "text", text: "Hello!" }],
+      metadata: {
+        totalMessageCost: 0.07,
+        totalMessageUsage: {
+          inputTokens: 100,
+          cachedInputTokens: 10,
+          outputTokens: 20,
+        },
+      },
+    };
+    vi.mocked(runAgentStep).mockResolvedValue({
+      finishReason: "stop",
+      responseMessage: responseMessage as never,
+    });
+
+    await runAgentWorkflow(baseInput);
+
+    expect(handleChatCredits).toHaveBeenCalledTimes(1);
+    expect(handleChatCredits).toHaveBeenCalledWith({
+      accountId: "acc-1",
+      model: "anthropic/claude-haiku-4.5",
+      source: "api",
+      gatewayCostUsd: 0.07,
+      usage: {
+        inputTokens: 100,
+        cachedInputTokens: 10,
+        outputTokens: 20,
+      },
+    });
+  });
+
+  it("calls handleChatCredits with zero usage when metadata is missing (lets the 1c floor apply)", async () => {
+    const responseMessage = {
+      id: "assistant-msg-xyz",
+      role: "assistant",
+      parts: [{ type: "text", text: "Hello!" }],
+      // no metadata
+    };
+    vi.mocked(runAgentStep).mockResolvedValue({
+      finishReason: "stop",
+      responseMessage: responseMessage as never,
+    });
+
+    await runAgentWorkflow(baseInput);
+
+    expect(handleChatCredits).toHaveBeenCalledTimes(1);
+    expect(handleChatCredits).toHaveBeenCalledWith({
+      accountId: "acc-1",
+      model: "anthropic/claude-haiku-4.5",
+      source: "api",
+      gatewayCostUsd: undefined,
+      usage: { inputTokens: 0, cachedInputTokens: 0, outputTokens: 0 },
+    });
+  });
+
+  it("does NOT call handleChatCredits when runAgentStep returns no responseMessage", async () => {
+    vi.mocked(runAgentStep).mockResolvedValue({
+      finishReason: "stop",
+      responseMessage: undefined,
+    });
+
+    await runAgentWorkflow(baseInput);
+
+    expect(handleChatCredits).not.toHaveBeenCalled();
+  });
+
+  it("does NOT call handleChatCredits when runAgentStep throws (no message to bill)", async () => {
+    vi.mocked(runAgentStep).mockRejectedValue(new Error("model exploded"));
+
+    await expect(runAgentWorkflow(baseInput)).rejects.toThrow("model exploded");
+
+    expect(handleChatCredits).not.toHaveBeenCalled();
+  });
 });

app/lib/workflows/runAgentWorkflow.ts

@@ -1,16 +1,31 @@
 import { getWorkflowMetadata, getWritable } from "workflow";
-import type { UIMessage, UIMessageChunk } from "ai";
+import type { LanguageModelUsage, UIMessage, UIMessageChunk } from "ai";
 import { closeChatStream } from "@/app/lib/workflows/closeChatStream";
 import { generateAssistantMessageId } from "@/app/lib/workflows/generateAssistantMessageId";
 import { runAgentStep } from "@/app/lib/workflows/runAgentStep";
 import { clearChatActiveStream } from "@/lib/chat/clearChatActiveStream";
 import { persistAssistantMessage } from "@/lib/chat/persistAssistantMessage";
+import { handleChatCredits } from "@/lib/credits/handleChatCredits";
+import type { AgentMessageMetadata } from "@/lib/agent/messageMetadata/AgentMessageMetadata";
 import type { DurableAgentContext } from "@/lib/agent/tools/AgentContext";
 
+const ZERO_USAGE: LanguageModelUsage = {
+  inputTokens: 0,
+  cachedInputTokens: 0,
+  outputTokens: 0,
+} as LanguageModelUsage;
+
 export type RunAgentWorkflowInput = {
   messages: UIMessage[];
   chatId: string;
   sessionId: string;
+  /**
+   * Authenticated account whose wallet absorbs the turn's cost. Resolved by
+   * the route handler via `validateChatWorkflow` so we never trust a
+   * caller-supplied id. Threaded into `recordChatUsage` after the assistant
+   * message is persisted.
+   */
+  accountId: string;
   modelId: string;
   /**
    * JSON-serializable subset of AgentContext that survives the durable
@@ -82,6 +97,24 @@ export async function runAgentWorkflow(input: RunAgentWorkflowInput): Promise<vo
     // mark the workflow run failed.
     if (result.responseMessage) {
       await persistAssistantMessage(input.chatId, result.responseMessage);
+
+      // Charge the account for this turn. Atomic wallet debit + audit
+      // row insert via the `deduct_credits_with_audit` Postgres function
+      // (wired into `handleChatCredits` → `recordCreditDeduction`).
+      // Fire-and-forget by contract; transient credits-table failures
+      // must not abort the chat workflow. Mirrors open-agents'
+      // `recordWorkflowUsage` main-agent path
+      // (apps/web/app/workflows/chat-post-finish.ts) and reuses the
+      // same `handleChatCredits` orchestrator that `handleChatStream`
+      // already uses for the non-workflow chat path.
+      const metadata = result.responseMessage.metadata as AgentMessageMetadata | undefined;
+      await handleChatCredits({
+        accountId: input.accountId,
+        model: input.modelId,
+        source: "api",
+        gatewayCostUsd: metadata?.totalMessageCost,
+        usage: metadata?.totalMessageUsage ?? ZERO_USAGE,
+      });
     }
   } finally {
     // Run two cleanup steps in parallel:

lib/chat/__tests__/integration/chatEndToEnd.test.ts

@@ -535,6 +535,7 @@ describe("Chat Integration Tests", () => {
       expect(mockGetCreditUsage).toHaveBeenCalledWith(
         { promptTokens: 1000, completionTokens: 500 },
         "gpt-4",
+        undefined,
       );
       expect(mockRecordCreditDeduction).toHaveBeenCalledWith(
         expect.objectContaining({

lib/chat/handleChatWorkflowStream.ts

@@ -122,6 +122,7 @@ export async function handleChatWorkflowStream(request: NextRequest): Promise<Re
       messages: validated.messages,
       chatId: validated.chatId,
       sessionId: validated.sessionId,
+      accountId: validated.accountId,
       modelId,
       agentContext: {
         sandbox: {

lib/credits/__tests__/getCreditUsage.test.ts

@@ -139,4 +139,80 @@ describe("getCreditUsage", () => {
       expect(cost).toBe(0);
     });
   });
+
+  describe("gateway cost short-circuit", () => {
+    it("returns gatewayCostUsd directly when it is a positive number (skips catalog lookup)", async () => {
+      const cost = await getCreditUsage(
+        { inputTokens: 1000, outputTokens: 500 },
+        "anthropic/claude-haiku-4.5",
+        0.07,
+      );
+      expect(cost).toBe(0.07);
+      expect(mockGetModel).not.toHaveBeenCalled();
+    });
+
+    it("falls through to token math when gatewayCostUsd is undefined", async () => {
+      mockGetModel.mockResolvedValue({
+        id: "gpt-4",
+        pricing: { input: "0.00003", output: "0.00006" },
+      } as any);
+
+      const cost = await getCreditUsage(
+        { inputTokens: 1000, outputTokens: 500 },
+        "gpt-4",
+        undefined,
+      );
+      // 1000 * 0.00003 + 500 * 0.00006 = 0.06
+      expect(cost).toBeCloseTo(0.06);
+      expect(mockGetModel).toHaveBeenCalledWith("gpt-4");
+    });
+
+    it("falls through to token math when gatewayCostUsd is 0", async () => {
+      mockGetModel.mockResolvedValue({
+        id: "gpt-4",
+        pricing: { input: "0.00003", output: "0.00006" },
+      } as any);
+
+      const cost = await getCreditUsage({ inputTokens: 1000, outputTokens: 500 }, "gpt-4", 0);
+      expect(cost).toBeCloseTo(0.06);
+    });
+
+    it("falls through to token math when gatewayCostUsd is negative", async () => {
+      mockGetModel.mockResolvedValue({
+        id: "gpt-4",
+        pricing: { input: "0.00003", output: "0.00006" },
+      } as any);
+
+      const cost = await getCreditUsage({ inputTokens: 1000, outputTokens: 500 }, "gpt-4", -1);
+      expect(cost).toBeCloseTo(0.06);
+    });
+
+    it("falls through to token math when gatewayCostUsd is NaN", async () => {
+      mockGetModel.mockResolvedValue({
+        id: "gpt-4",
+        pricing: { input: "0.00003", output: "0.00006" },
+      } as any);
+
+      const cost = await getCreditUsage(
+        { inputTokens: 1000, outputTokens: 500 },
+        "gpt-4",
+        Number.NaN,
+      );
+      expect(cost).toBeCloseTo(0.06);
+    });
+
+    it("falls through to token math when gatewayCostUsd is Infinity", async () => {
+      mockGetModel.mockResolvedValue({
+        id: "gpt-4",
+        pricing: { input: "0.00003", output: "0.00006" },
+      } as any);
+
+      const cost = await getCreditUsage(
+        { inputTokens: 1000, outputTokens: 500 },
+        "gpt-4",
+        Number.POSITIVE_INFINITY,
+      );
+      expect(cost).toBeCloseTo(0.06);
+    });
+  });
 });

lib/credits/__tests__/handleChatCredits.test.ts

@@ -42,7 +42,7 @@ describe("handleChatCredits", () => {
         accountId: "account-123",
       });
 
-      expect(mockGetCreditUsage).toHaveBeenCalledWith(USAGE, "gpt-4");
+      expect(mockGetCreditUsage).toHaveBeenCalledWith(USAGE, "gpt-4", undefined);
       expect(mockRecordCreditDeduction).toHaveBeenCalledWith({
         accountId: "account-123",
         creditsToDeduct: 5,
@@ -154,4 +154,51 @@ describe("handleChatCredits", () => {
       expect(consoleSpy).toHaveBeenCalledWith("Failed to handle chat credits:", expect.any(Error));
     });
   });
+
+  describe("gateway cost + source extensions", () => {
+    it("forwards gatewayCostUsd to getCreditUsage when provided", async () => {
+      mockGetCreditUsage.mockResolvedValue(0.07);
+      mockRecordCreditDeduction.mockResolvedValue({ success: true, newBalance: 93 });
+
+      await handleChatCredits({
+        usage: USAGE,
+        model: "anthropic/claude-haiku-4.5",
+        accountId: "account-123",
+        gatewayCostUsd: 0.07,
+      });
+
+      expect(mockGetCreditUsage).toHaveBeenCalledWith(USAGE, "anthropic/claude-haiku-4.5", 0.07);
+    });
+
+    it("defaults source to 'web' when not provided (backwards compatible)", async () => {
+      mockGetCreditUsage.mockResolvedValue(0.05);
+      mockRecordCreditDeduction.mockResolvedValue({ success: true, newBalance: 95 });
+
+      await handleChatCredits({
+        usage: USAGE,
+        model: "gpt-4",
+        accountId: "account-123",
+      });
+
+      expect(mockRecordCreditDeduction).toHaveBeenCalledWith(
+        expect.objectContaining({ source: "web" }),
+      );
+    });
+
+    it("propagates source='api' when caller is the chat workflow", async () => {
+      mockGetCreditUsage.mockResolvedValue(0.05);
+      mockRecordCreditDeduction.mockResolvedValue({ success: true, newBalance: 95 });
+
+      await handleChatCredits({
+        usage: USAGE,
+        model: "anthropic/claude-haiku-4.5",
+        accountId: "account-123",
+        source: "api",
+      });
+
+      expect(mockRecordCreditDeduction).toHaveBeenCalledWith(
+        expect.objectContaining({ source: "api" }),
+      );
+    });
+  });
 });