|
| 1 | +import { createServer, type Server } from "node:http"; |
| 2 | +import type { AddressInfo } from "node:net"; |
| 3 | +import type { Model } from "@earendil-works/pi-ai"; |
| 4 | +import { afterEach, describe, expect, it } from "vitest"; |
| 5 | +import { BillingLedger, formatBillingSummary, parseBillingReceipt } from "./codebase-billing.js"; |
| 6 | +import { createCodebaseStream } from "./codebase-stream.js"; |
| 7 | + |
| 8 | +let server: Server | undefined; |
| 9 | +afterEach(async () => { |
| 10 | + server?.closeAllConnections(); |
| 11 | + if (server) await new Promise<void>((resolve) => server!.close(() => resolve())); |
| 12 | +}); |
| 13 | +async function serve(chunks: unknown[], status = 200) { |
| 14 | + let request: Record<string, unknown> | undefined; |
| 15 | + server = createServer(async (req, res) => { |
| 16 | + let raw = ""; |
| 17 | + for await (const bytes of req) raw += bytes; |
| 18 | + request = JSON.parse(raw); |
| 19 | + res.writeHead(status, { "content-type": status === 200 ? "text/event-stream" : "application/json" }); |
| 20 | + if (status !== 200) return res.end(JSON.stringify(chunks[0])); |
| 21 | + for (const chunk of chunks) res.write(`data: ${JSON.stringify(chunk)}\n\n`); |
| 22 | + res.end("data: [DONE]\n\n"); |
| 23 | + }); |
| 24 | + await new Promise<void>((resolve) => server!.listen(0, "127.0.0.1", resolve)); |
| 25 | + const model: Model<string> = { |
| 26 | + id: "d4f", |
| 27 | + name: "Auto", |
| 28 | + provider: "codebase", |
| 29 | + api: "openai-completions", |
| 30 | + baseUrl: `http://127.0.0.1:${(server.address() as AddressInfo).port}`, |
| 31 | + reasoning: false, |
| 32 | + input: ["text"], |
| 33 | + contextWindow: 131072, |
| 34 | + maxTokens: 8192, |
| 35 | + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, |
| 36 | + }; |
| 37 | + const ledger = new BillingLedger(); |
| 38 | + const stream = createCodebaseStream(ledger, model.baseUrl!); |
| 39 | + return { |
| 40 | + ledger, |
| 41 | + model, |
| 42 | + stream, |
| 43 | + request: () => request, |
| 44 | + run: (signal?: AbortSignal) => |
| 45 | + stream( |
| 46 | + model, |
| 47 | + { systemPrompt: "Be concise", messages: [{ role: "user", content: "hi", timestamp: 1 }] }, |
| 48 | + { apiKey: "fake-key", signal }, |
| 49 | + ), |
| 50 | + }; |
| 51 | +} |
| 52 | +const final = { |
| 53 | + model: "deepseek-flash", |
| 54 | + choices: [], |
| 55 | + usage: { |
| 56 | + prompt_tokens: 100, |
| 57 | + completion_tokens: 10, |
| 58 | + total_tokens: 110, |
| 59 | + prompt_tokens_details: { cached_tokens: 80, cache_write_tokens: 5 }, |
| 60 | + }, |
| 61 | + codebase_billing: { credits_charged: 3, usage_estimated: false }, |
| 62 | +}; |
| 63 | + |
| 64 | +describe("Codebase receipt transport", () => { |
| 65 | + it("uses the SDK stream with terminal usage-only receipts and the actual fallback model", async () => { |
| 66 | + const h = await serve([ |
| 67 | + { model: "deepseek-flash", choices: [{ delta: { content: "Hello" }, finish_reason: "stop" }] }, |
| 68 | + final, |
| 69 | + final, |
| 70 | + ]); |
| 71 | + const result = await h.run().result(); |
| 72 | + expect(result.content).toEqual([{ type: "text", text: "Hello" }]); |
| 73 | + expect(result.responseModel).toBe("deepseek-flash"); |
| 74 | + expect(result.usage).toMatchObject({ input: 15, output: 10, cacheRead: 80, cacheWrite: 5 }); |
| 75 | + expect(h.ledger.snapshot()).toEqual({ |
| 76 | + creditsCharged: 3, |
| 77 | + requests: 1, |
| 78 | + estimated: 0, |
| 79 | + unconfirmed: 0, |
| 80 | + models: ["deepseek-flash"], |
| 81 | + }); |
| 82 | + expect(h.request()).toMatchObject({ stream: true, stream_options: { include_usage: true } }); |
| 83 | + }); |
| 84 | + it("never calls a missing or null receipt free", async () => { |
| 85 | + const h = await serve([{ choices: [{ delta: { content: "Hello" }, finish_reason: "stop" }] }]); |
| 86 | + await h.run().result(); |
| 87 | + expect(formatBillingSummary(h.ledger.snapshot())).toBe("0 credits charged + 1 unconfirmed"); |
| 88 | + expect(parseBillingReceipt({ credits_charged: null, usage_estimated: true }, "d4f")?.creditsCharged).toBeNull(); |
| 89 | + }); |
| 90 | + it("keeps the receipt sent before an interrupted-stream SDK error", async () => { |
| 91 | + const h = await serve([final, { error: { type: "provider_error", message: "upstream unavailable" } }]); |
| 92 | + expect((await h.run().result()).stopReason).toBe("error"); |
| 93 | + expect(h.ledger.snapshot()).toMatchObject({ creditsCharged: 3, requests: 1, unconfirmed: 0 }); |
| 94 | + }); |
| 95 | + it("shows an admission denial without a retry or invented debit", async () => { |
| 96 | + const h = await serve([{ error: "insufficient_credits", error_description: "Not enough credits" }], 402); |
| 97 | + expect((await h.run().result()).stopReason).toBe("error"); |
| 98 | + expect(h.ledger.snapshot()).toMatchObject({ creditsCharged: 0, requests: 1, unconfirmed: 1 }); |
| 99 | + }); |
| 100 | + it("keeps cancellation unconfirmed rather than zero-cost", async () => { |
| 101 | + const h = await serve([]); |
| 102 | + const controller = new AbortController(); |
| 103 | + controller.abort(); |
| 104 | + expect((await h.run(controller.signal).result()).stopReason).toBe("aborted"); |
| 105 | + expect(h.ledger.snapshot().unconfirmed).toBe(1); |
| 106 | + }); |
| 107 | + it("preserves SDK tool deltas", async () => { |
| 108 | + const h = await serve([ |
| 109 | + { |
| 110 | + choices: [ |
| 111 | + { |
| 112 | + delta: { |
| 113 | + tool_calls: [{ index: 0, id: "call1", function: { name: "write_file", arguments: '{"path":' } }], |
| 114 | + }, |
| 115 | + }, |
| 116 | + ], |
| 117 | + }, |
| 118 | + { |
| 119 | + choices: [ |
| 120 | + { |
| 121 | + delta: { tool_calls: [{ index: 0, function: { arguments: '"a.txt"}' } }] }, |
| 122 | + finish_reason: "tool_calls", |
| 123 | + }, |
| 124 | + ], |
| 125 | + }, |
| 126 | + final, |
| 127 | + ]); |
| 128 | + const result = await h.run().result(); |
| 129 | + expect(result.stopReason).toBe("toolUse"); |
| 130 | + expect(result.content).toEqual([ |
| 131 | + { type: "toolCall", id: "call1", name: "write_file", arguments: { path: "a.txt" } }, |
| 132 | + ]); |
| 133 | + }); |
| 134 | + it("rejects incomplete tool JSON without a completed tool event", async () => { |
| 135 | + const h = await serve([ |
| 136 | + { |
| 137 | + choices: [ |
| 138 | + { |
| 139 | + delta: { |
| 140 | + tool_calls: [{ index: 0, id: "call1", function: { name: "write_file", arguments: '{"path":' } }], |
| 141 | + }, |
| 142 | + finish_reason: "length", |
| 143 | + }, |
| 144 | + ], |
| 145 | + }, |
| 146 | + final, |
| 147 | + ]); |
| 148 | + const events = []; |
| 149 | + for await (const event of h.run()) events.push(event); |
| 150 | + expect(events.some((event) => event.type === "toolcall_end")).toBe(false); |
| 151 | + expect(events.at(-1)?.type).toBe("error"); |
| 152 | + expect(h.ledger.snapshot().requests).toBe(1); |
| 153 | + }); |
| 154 | + it("never sends a Codebase key to a helper endpoint override", async () => { |
| 155 | + const h = await serve([]); |
| 156 | + const result = await h |
| 157 | + .stream({ ...h.model, baseUrl: "https://example.com/v1" }, { messages: [] }, { apiKey: "fake-key" }) |
| 158 | + .result(); |
| 159 | + expect(result.errorMessage).toContain("different endpoint"); |
| 160 | + expect(h.request()).toBeUndefined(); |
| 161 | + }); |
| 162 | + it.each([-1, 0.5, Number.NaN, "3", undefined])("rejects invalid receipt amount %s", (value) => { |
| 163 | + expect(parseBillingReceipt({ credits_charged: value, usage_estimated: false }, "d4f")).toBeUndefined(); |
| 164 | + }); |
| 165 | +}); |
0 commit comments