diff --git a/app/lib/workflows/__tests__/gateWorkflowCredits.test.ts b/app/lib/workflows/__tests__/gateWorkflowCredits.test.ts new file mode 100644 index 00000000..6d16a528 --- /dev/null +++ b/app/lib/workflows/__tests__/gateWorkflowCredits.test.ts @@ -0,0 +1,63 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { NextResponse } from "next/server"; +import { gateWorkflowCredits } from "@/app/lib/workflows/gateWorkflowCredits"; +import { ensureCreditsOrShortCircuit } from "@/lib/credits/ensureCreditsOrShortCircuit"; +import { alertCreditShortfall } from "@/lib/credits/alertCreditShortfall"; + +vi.mock("@/lib/credits/ensureCreditsOrShortCircuit", () => ({ + ensureCreditsOrShortCircuit: vi.fn(), +})); +vi.mock("@/lib/credits/alertCreditShortfall", () => ({ + alertCreditShortfall: vi.fn(), +})); + +const base = { accountId: "acc-1", chatId: "chat-1", sessionId: "sess-1" }; + +describe("gateWorkflowCredits", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("returns hasCredits:true and does NOT alert when credits are available", async () => { + vi.mocked(ensureCreditsOrShortCircuit).mockResolvedValue(null); + + const result = await gateWorkflowCredits(base); + + expect(result).toEqual({ hasCredits: true }); + expect(alertCreditShortfall).not.toHaveBeenCalled(); + }); + + it("gates on the account id with a 1-credit floor", async () => { + vi.mocked(ensureCreditsOrShortCircuit).mockResolvedValue(null); + + await gateWorkflowCredits(base); + + expect(ensureCreditsOrShortCircuit).toHaveBeenCalledWith( + expect.objectContaining({ accountId: "acc-1", creditsToDeduct: 1 }), + ); + }); + + it("returns hasCredits:false and fires a Telegram alert on shortfall", async () => { + vi.mocked(ensureCreditsOrShortCircuit).mockResolvedValue( + NextResponse.json({ error: "insufficient" }, { status: 402 }), + ); + + const result = await gateWorkflowCredits(base); + + expect(result).toEqual({ hasCredits: false }); + expect(alertCreditShortfall).toHaveBeenCalledWith({ + accountId: "acc-1", + chatId: "chat-1", + sessionId: "sess-1", + }); + }); + + it("fails open (hasCredits:true, no alert) when the credit check throws", async () => { + vi.mocked(ensureCreditsOrShortCircuit).mockRejectedValue(new Error("stripe down")); + + const result = await gateWorkflowCredits(base); + + expect(result).toEqual({ hasCredits: true }); + expect(alertCreditShortfall).not.toHaveBeenCalled(); + }); +}); diff --git a/app/lib/workflows/__tests__/runAgentWorkflow.test.ts b/app/lib/workflows/__tests__/runAgentWorkflow.test.ts index be1adf5d..f3e3c773 100644 --- a/app/lib/workflows/__tests__/runAgentWorkflow.test.ts +++ b/app/lib/workflows/__tests__/runAgentWorkflow.test.ts @@ -7,10 +7,14 @@ import { generateAssistantMessageId } from "@/app/lib/workflows/generateAssistan import { handleChatCredits } from "@/lib/credits/handleChatCredits"; import { autoCommitChatTurn } from "@/lib/chat/auto-commit/autoCommitChatTurn"; import { deleteEphemeralKeyStep } from "@/app/lib/workflows/deleteEphemeralKeyStep"; +import { gateWorkflowCredits } from "@/app/lib/workflows/gateWorkflowCredits"; vi.mock("@/app/lib/workflows/deleteEphemeralKeyStep", () => ({ deleteEphemeralKeyStep: vi.fn(), })); +vi.mock("@/app/lib/workflows/gateWorkflowCredits", () => ({ + gateWorkflowCredits: vi.fn(), +})); vi.mock("@/app/lib/workflows/runAgentStep", () => ({ runAgentStep: vi.fn(), })); @@ -45,6 +49,8 @@ vi.mock("workflow", () => ({ beforeEach(() => { vi.clearAllMocks(); vi.mocked(generateAssistantMessageId).mockResolvedValue("asst-fresh-id"); + // Default: account is funded so existing behavior is unchanged. + vi.mocked(gateWorkflowCredits).mockResolvedValue({ hasCredits: true }); }); const baseInput = { @@ -359,6 +365,65 @@ describe("runAgentWorkflow", () => { }); }); + describe("credit gate", () => { + it("gates on credits before running the step, keyed to the account/chat/session", async () => { + vi.mocked(runAgentStep).mockResolvedValue({ + finishReason: "stop", + aborted: false, + responseMessage: undefined, + }); + + await runAgentWorkflow(baseInput); + + expect(gateWorkflowCredits).toHaveBeenCalledWith({ + accountId: "acc-1", + chatId: "chat-1", + sessionId: "session-1", + }); + }); + + it("skips runAgentStep + billing when the gate reports no credits", async () => { + vi.mocked(gateWorkflowCredits).mockResolvedValue({ hasCredits: false }); + + await runAgentWorkflow(baseInput); + + expect(runAgentStep).not.toHaveBeenCalled(); + expect(handleChatCredits).not.toHaveBeenCalled(); + expect(autoCommitChatTurn).not.toHaveBeenCalled(); + }); + + it("does not generate an assistant message id when the gate blocks the run", async () => { + vi.mocked(gateWorkflowCredits).mockResolvedValue({ hasCredits: false }); + + await runAgentWorkflow(baseInput); + + expect(generateAssistantMessageId).not.toHaveBeenCalled(); + }); + + it("still runs cleanup (clearActiveStream + closeChatStream) when the gate blocks the run", async () => { + vi.mocked(gateWorkflowCredits).mockResolvedValue({ hasCredits: false }); + + await runAgentWorkflow(baseInput); + + expect(clearChatActiveStream).toHaveBeenCalledTimes(1); + expect(closeChatStream).toHaveBeenCalledTimes(1); + }); + + it("revokes the ephemeral key when the gate blocks a headless run", async () => { + vi.mocked(gateWorkflowCredits).mockResolvedValue({ hasCredits: false }); + + await runAgentWorkflow({ + ...baseInput, + agentContext: { + sandbox: { state: { type: "vercel" }, workingDirectory: "/sandbox/mono" }, + ephemeralKeyId: "ephem-key-1", + } as never, + }); + + expect(deleteEphemeralKeyStep).toHaveBeenCalledWith("ephem-key-1"); + }); + }); + describe("user-abort path", () => { const abortedResponseMessage = { id: "assistant-msg-xyz", diff --git a/app/lib/workflows/gateWorkflowCredits.ts b/app/lib/workflows/gateWorkflowCredits.ts new file mode 100644 index 00000000..7bb7cde1 --- /dev/null +++ b/app/lib/workflows/gateWorkflowCredits.ts @@ -0,0 +1,70 @@ +import { NextResponse } from "next/server"; +import { ensureCreditsOrShortCircuit } from "@/lib/credits/ensureCreditsOrShortCircuit"; +import { alertCreditShortfall } from "@/lib/credits/alertCreditShortfall"; +import { CREDIT_AUTO_RECHARGE_FALLBACK_SUCCESS_URL } from "@/lib/credits/const"; + +/** + * Minimum credit balance required before an agent run is allowed to start. + * A 1-credit floor turns the "charge after the run" model into "don't start a + * run for an account that's already out of credits (and can't auto-recharge)". + * Accounts with a saved card top up transparently inside the gate; burners with + * no card are stopped here instead of draining into a deep negative balance + * (recoupable/chat#1885 — WAVS Digital hit -1,563). + */ +const SCHEDULED_RUN_MIN_CREDITS = 1; + +export type GateWorkflowCreditsParams = { + accountId: string; + chatId: string; + sessionId: string; +}; + +export type GateWorkflowCreditsResult = { + /** True when the run may proceed; false when it should be skipped. */ + hasCredits: boolean; +}; + +/** + * Pre-run credit gate for `runAgentWorkflow`, run as a `"use step"` so its DB / + * Stripe I/O is legal inside the durable workflow and cached across replays. + * + * Reuses the request-path gate (`ensureCreditsOrShortCircuit`, which attempts a + * silent auto-recharge) so scheduled runs enforce the same credit rules as + * `/api/research` and the social-scrape routes. On shortfall it fires a + * Telegram alert to the team and returns `hasCredits:false` so the caller skips + * the run before any model tokens are spent. + * + * Fails OPEN: if the credit check itself throws (transient Stripe/Supabase + * error), we let the run proceed rather than break every run on an infra blip — + * the post-run charge still applies. + */ +export async function gateWorkflowCredits( + params: GateWorkflowCreditsParams, +): Promise { + "use step"; + + const { accountId, chatId, sessionId } = params; + + try { + const short = await ensureCreditsOrShortCircuit({ + accountId, + creditsToDeduct: SCHEDULED_RUN_MIN_CREDITS, + successUrl: CREDIT_AUTO_RECHARGE_FALLBACK_SUCCESS_URL, + }); + + if (short instanceof NextResponse) { + console.warn("[gateWorkflowCredits] insufficient credits, skipping run", { + accountId, + chatId, + sessionId, + }); + await alertCreditShortfall({ accountId, chatId, sessionId }); + return { hasCredits: false }; + } + + return { hasCredits: true }; + } catch (error) { + console.error("[gateWorkflowCredits] credit check failed, proceeding (fail-open):", error); + return { hasCredits: true }; + } +} diff --git a/app/lib/workflows/runAgentWorkflow.ts b/app/lib/workflows/runAgentWorkflow.ts index b7335372..b034d690 100644 --- a/app/lib/workflows/runAgentWorkflow.ts +++ b/app/lib/workflows/runAgentWorkflow.ts @@ -3,6 +3,7 @@ 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 { gateWorkflowCredits } from "@/app/lib/workflows/gateWorkflowCredits"; import { clearChatActiveStream } from "@/lib/chat/clearChatActiveStream"; import { deleteEphemeralKeyStep } from "@/app/lib/workflows/deleteEphemeralKeyStep"; import { handleChatCredits } from "@/lib/credits/handleChatCredits"; @@ -84,20 +85,41 @@ export async function runAgentWorkflow(input: RunAgentWorkflowInput): Promise(); - // Pick or generate a stable id for the assistant message. If the - // last message in the conversation is already an assistant message - // (we're resuming an in-progress turn after a tool-call interaction) - // reuse its id so chunks append to the same `chat_messages` row. - // Otherwise generate a fresh id once via a `"use step"` so the - // value is durable across workflow replays. Mirrors open-agents' - // pattern in `apps/web/app/workflows/chat.ts` where the id is - // generated in the workflow body and threaded into every - // `runAgentStep` call. - const latestMessage = input.messages.at(-1); - const assistantMessageId = - latestMessage?.role === "assistant" ? latestMessage.id : await generateAssistantMessageId(); - try { + // Pre-run credit gate. `runAgentStep` charges AFTER the turn, so without a + // pre-gate an unfunded account (out of credits, no card to auto-recharge) + // keeps running scheduled turns into a deep negative balance + // (recoupable/chat#1885 — WAVS Digital hit -1,563). Gate first — before we + // even mint an assistant message id or spend a model token — and skip the + // run on shortfall, firing a Telegram alert. `finally` still runs cleanup. + const gate = await gateWorkflowCredits({ + accountId: input.accountId, + chatId: input.chatId, + sessionId: input.sessionId, + }); + if (!gate.hasCredits) { + console.warn("[runAgentWorkflow] skipping run — insufficient credits", { + accountId: input.accountId, + chatId: input.chatId, + sessionId: input.sessionId, + workflowRunId, + }); + return; + } + + // Pick or generate a stable id for the assistant message. If the + // last message in the conversation is already an assistant message + // (we're resuming an in-progress turn after a tool-call interaction) + // reuse its id so chunks append to the same `chat_messages` row. + // Otherwise generate a fresh id once via a `"use step"` so the + // value is durable across workflow replays. Mirrors open-agents' + // pattern in `apps/web/app/workflows/chat.ts` where the id is + // generated in the workflow body and threaded into every + // `runAgentStep` call. + const latestMessage = input.messages.at(-1); + const assistantMessageId = + latestMessage?.role === "assistant" ? latestMessage.id : await generateAssistantMessageId(); + const result = await runAgentStep({ ...input, writable, diff --git a/lib/credits/__tests__/alertCreditShortfall.test.ts b/lib/credits/__tests__/alertCreditShortfall.test.ts new file mode 100644 index 00000000..a19267dd --- /dev/null +++ b/lib/credits/__tests__/alertCreditShortfall.test.ts @@ -0,0 +1,32 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { alertCreditShortfall } from "@/lib/credits/alertCreditShortfall"; +import { sendMessage } from "@/lib/telegram/sendMessage"; + +vi.mock("@/lib/telegram/sendMessage", () => ({ + sendMessage: vi.fn(), +})); + +describe("alertCreditShortfall", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("sends a Telegram alert naming the account and chat", async () => { + vi.mocked(sendMessage).mockResolvedValue({} as never); + + await alertCreditShortfall({ accountId: "acc-1", chatId: "chat-9", sessionId: "sess-9" }); + + expect(sendMessage).toHaveBeenCalledTimes(1); + const text = vi.mocked(sendMessage).mock.calls[0]?.[0] as string; + expect(text).toContain("acc-1"); + expect(text).toContain("chat-9"); + }); + + it("never throws when Telegram send fails (alerting must not break the workflow)", async () => { + vi.mocked(sendMessage).mockRejectedValue(new Error("telegram down")); + + await expect( + alertCreditShortfall({ accountId: "acc-1", chatId: "chat-9", sessionId: "sess-9" }), + ).resolves.toBeUndefined(); + }); +}); diff --git a/lib/credits/alertCreditShortfall.ts b/lib/credits/alertCreditShortfall.ts new file mode 100644 index 00000000..a38b9df0 --- /dev/null +++ b/lib/credits/alertCreditShortfall.ts @@ -0,0 +1,39 @@ +import { sendMessage } from "@/lib/telegram/sendMessage"; + +export type CreditShortfallAlert = { + /** Account whose wallet couldn't cover the run (and couldn't auto-recharge). */ + accountId: string; + /** Chat the blocked run belonged to. */ + chatId: string; + /** Session the blocked run belonged to. */ + sessionId: string; +}; + +/** + * Fire a Telegram alert to the team when a scheduled/agent run is skipped for + * insufficient credits (the account is out of credits AND auto-recharge failed + * or is opted out). Surfaces burners that would otherwise drain into a deep + * negative balance (e.g. WAVS Digital hit -1,563) before we notice. + * + * Never throws — a Telegram outage must not turn into a workflow failure. Fire + * and (best-effort) forget, mirroring `sendSalesNotification`. + */ +export async function alertCreditShortfall(alert: CreditShortfallAlert): Promise { + const { accountId, chatId, sessionId } = alert; + const text = [ + "⚠️ *Scheduled run blocked — insufficient credits*", + "", + `*Account:* ${accountId}`, + `*Chat:* ${chatId}`, + `*Session:* ${sessionId}`, + "", + "The account is out of credits and auto-recharge did not cover the run.", + `*Time:* ${new Date().toISOString()}`, + ].join("\n"); + + try { + await sendMessage(text, { parse_mode: "Markdown" }); + } catch (error) { + console.error("[alertCreditShortfall] failed to send Telegram alert:", error); + } +}