From 7b75476c4aab17afd07ffa8a93930a7eb4e87b8b Mon Sep 17 00:00:00 2001 From: diegosouzapw Date: Mon, 6 Apr 2026 09:47:45 -0300 Subject: [PATCH] fix(core): preserve primary failures across chat fallbacks Keep the original combo and budget exhaustion errors when global or emergency fallbacks also fail so callers see the real upstream cause. Also preserve translated responses for memory extraction before output post-processing, track pending rate-limit async work for deterministic test resets, and expose usage helpers needed for deeper branch coverage. Expand unit coverage across moderation, media generation, streaming, response logging, usage helpers, and fallback proxy error handling. --- open-sse/handlers/chatCore.ts | 35 +-- open-sse/services/rateLimitManager.ts | 30 ++- open-sse/services/usage.ts | 8 + tests/unit/chat-route-coverage.test.mjs | 118 ++++++++++ .../unit/chatcore-translation-paths.test.mjs | 8 +- tests/unit/moderations-handler.test.mjs | 111 +++++++++ tests/unit/music-generation-handler.test.mjs | 98 ++++++++ tests/unit/rate-limit-manager.test.mjs | 2 +- tests/unit/responses-transformer.test.mjs | 57 ++++- tests/unit/stream-utils.test.mjs | 160 +++++++++++++ tests/unit/usage-service-hardening.test.mjs | 161 +++++++++++++ tests/unit/video-generation-handler.test.mjs | 218 ++++++++++++++++++ 12 files changed, 978 insertions(+), 28 deletions(-) create mode 100644 tests/unit/moderations-handler.test.mjs diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 5dd5dea866..fbcff1bfc8 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -1246,22 +1246,9 @@ export async function handleChatCore({ log?: unknown; upstreamExtraHeaders?: Record | null; }) => { + let result; try { - const result = await nativeExec.execute(input); - if (isRetryableStatus(result.response.status)) { - log?.info?.( - "UPSTREAM_PROXY", - `${prov} native failed (${result.response.status}), retrying via CLIProxyAPI` - ); - try { - return await proxyExec.execute(input); - } catch (proxyErr) { - const proxyMsg = proxyErr instanceof Error ? proxyErr.message : String(proxyErr); - log?.error?.("UPSTREAM_PROXY", `${prov} CLIProxyAPI fallback also failed: ${proxyMsg}`); - throw proxyErr; - } - } - return result; + result = await nativeExec.execute(input); } catch (err) { const errMsg = err instanceof Error ? err.message : String(err); log?.info?.("UPSTREAM_PROXY", `${prov} native error (${errMsg}), retrying via CLIProxyAPI`); @@ -1273,6 +1260,21 @@ export async function handleChatCore({ throw proxyErr; } } + + if (!isRetryableStatus(result.response.status)) { + return result; + } + log?.info?.( + "UPSTREAM_PROXY", + `${prov} native failed (${result.response.status}), retrying via CLIProxyAPI` + ); + try { + return await proxyExec.execute(input); + } catch (proxyErr) { + const proxyMsg = proxyErr instanceof Error ? proxyErr.message : String(proxyErr); + log?.error?.("UPSTREAM_PROXY", `${prov} CLIProxyAPI fallback also failed: ${proxyMsg}`); + throw proxyErr; + } }; return wrapper; }; @@ -2125,6 +2127,7 @@ export async function handleChatCore({ toolNameMap as Map | null ) : responseBody; + const memoryExtractionResponse = translatedResponse; // T26: Strip markdown code blocks if provider format is Claude if (sourceFormat === "claude" && !stream) { @@ -2180,7 +2183,7 @@ export async function handleChatCore({ )) || skillRequestId; if (apiKeyInfo?.id && memorySettings?.enabled && memorySettings.maxTokens > 0) { - const memoryText = extractMemoryTextFromResponse(translatedResponse); + const memoryText = extractMemoryTextFromResponse(memoryExtractionResponse); if (memoryText) { extractFacts(memoryText, apiKeyInfo.id, pipelineSessionId); } diff --git a/open-sse/services/rateLimitManager.ts b/open-sse/services/rateLimitManager.ts index 4e8adc5a57..48e8e46b6a 100644 --- a/open-sse/services/rateLimitManager.ts +++ b/open-sse/services/rateLimitManager.ts @@ -55,6 +55,7 @@ const enabledConnections = new Set(); // Store learned limits for persistence (debounced) const learnedLimits: Record = {}; let persistTimer: ReturnType | null = null; +const pendingAsyncOperations = new Set>(); const PERSIST_DEBOUNCE_MS = 60_000; // Debounce persistence to every 60s max // Track initialization @@ -75,6 +76,14 @@ const DEFAULT_SETTINGS = { maxWait: MAX_WAIT_MS, // Fail-fast: don't queue forever on 429 exhaustion }; +function trackAsyncOperation(promise: Promise): Promise { + pendingAsyncOperations.add(promise); + promise.finally(() => { + pendingAsyncOperations.delete(promise); + }); + return promise; +} + /** * Initialize rate limit protection from persisted connection settings. * Called once on app startup. @@ -363,9 +372,11 @@ export function updateFromHeaders(provider, connectionId, headers, status, model // be hours for providers like Codex with long rate limit windows). // This lets upstream callers (e.g. LiteLLM) trigger fallback to other providers. // After stop, delete from Map so getLimiter() creates a fresh instance. - limiter.stop({ dropWaitingJobs: true }).finally(() => { - limiters.delete(limiterKey); - }); + trackAsyncOperation( + limiter.stop({ dropWaitingJobs: true }).finally(() => { + limiters.delete(limiterKey); + }) + ); return; } @@ -504,7 +515,7 @@ function recordLearnedLimit( if (!persistTimer) { persistTimer = setTimeout(async () => { persistTimer = null; - await persistLearnedLimitsNow(); + await trackAsyncOperation(persistLearnedLimitsNow()); }, PERSIST_DEBOUNCE_MS); } } @@ -514,10 +525,13 @@ export async function __flushLearnedLimitsForTests() { clearTimeout(persistTimer); persistTimer = null; } - await persistLearnedLimitsNow(); + await trackAsyncOperation(persistLearnedLimitsNow()); + if (pendingAsyncOperations.size > 0) { + await Promise.allSettled(Array.from(pendingAsyncOperations)); + } } -export function __resetRateLimitManagerForTests() { +export async function __resetRateLimitManagerForTests() { if (persistTimer) { clearTimeout(persistTimer); persistTimer = null; @@ -533,6 +547,10 @@ export function __resetRateLimitManagerForTests() { for (const key of Object.keys(learnedLimits)) { delete learnedLimits[key]; } + + if (pendingAsyncOperations.size > 0) { + await Promise.allSettled(Array.from(pendingAsyncOperations)); + } } /** diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index 84a09ab828..f07530ef3e 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -1344,3 +1344,11 @@ async function getIflowUsage(accessToken) { return { message: "Unable to fetch Qoder usage." }; } } + +export const __testing = { + parseResetTime, + formatGitHubQuotaSnapshot, + inferGitHubPlanName, + getGeminiCliPlanLabel, + getAntigravityPlanLabel, +}; diff --git a/tests/unit/chat-route-coverage.test.mjs b/tests/unit/chat-route-coverage.test.mjs index 663416a0cb..411ce860e8 100644 --- a/tests/unit/chat-route-coverage.test.mjs +++ b/tests/unit/chat-route-coverage.test.mjs @@ -287,6 +287,47 @@ test("handleChat routes exact combo names and can recover via global fallback", assert.equal(json.choices[0].message.content, "Global fallback answered"); }); +test("handleChat keeps the combo error when the global fallback throws", async () => { + await seedConnection("openai", { apiKey: "sk-openai-combo-fail" }); + await seedConnection("claude", { apiKey: "sk-claude-fallback-throw" }); + await combosDb.createCombo({ + name: "router-global-fallback-throw", + strategy: "priority", + config: { maxRetries: 0, retryDelayMs: 0 }, + models: ["openai/gpt-4o-mini"], + }); + await settingsDb.updateSettings({ + globalFallbackModel: "claude/claude-3-5-sonnet-20241022", + }); + + let attempts = 0; + globalThis.fetch = async () => { + attempts += 1; + if (attempts === 1) { + return new Response(JSON.stringify({ error: { message: "primary combo failed" } }), { + status: 503, + headers: { "Content-Type": "application/json" }, + }); + } + throw new Error("fallback transport crashed"); + }; + + const response = await handleChat( + buildRequest({ + body: { + model: "router-global-fallback-throw", + stream: false, + messages: [{ role: "user", content: "Use combo fallback but force a throw" }], + }, + }) + ); + const json = await response.json(); + + assert.equal(response.status, 503); + assert.equal(attempts, 2); + assert.match(json.error.message, /primary combo failed/i); +}); + test("handleChat returns 400 when no provider credentials exist", async () => { const response = await handleChat( buildRequest({ @@ -367,6 +408,83 @@ test("handleChat maps upstream timeouts to HTTP 504", async () => { assert.match(json.error.message, /\[504\]: upstream timed out/); }); +test("handleChat uses the emergency fallback model on budget exhaustion", async () => { + await seedConnection("openai", { apiKey: "sk-openai-billing" }); + await seedConnection("nvidia", { apiKey: "sk-nvidia-fallback" }); + const seenBodies = []; + + globalThis.fetch = async (_url, init = {}) => { + const body = JSON.parse(String(init.body)); + seenBodies.push(body); + + if (seenBodies.length === 1) { + return new Response(JSON.stringify({ error: { message: "billing limit exceeded" } }), { + status: 402, + headers: { "Content-Type": "application/json" }, + }); + } + + return buildOpenAIResponse("Emergency fallback answered", "gpt-oss-120b"); + }; + + const response = await handleChat( + buildRequest({ + body: { + model: "openai/gpt-4o-mini", + stream: false, + max_tokens: 9000, + messages: [{ role: "user", content: "budget exhausted" }], + }, + }) + ); + const json = await response.json(); + + assert.equal(response.status, 200); + assert.equal(seenBodies.length, 2); + assert.equal(seenBodies[1].model, "openai/gpt-oss-120b"); + assert.equal(seenBodies[1].max_tokens, 4096); + assert.equal(seenBodies[1].max_completion_tokens, 4096); + assert.equal(json.choices[0].message.content, "Emergency fallback answered"); +}); + +test("handleChat returns the primary budget error when emergency fallback also fails", async () => { + await seedConnection("openai", { apiKey: "sk-openai-billing-fail" }); + await seedConnection("nvidia", { apiKey: "sk-nvidia-fallback-fail" }); + const seenModels = []; + + globalThis.fetch = async (_url, init = {}) => { + const body = JSON.parse(String(init.body)); + seenModels.push(body.model); + + if (seenModels.length === 1) { + return new Response(JSON.stringify({ error: { message: "quota exceeded" } }), { + status: 402, + headers: { "Content-Type": "application/json" }, + }); + } + + return new Response(JSON.stringify({ error: { message: "fallback unavailable" } }), { + status: 503, + headers: { "Content-Type": "application/json" }, + }); + }; + + const response = await handleChat( + buildRequest({ + body: { + model: "openai/gpt-4o-mini", + stream: false, + messages: [{ role: "user", content: "budget exhausted again" }], + }, + }) + ); + const json = await response.json(); + + assert.equal(response.status, 402); + assert.deepEqual(seenModels, ["gpt-4o-mini", "openai/gpt-oss-120b"]); + assert.match(json.error.message, /quota exceeded/i); +}); + test("handleChat rejects models that are not allowed by the caller API key policy", async () => { await seedConnection("openai", { apiKey: "sk-openai-policy" }); const apiKey = await seedApiKey({ diff --git a/tests/unit/chatcore-translation-paths.test.mjs b/tests/unit/chatcore-translation-paths.test.mjs index 76c8e17611..d97dfd77ed 100644 --- a/tests/unit/chatcore-translation-paths.test.mjs +++ b/tests/unit/chatcore-translation-paths.test.mjs @@ -824,8 +824,8 @@ test("chatCore fallback proxy mode surfaces CLIProxyAPI errors after a retryable assert.equal(calls.length, 2); assert.equal(result.success, false); - assert.equal(result.status, 500); - assert.equal(result.error, "cliproxy retry failed"); + assert.equal(result.status, 502); + assert.equal(result.error, "[502]: cliproxy retry failed"); }); test("chatCore fallback proxy mode surfaces CLIProxyAPI errors after native executor throws", async () => { @@ -855,8 +855,8 @@ test("chatCore fallback proxy mode surfaces CLIProxyAPI errors after native exec assert.equal(calls.length, 2); assert.equal(result.success, false); - assert.equal(result.status, 500); - assert.equal(result.error, "cliproxy transport exploded"); + assert.equal(result.status, 502); + assert.equal(result.error, "[502]: cliproxy transport exploded"); }); test("chatCore serves a cached idempotent response without hitting the provider twice", async () => { diff --git a/tests/unit/moderations-handler.test.mjs b/tests/unit/moderations-handler.test.mjs new file mode 100644 index 0000000000..54fcb2bc5f --- /dev/null +++ b/tests/unit/moderations-handler.test.mjs @@ -0,0 +1,111 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +const { handleModeration } = await import("../../open-sse/handlers/moderations.ts"); + +const originalFetch = globalThis.fetch; + +test.afterEach(() => { + globalThis.fetch = originalFetch; +}); + +test("handleModeration requires input", async () => { + const response = await handleModeration({ + body: { model: "openai/omni-moderation-latest" }, + credentials: { apiKey: "sk-test" }, + }); + const payload = await response.json(); + + assert.equal(response.status, 400); + assert.equal(payload.error.message, "input is required"); +}); + +test("handleModeration rejects unknown moderation models", async () => { + const response = await handleModeration({ + body: { model: "mystery/moderation", input: "hello" }, + credentials: { apiKey: "sk-test" }, + }); + const payload = await response.json(); + + assert.equal(response.status, 400); + assert.match(payload.error.message, /No moderation provider found/); +}); + +test("handleModeration requires credentials for the resolved provider", async () => { + const response = await handleModeration({ + body: { input: "hello" }, + credentials: null, + }); + const payload = await response.json(); + + assert.equal(response.status, 401); + assert.equal(payload.error.message, "No credentials for moderation provider: openai"); +}); + +test("handleModeration proxies successful requests with default model and accessToken fallback", async () => { + let captured; + + globalThis.fetch = async (url, options = {}) => { + captured = { + url: String(url), + headers: options.headers, + body: JSON.parse(String(options.body || "{}")), + }; + + return Response.json({ + id: "modr-1", + results: [{ flagged: false }], + }); + }; + + const response = await handleModeration({ + body: { input: "all clear" }, + credentials: { accessToken: "oauth-token" }, + }); + + assert.equal(captured.url, "https://api.openai.com/v1/moderations"); + assert.equal(captured.headers.Authorization, "Bearer oauth-token"); + assert.deepEqual(captured.body, { + model: "omni-moderation-latest", + input: "all clear", + }); + assert.equal(response.status, 200); + assert.equal(response.headers.get("access-control-allow-origin"), "*"); + assert.deepEqual(await response.json(), { + id: "modr-1", + results: [{ flagged: false }], + }); +}); + +test("handleModeration returns upstream error payloads with CORS headers", async () => { + globalThis.fetch = async () => + new Response('{"error":"busy"}', { + status: 429, + headers: { "content-type": "application/json" }, + }); + + const response = await handleModeration({ + body: { model: "openai/text-moderation-latest", input: "check this" }, + credentials: { apiKey: "sk-test" }, + }); + + assert.equal(response.status, 429); + assert.equal(await response.text(), '{"error":"busy"}'); + assert.equal(response.headers.get("content-type"), "application/json"); + assert.equal(response.headers.get("access-control-allow-origin"), "*"); +}); + +test("handleModeration returns a 500 when the upstream request throws", async () => { + globalThis.fetch = async () => { + throw new Error("socket closed"); + }; + + const response = await handleModeration({ + body: { model: "openai/text-moderation-latest", input: "check this" }, + credentials: { apiKey: "sk-test" }, + }); + const payload = await response.json(); + + assert.equal(response.status, 500); + assert.match(payload.error.message, /Moderation request failed: socket closed/); +}); diff --git a/tests/unit/music-generation-handler.test.mjs b/tests/unit/music-generation-handler.test.mjs index bad245ea1d..690eb708f4 100644 --- a/tests/unit/music-generation-handler.test.mjs +++ b/tests/unit/music-generation-handler.test.mjs @@ -134,3 +134,101 @@ test("handleMusicGeneration rejects unsupported provider formats", async () => { } } }); + +test("handleMusicGeneration returns unknown provider when registry lookup disappears after parsing", async () => { + Object.defineProperty(MUSIC_PROVIDERS, "flakyprovider", { + configurable: true, + enumerable: true, + get() { + delete MUSIC_PROVIDERS.flakyprovider; + return { + id: "flakyprovider", + baseUrl: "http://localhost:9999", + authType: "none", + authHeader: "none", + format: "comfyui", + models: [{ id: "ghost-model", name: "Ghost Model" }], + }; + }, + }); + + const result = await handleMusicGeneration({ + body: { model: "flakyprovider/ghost-model", prompt: "x" }, + credentials: null, + log: null, + }); + + assert.equal(result.success, false); + assert.equal(result.status, 400); + assert.match(result.error, /Unknown music provider: flakyprovider/); +}); + +test("handleMusicGeneration returns provider errors for ComfyUI failures and logs defaults", async () => { + const originalFetch = globalThis.fetch; + const originalSetTimeout = globalThis.setTimeout; + const logEntries = []; + let promptBody; + + globalThis.setTimeout = immediateTimeout; + globalThis.fetch = async (url, options = {}) => { + const stringUrl = String(url); + + if (stringUrl === "http://localhost:8188/prompt") { + promptBody = JSON.parse(String(options.body || "{}")); + return new Response(JSON.stringify({ prompt_id: "music-fail" }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + + if (stringUrl === "http://localhost:8188/history/music-fail") { + return new Response( + JSON.stringify({ + "music-fail": { + outputs: { + 7: { + audio: [{ filename: "broken.wav", subfolder: "out", type: "output" }], + }, + }, + }, + }), + { status: 200, headers: { "content-type": "application/json" } } + ); + } + + if (stringUrl.includes("/view?")) { + return new Response("missing output", { status: 500 }); + } + + throw new Error(`Unexpected URL: ${stringUrl}`); + }; + + try { + const result = await handleMusicGeneration({ + body: { + model: "comfyui/musicgen-medium", + prompt: "slow piano", + }, + credentials: null, + log: { + info: (...args) => logEntries.push(["info", ...args]), + error: (...args) => logEntries.push(["error", ...args]), + }, + }); + + assert.equal(promptBody.prompt["4"].inputs.seconds, 10); + assert.equal(promptBody.prompt["5"].inputs.steps, 100); + assert.equal(promptBody.prompt["5"].inputs.cfg, 7); + assert.equal(result.success, false); + assert.equal(result.status, 502); + assert.match(result.error, /ComfyUI fetch output failed \(500\)/); + assert.deepEqual( + logEntries.map((entry) => entry[0]), + ["info", "error"] + ); + assert.match(logEntries[1][2], /comfyui error/i); + } finally { + globalThis.fetch = originalFetch; + globalThis.setTimeout = originalSetTimeout; + } +}); diff --git a/tests/unit/rate-limit-manager.test.mjs b/tests/unit/rate-limit-manager.test.mjs index e3ad2ca199..5fbaee44ca 100644 --- a/tests/unit/rate-limit-manager.test.mjs +++ b/tests/unit/rate-limit-manager.test.mjs @@ -9,7 +9,7 @@ function wait(ms) { } test.afterEach(async () => { - rateLimitManager.__resetRateLimitManagerForTests(); + await rateLimitManager.__resetRateLimitManagerForTests(); await wait(5); }); diff --git a/tests/unit/responses-transformer.test.mjs b/tests/unit/responses-transformer.test.mjs index e73518b319..1f708de319 100644 --- a/tests/unit/responses-transformer.test.mjs +++ b/tests/unit/responses-transformer.test.mjs @@ -1,6 +1,6 @@ import test from "node:test"; import assert from "node:assert/strict"; -import { mkdtempSync, readdirSync, readFileSync } from "node:fs"; +import { mkdtempSync, readdirSync, readFileSync, rmSync, unlinkSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -179,3 +179,58 @@ test("createResponsesLogger persists input and output event logs on flush", asyn assert.match(outputLog, /response\.completed/); assert.match(output, /data: \[DONE]/); }); + +test("createResponsesApiTransformStream ignores malformed events and preserves usage-only chunks", async () => { + const output = await runTransformStream([ + "event: ping\n\n", + "data: [DONE]\n\n", + 'data: {"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}}\n\n', + "data: {not-json}\n\n", + 'data: {"id":"chatcmpl_edge","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]}\n\n', + ]); + + const events = parseSseOutput(output); + const completed = JSON.parse( + events.find((event) => event.event === "response.completed").data + ).response; + + assert.equal(completed.id, "resp_chatcmpl_edge"); + assert.equal(completed.output[0].content[0].text, "ok"); + assert.deepEqual(completed.usage, { + prompt_tokens: 2, + completion_tokens: 1, + total_tokens: 3, + }); +}); + +test("createResponsesLogger returns null for invalid base paths and swallows flush write failures", () => { + const blockedPath = join(tmpdir(), `responses-transformer-blocked-${Date.now()}`); + writeFileSync(blockedPath, "blocked"); + + try { + const blockedLogger = createResponsesLogger("gpt-4o", blockedPath); + assert.equal(blockedLogger, null); + } finally { + unlinkSync(blockedPath); + } + + const logsDir = mkdtempSync(join(tmpdir(), "responses-transformer-broken-")); + const logger = createResponsesLogger("gpt-4o", logsDir); + const capturedLogs = []; + const originalConsoleLog = console.log; + + logger.logInput("input"); + logger.logOutput("output"); + + const sessionDir = readdirSync(join(logsDir, "logs"))[0]; + rmSync(join(logsDir, "logs", sessionDir), { recursive: true, force: true }); + console.log = (...args) => capturedLogs.push(args.join(" ")); + + try { + logger.flush(); + } finally { + console.log = originalConsoleLog; + } + + assert.ok(capturedLogs.some((entry) => entry.includes("[RESPONSES] Failed to write logs:"))); +}); diff --git a/tests/unit/stream-utils.test.mjs b/tests/unit/stream-utils.test.mjs index ca165161d5..e9c5351fba 100644 --- a/tests/unit/stream-utils.test.mjs +++ b/tests/unit/stream-utils.test.mjs @@ -298,6 +298,166 @@ test("createSSEStream passthrough fixes generic ids and normalizes reasoning ali assert.equal(text.includes('"reasoning":"Let me think first"'), false); }); +test("createSSEStream passthrough splits mixed reasoning and content deltas and estimates usage", async () => { + let onCompletePayload = null; + const text = await readTransformed( + [ + `data: ${JSON.stringify({ + id: "chatcmpl_reasoning", + object: "chat.completion.chunk", + created: 1, + model: "gpt-4.1-mini", + choices: [ + { + index: 0, + delta: { + reasoning_content: "First think", + content: "Then answer", + }, + }, + ], + })}\n\n`, + `data: ${JSON.stringify({ + id: "chatcmpl_reasoning", + object: "chat.completion.chunk", + created: 1, + model: "gpt-4.1-mini", + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + })}\n\n`, + ], + { + mode: "passthrough", + sourceFormat: FORMATS.OPENAI, + provider: "openai", + model: "gpt-4.1-mini", + body: { + messages: [{ role: "user", content: "hello world" }], + }, + onComplete(payload) { + onCompletePayload = payload; + }, + } + ); + + const reasoningIndex = text.indexOf('"reasoning_content":"First think"'); + const contentIndex = text.indexOf('"content":"Then answer"'); + + assert.ok(reasoningIndex >= 0); + assert.ok(contentIndex > reasoningIndex); + assert.match(text, /"total_tokens":\d+/); + assert.equal(onCompletePayload.responseBody.choices[0].message.reasoning_content, "First think"); + assert.equal(onCompletePayload.responseBody.choices[0].message.content, "Then answer"); + assert.ok(onCompletePayload.responseBody.usage.total_tokens > 0); +}); + +test("createSSEStream passthrough merges Claude usage chunks and restores mapped tool names", async () => { + let onCompletePayload = null; + const text = await readTransformed( + [ + `data: ${JSON.stringify({ + type: "message_start", + message: { + id: "msg_passthrough", + model: "claude-sonnet-4", + role: "assistant", + usage: { input_tokens: 6 }, + }, + })}\n\n`, + `data: ${JSON.stringify({ + type: "content_block_start", + index: 0, + content_block: { + type: "tool_use", + id: "tool_1", + name: "tool_alias", + input: { path: "/tmp/a" }, + }, + })}\n\n`, + `data: ${JSON.stringify({ + type: "content_block_delta", + index: 1, + delta: { text: "Claude says hi" }, + })}\n\n`, + `data: ${JSON.stringify({ + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { output_tokens: 4 }, + })}\n\n`, + ], + { + mode: "passthrough", + sourceFormat: FORMATS.CLAUDE, + provider: "claude", + model: "claude-sonnet-4", + toolNameMap: new Map([["tool_alias", "read_file"]]), + body: { + messages: [{ role: "user", content: "hello" }], + }, + onComplete(payload) { + onCompletePayload = payload; + }, + } + ); + + assert.match(text, /"name":"read_file"/); + assert.equal(text.includes('"name":"tool_alias"'), false); + assert.equal(onCompletePayload.responseBody.choices[0].message.content, "Claude says hi"); + assert.equal(onCompletePayload.responseBody.usage.prompt_tokens, 6); + assert.equal(onCompletePayload.responseBody.usage.completion_tokens, 4); + assert.equal(onCompletePayload.responseBody.usage.total_tokens, 10); +}); + +test("createSSETransformStreamWithLogger flushes a trailing Claude usage event without a newline", async () => { + let onCompletePayload = null; + const text = await readWithTransform( + [ + `data: ${JSON.stringify({ + type: "message_start", + message: { + id: "msg_tail", + model: "claude-sonnet-4", + role: "assistant", + usage: { input_tokens: 3 }, + }, + })}\n\n`, + `data: ${JSON.stringify({ + type: "content_block_start", + index: 0, + content_block: { type: "text", text: "" }, + })}\n\n`, + `data: ${JSON.stringify({ + type: "content_block_delta", + index: 0, + delta: { type: "text_delta", text: "Buffered tail" }, + })}\n\n`, + `data: ${JSON.stringify({ + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { output_tokens: 5 }, + })}`, + ], + createSSETransformStreamWithLogger( + FORMATS.CLAUDE, + FORMATS.OPENAI, + "claude", + null, + null, + "claude-sonnet-4", + null, + { messages: [{ role: "user", content: "hello" }] }, + (payload) => { + onCompletePayload = payload; + } + ) + ); + + assert.match(text, /Buffered tail/); + assert.match(text, /\[DONE\]/); + assert.equal(onCompletePayload.responseBody.choices[0].message.content, "Buffered tail"); + assert.equal(onCompletePayload.responseBody.usage.completion_tokens, 5); + assert.equal(onCompletePayload.responseBody.usage.total_tokens, 5); +}); + test("buildStreamSummaryFromEvents compacts Responses API deltas into a synthetic response", () => { const summary = buildStreamSummaryFromEvents( [ diff --git a/tests/unit/usage-service-hardening.test.mjs b/tests/unit/usage-service-hardening.test.mjs index fc9230bdc3..b92f9799f2 100644 --- a/tests/unit/usage-service-hardening.test.mjs +++ b/tests/unit/usage-service-hardening.test.mjs @@ -2,6 +2,7 @@ import test from "node:test"; import assert from "node:assert/strict"; const usageService = await import("../../open-sse/services/usage.ts"); +const { __testing } = usageService; const originalFetch = globalThis.fetch; @@ -755,3 +756,163 @@ test("usage service covers Qwen, Qoder and GLM branches", async () => { /Invalid API key/ ); }); + +test("usage helper branches cover reset parsing, GitHub quota math, and plan inference fallbacks", () => { + const fixedDate = new Date("2026-01-02T03:04:05.000Z"); + + assert.equal(__testing.parseResetTime(null), null); + assert.equal(__testing.parseResetTime(0), null); + assert.equal(__testing.parseResetTime(fixedDate), fixedDate.toISOString()); + assert.equal(__testing.parseResetTime(fixedDate.getTime()), fixedDate.toISOString()); + assert.equal(__testing.parseResetTime("not-a-date"), null); + + assert.equal(__testing.formatGitHubQuotaSnapshot({}), null); + assert.deepEqual( + __testing.formatGitHubQuotaSnapshot({ entitlement: 20, remaining: 5 }, fixedDate.toISOString()), + { + used: 15, + total: 20, + remaining: 5, + remainingPercentage: 25, + resetAt: fixedDate.toISOString(), + unlimited: false, + } + ); + assert.deepEqual(__testing.formatGitHubQuotaSnapshot({ total: 10, used: 4 }), { + used: 4, + total: 10, + remaining: 6, + remainingPercentage: 60, + resetAt: null, + unlimited: false, + }); + assert.deepEqual(__testing.formatGitHubQuotaSnapshot({ percent_remaining: 30 }), { + used: 70, + total: 100, + remaining: 30, + remainingPercentage: 30, + resetAt: null, + unlimited: false, + }); + assert.deepEqual(__testing.formatGitHubQuotaSnapshot({ unlimited: true }), { + used: 0, + total: 0, + remaining: undefined, + remainingPercentage: undefined, + resetAt: null, + unlimited: true, + }); + + assert.equal( + __testing.inferGitHubPlanName( + { access_type_sku: "copilot_pro_plus" }, + { used: 0, total: 0, resetAt: null, unlimited: false } + ), + "Copilot Pro+" + ); + assert.equal( + __testing.inferGitHubPlanName( + { copilot_plan: "enterprise" }, + { used: 0, total: 0, resetAt: null, unlimited: false } + ), + "Copilot Enterprise" + ); + assert.equal( + __testing.inferGitHubPlanName( + { + copilot_plan: "individual", + monthly_quotas: { premium_interactions: 300 }, + }, + { used: 10, total: 300, resetAt: null, unlimited: false } + ), + "Copilot Pro" + ); + assert.equal( + __testing.inferGitHubPlanName( + { + monthly_quotas: { premium_interactions: 300 }, + }, + { used: 10, total: 300, resetAt: null, unlimited: false } + ), + "Copilot Business" + ); + assert.equal( + __testing.inferGitHubPlanName( + { + monthly_quotas: { chat: 50 }, + }, + null + ), + "Copilot Free" + ); + assert.equal( + __testing.inferGitHubPlanName( + { + access_type_sku: "student_seat", + }, + null + ), + "Copilot Student" + ); + assert.equal(__testing.inferGitHubPlanName({}, null), "GitHub Copilot"); +}); + +test("usage helper branches cover Gemini CLI and Antigravity plan label fallbacks", () => { + assert.equal(__testing.getGeminiCliPlanLabel(null), "Free"); + assert.equal( + __testing.getGeminiCliPlanLabel({ + allowedTiers: [{ id: "tier_ultra", isDefault: true }], + }), + "Ultra" + ); + assert.equal( + __testing.getGeminiCliPlanLabel({ + currentTier: { id: "tier_business" }, + }), + "Business" + ); + assert.equal( + __testing.getGeminiCliPlanLabel({ + subscriptionType: "enterprise", + }), + "Enterprise" + ); + assert.equal( + __testing.getGeminiCliPlanLabel({ + currentTier: { upgradeSubscriptionType: "tier_pro" }, + }), + "Free" + ); + assert.equal( + __testing.getGeminiCliPlanLabel({ + currentTier: { name: "custom neon" }, + }), + "Custom neon" + ); + + assert.equal(__testing.getAntigravityPlanLabel(null), "Free"); + assert.equal( + __testing.getAntigravityPlanLabel({ + allowedTiers: [{ id: "tier_pro", isDefault: true }], + }), + "Pro" + ); + assert.equal( + __testing.getAntigravityPlanLabel({ + currentTier: { displayName: "Standard" }, + }), + "Business" + ); + assert.equal( + __testing.getAntigravityPlanLabel({ + currentTier: { id: "tier_legacy" }, + }), + "Free" + ); + assert.equal( + __testing.getAntigravityPlanLabel({ + currentTier: { name: "custom sky" }, + }), + "Custom sky" + ); +}); diff --git a/tests/unit/video-generation-handler.test.mjs b/tests/unit/video-generation-handler.test.mjs index 04153cf907..2c85ed396f 100644 --- a/tests/unit/video-generation-handler.test.mjs +++ b/tests/unit/video-generation-handler.test.mjs @@ -7,6 +7,7 @@ import { join } from "node:path"; process.env.DATA_DIR = mkdtempSync(join(tmpdir(), "omniroute-video-")); const { handleVideoGeneration } = await import("../../open-sse/handlers/videoGeneration.ts"); +const { VIDEO_PROVIDERS } = await import("../../open-sse/config/videoRegistry.ts"); function immediateTimeout(callback, _ms, ...args) { if (typeof callback === "function") callback(...args); @@ -158,3 +159,220 @@ test("handleVideoGeneration executes ComfyUI workflow and returns fetched output globalThis.setTimeout = originalSetTimeout; } }); + +test("handleVideoGeneration returns unknown provider when registry lookup disappears after parsing", async () => { + Object.defineProperty(VIDEO_PROVIDERS, "flakyprovider", { + configurable: true, + enumerable: true, + get() { + delete VIDEO_PROVIDERS.flakyprovider; + return { + id: "flakyprovider", + baseUrl: "http://localhost:9999", + authType: "none", + authHeader: "none", + format: "comfyui", + models: [{ id: "ghost-model", name: "Ghost Model" }], + }; + }, + }); + + const result = await handleVideoGeneration({ + body: { model: "flakyprovider/ghost-model", prompt: "x" }, + credentials: null, + log: null, + }); + + assert.equal(result.success, false); + assert.equal(result.status, 400); + assert.match(result.error, /Unknown video provider: flakyprovider/); +}); + +test("handleVideoGeneration rejects unsupported provider formats", async () => { + const originalProvider = VIDEO_PROVIDERS.fakeprovider; + + VIDEO_PROVIDERS.fakeprovider = { + id: "fakeprovider", + baseUrl: "http://localhost:9999", + authType: "none", + authHeader: "none", + format: "custom-video", + models: [{ id: "broken-model", name: "Broken Model" }], + }; + + try { + const result = await handleVideoGeneration({ + body: { model: "fakeprovider/broken-model", prompt: "x" }, + credentials: null, + log: null, + }); + + assert.equal(result.success, false); + assert.equal(result.status, 400); + assert.match(result.error, /Unsupported video format: custom-video/); + } finally { + if (originalProvider) { + VIDEO_PROVIDERS.fakeprovider = originalProvider; + } else { + delete VIDEO_PROVIDERS.fakeprovider; + } + } +}); + +test("handleVideoGeneration normalizes SD WebUI image arrays and applies default dimensions", async () => { + const originalFetch = globalThis.fetch; + const logEntries = []; + let captured; + + globalThis.fetch = async (url, options = {}) => { + captured = { + url: String(url), + body: JSON.parse(String(options.body || "{}")), + }; + + return new Response( + JSON.stringify({ + images: ["ZnJhbWUtMQ==", { image: "ZnJhbWUtMg==" }], + }), + { status: 200, headers: { "content-type": "application/json" } } + ); + }; + + try { + const result = await handleVideoGeneration({ + body: { + model: "sdwebui/animatediff-webui", + prompt: "forest path", + }, + credentials: null, + log: { + info: (...args) => logEntries.push(["info", ...args]), + error: (...args) => logEntries.push(["error", ...args]), + }, + }); + + assert.equal(captured.url, "http://localhost:7860/animatediff/v1/generate"); + assert.deepEqual(captured.body, { + prompt: "forest path", + negative_prompt: "", + width: 512, + height: 512, + steps: 20, + cfg_scale: 7, + frames: 16, + fps: 8, + }); + assert.equal(result.success, true); + assert.deepEqual(result.data.data, [ + { b64_json: "ZnJhbWUtMQ==", format: "mp4" }, + { b64_json: "ZnJhbWUtMg==", format: "mp4" }, + ]); + assert.equal(logEntries[0][0], "info"); + assert.match(logEntries[0][2], /sdwebui\/animatediff-webui \(sdwebui\)/); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("handleVideoGeneration returns SD WebUI upstream errors and logs them", async () => { + const originalFetch = globalThis.fetch; + const logEntries = []; + + globalThis.fetch = async () => new Response("provider busy", { status: 503 }); + + try { + const result = await handleVideoGeneration({ + body: { + model: "sdwebui/animatediff-webui", + prompt: "storm", + }, + credentials: null, + log: { + info: (...args) => logEntries.push(["info", ...args]), + error: (...args) => logEntries.push(["error", ...args]), + }, + }); + + assert.equal(result.success, false); + assert.equal(result.status, 503); + assert.equal(result.error, "provider busy"); + assert.deepEqual( + logEntries.map((entry) => entry[0]), + ["info", "error"] + ); + assert.match(logEntries[1][2], /provider busy/); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("handleVideoGeneration returns provider errors for ComfyUI failures and logs defaults", async () => { + const originalFetch = globalThis.fetch; + const originalSetTimeout = globalThis.setTimeout; + const logEntries = []; + let promptBody; + + globalThis.setTimeout = immediateTimeout; + globalThis.fetch = async (url, options = {}) => { + const stringUrl = String(url); + + if (stringUrl === "http://localhost:8188/prompt") { + promptBody = JSON.parse(String(options.body || "{}")); + return new Response(JSON.stringify({ prompt_id: "video-fail" }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + + if (stringUrl === "http://localhost:8188/history/video-fail") { + return new Response( + JSON.stringify({ + "video-fail": { + outputs: { + 7: { + gifs: [{ filename: "broken.webp", subfolder: "out", type: "output" }], + }, + }, + }, + }), + { status: 200, headers: { "content-type": "application/json" } } + ); + } + + if (stringUrl.includes("/view?")) { + return new Response("missing output", { status: 500 }); + } + + throw new Error(`Unexpected URL: ${stringUrl}`); + }; + + try { + const result = await handleVideoGeneration({ + body: { + model: "comfyui/animatediff", + prompt: "night drive", + }, + credentials: null, + log: { + info: (...args) => logEntries.push(["info", ...args]), + error: (...args) => logEntries.push(["error", ...args]), + }, + }); + + assert.equal(promptBody.prompt["4"].inputs.width, 512); + assert.equal(promptBody.prompt["4"].inputs.height, 512); + assert.equal(promptBody.prompt["4"].inputs.batch_size, 16); + assert.equal(promptBody.prompt["7"].inputs.fps, 8); + assert.equal(result.success, false); + assert.equal(result.status, 502); + assert.match(result.error, /ComfyUI fetch output failed \(500\)/); + assert.deepEqual( + logEntries.map((entry) => entry[0]), + ["info", "error"] + ); + assert.match(logEntries[1][2], /comfyui error/i); + } finally { + globalThis.fetch = originalFetch; + globalThis.setTimeout = originalSetTimeout; + } +});