Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 26 additions & 4 deletions app/lib/workflows/__tests__/runAgentWorkflow.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ describe("runAgentWorkflow", () => {
});
});

it("does NOT call handleChatCredits when runAgentStep returns no responseMessage", async () => {
it("STILL bills the turn (1c floor) when runAgentStep returns no responseMessage — a mid-turn side-effect (e.g. send_email) may already have fired", async () => {
vi.mocked(runAgentStep).mockResolvedValue({
finishReason: "stop",
aborted: false,
Expand All @@ -271,15 +271,37 @@ describe("runAgentWorkflow", () => {

await runAgentWorkflow(baseInput);

expect(handleChatCredits).not.toHaveBeenCalled();
// Previously this path skipped billing entirely — a turn that emailed
// the customer but returned no responseMessage was free + left no
// usage_events audit row. Now it bills once with zero usage → 1c floor
// + a model_id row.
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 throws (no message to bill)", async () => {
it("STILL bills the turn (1c floor) when runAgentStep throws after a side-effect — email out must never escape the wallet debit + audit row", async () => {
vi.mocked(runAgentStep).mockRejectedValue(new Error("model exploded"));

// The error still propagates (workflow fails)...
await expect(runAgentWorkflow(baseInput)).rejects.toThrow("model exploded");

expect(handleChatCredits).not.toHaveBeenCalled();
// ...but the `finally` backstop bills the turn first, so a turn that ran
// a customer-facing tool (send_email) then threw is charged + audited
// instead of emailing-but-not-billing.
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 },
});
});

describe("auto-commit", () => {
Expand Down
66 changes: 44 additions & 22 deletions app/lib/workflows/runAgentWorkflow.ts
Original file line number Diff line number Diff line change
Expand Up @@ -97,35 +97,47 @@ export async function runAgentWorkflow(input: RunAgentWorkflowInput): Promise<vo
const assistantMessageId =
latestMessage?.role === "assistant" ? latestMessage.id : await generateAssistantMessageId();

// Charge the account for this turn: atomic wallet debit + audit row via the
// `deduct_credits_with_audit` Postgres function (`handleChatCredits` →
// `recordCreditDeduction`), using the usage on `responseMessage.metadata`
// when present.
//
// This MUST run even when the step throws or returns no `responseMessage`.
// A turn can run a customer-facing tool (e.g. `send_email`) mid-step, so if
// billing were gated on a clean `responseMessage` return, a turn that emails
// the customer and then fails would be free AND leave no `usage_events` row
// (observed 2026-07-23: a scheduled briefing emailed the customer but wrote
// zero LLM rows). `chargeTurn` is idempotent — the normal path bills exactly
// once below; the `finally` backstop only fires when a throw skipped it.
// ZERO_USAGE with no gateway cost floors to the 1c minimum in
// `handleChatCredits`, so a failed turn still writes a `model_id` audit row.
let result: Awaited<ReturnType<typeof runAgentStep>> | undefined;
let charged = false;
const chargeTurn = async () => {
if (charged) return;
charged = true;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: A transient credit-recording failure still leaves a side-effecting turn unbilled: this flag is set before the debit outcome is known, while handleChatCredits swallows failures. Return/propagate a success result and mark the turn charged only after a successful debit, so the finally path can retry a failed normal-path attempt.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At app/lib/workflows/runAgentWorkflow.ts, line 118:

<comment>A transient credit-recording failure still leaves a side-effecting turn unbilled: this flag is set before the debit outcome is known, while `handleChatCredits` swallows failures. Return/propagate a success result and mark the turn charged only after a successful debit, so the finally path can retry a failed normal-path attempt.</comment>

<file context>
@@ -97,35 +97,47 @@ export async function runAgentWorkflow(input: RunAgentWorkflowInput): Promise<vo
+  let charged = false;
+  const chargeTurn = async () => {
+    if (charged) return;
+    charged = true;
+    const metadata = result?.responseMessage?.metadata as AgentMessageMetadata | undefined;
+    await handleChatCredits({
</file context>

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,
});
};

try {
const result = await runAgentStep({
result = await runAgentStep({
...input,
writable,
assistantMessageId,
});
console.log("[runAgentWorkflow] finish", { finishReason: result.finishReason });

// The assistant message is persisted per step inside `runAgentStep`, so
// it's not written here. We still use the final `responseMessage` to
// charge the account for this turn: atomic wallet debit + audit row via
// the `deduct_credits_with_audit` Postgres function (`handleChatCredits`
// → `recordCreditDeduction`).
//
// Charge on user-stop too — the provider already billed us for the
// tokens consumed, and the assistant message (including partial tool
// runs) is persisted, so the user owes the charge regardless of how
// the turn ended. `result.responseMessage.metadata` carries the
// usage actually consumed up to the abort point.
if (result.responseMessage) {
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,
});
}
// Charge on user-stop too — the provider already billed us for the tokens
// consumed, and the assistant message (including partial tool runs) is
// persisted, so the user owes the charge regardless of how the turn ended.
await chargeTurn();

// Auto-commit + push only after a natural finish. Skip on user-stop —
// don't push half-done work, and `autoCommitChatTurn` can run 30+
Expand All @@ -145,6 +157,16 @@ export async function runAgentWorkflow(input: RunAgentWorkflowInput): Promise<vo
});
}
} finally {
// Billing backstop: if `runAgentStep` threw AFTER a customer-facing
// side-effect already fired (e.g. the agent's `send_email` tool sent the
// email, then the turn errored), the in-`try` charge above was skipped.
// Bill here — `chargeTurn` is idempotent, so the success path is never
// double-charged; a failed turn still lands a wallet debit + `usage_events`
// audit row instead of emailing-but-not-billing. Runs before cleanup;
// `handleChatCredits` swallows its own errors so it can't break the
// cleanup steps below or mask the original throw.
await chargeTurn();

// Run two cleanup steps in parallel:
// 1) `clearChatActiveStream` — CAS-gated DB clear of the chat's
// `active_stream_id` so the recovery probe flips to false.
Expand Down
Loading