Skip to content
Merged
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
17 changes: 2 additions & 15 deletions open-sse/handlers/chatCore.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { extractSystemRoleMessages } from "./chatCore/claudeSystemRole.ts";
export { extractSystemRoleMessages } from "./chatCore/claudeSystemRole.ts";
import { checkIdempotencyCache } from "./chatCore/idempotency.ts";
import { checkSemanticCache } from "./chatCore/semanticCache.ts";
import { applyClientUsageBuffer } from "./chatCore/clientUsageBuffer.ts";
import { sanitizeChatRequestBody } from "./chatCore/sanitization.ts";
import {
getHeaderValueCaseInsensitive,
Expand Down Expand Up @@ -68,7 +69,6 @@ import { resolveStreamReadinessTimeout } from "../utils/streamReadinessPolicy.ts
import { createStreamController, pipeWithDisconnect } from "../utils/streamHandler.ts";
import * as streamFailure from "../utils/streamFailureFinalization.ts";
import { createSseHeartbeatTransform, shapeForClientFormat } from "../utils/sseHeartbeat.ts";
import { addBufferToUsage, filterUsageForFormat, estimateUsage } from "../utils/usageTracking.ts";
import {
refreshWithRetry,
isUnrecoverableRefreshError,
Expand Down Expand Up @@ -3470,20 +3470,7 @@ export async function handleChatCore({
translatedResponse = sanitizeOpenAIResponse(translatedResponse);
}

// Add buffer and filter usage for client (to prevent CLI context errors)
if (translatedResponse?.usage) {
const buffered = addBufferToUsage(translatedResponse.usage);
translatedResponse.usage = filterUsageForFormat(buffered, clientResponseFormat);
} else {
// Fallback: estimate usage when provider returned no usage block
const contentLength = JSON.stringify(
translatedResponse?.choices?.[0]?.message?.content || ""
).length;
if (contentLength > 0) {
const estimated = estimateUsage(body, contentLength, clientResponseFormat);
translatedResponse.usage = filterUsageForFormat(estimated, clientResponseFormat);
}
}
applyClientUsageBuffer(translatedResponse, body, clientResponseFormat);

if (memoryOwnerId && memorySettings?.enabled && memorySettings.maxTokens > 0) {
const requestMemoryText = extractMemoryTextFromRequestBody(body as Record<string, unknown>);
Expand Down
54 changes: 54 additions & 0 deletions open-sse/handlers/chatCore/clientUsageBuffer.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
/**
* chatCore client usage buffer/estimate (Quality Gate v2 / Fase 9 — chatCore god-file
* decomposition, #3501).
*
* Extracted from handleChatCore's non-streaming success path: add a buffer to the response usage
* and filter it for the client format (to prevent CLI context errors); if the provider returned no
* usage block, fall back to estimating from the serialized content length. Mutates
* `translatedResponse.usage` in place — byte-identical to the previous inline block, including the
* `?.usage` guard, the `JSON.stringify(... || "")` content-length, and the `> 0` estimate gate.
*/
import {
addBufferToUsage as defaultAddBuffer,
filterUsageForFormat as defaultFilterUsage,
estimateUsage as defaultEstimateUsage,
} from "../../utils/usageTracking.ts";

type ResponseLike = {
usage?: unknown;
choices?: Array<{ message?: { content?: unknown } }>;
} | null | undefined;

export interface ClientUsageBufferDeps {
addBufferToUsage: typeof defaultAddBuffer;
filterUsageForFormat: typeof defaultFilterUsage;
estimateUsage: typeof defaultEstimateUsage;
}

const DEFAULT_DEPS: ClientUsageBufferDeps = {
addBufferToUsage: defaultAddBuffer,
filterUsageForFormat: defaultFilterUsage,
estimateUsage: defaultEstimateUsage,
};

export function applyClientUsageBuffer(
translatedResponse: ResponseLike,
body: unknown,
clientResponseFormat: unknown,
deps: ClientUsageBufferDeps = DEFAULT_DEPS
): void {
// Add buffer and filter usage for client (to prevent CLI context errors)
if (translatedResponse?.usage) {
const buffered = deps.addBufferToUsage(translatedResponse.usage);
translatedResponse.usage = deps.filterUsageForFormat(buffered, clientResponseFormat);
} else {
// Fallback: estimate usage when provider returned no usage block
const contentLength = JSON.stringify(
translatedResponse?.choices?.[0]?.message?.content || ""
).length;
if (contentLength > 0) {
const estimated = deps.estimateUsage(body, contentLength, clientResponseFormat);
translatedResponse.usage = deps.filterUsageForFormat(estimated, clientResponseFormat);
}
}
}
78 changes: 78 additions & 0 deletions tests/unit/chatcore-client-usage-buffer.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
// Characterization of applyClientUsageBuffer — the non-streaming usage buffer/estimate block
// extracted from handleChatCore (chatCore god-file decomposition, #3501). Deps are injected so the
// buffer-vs-estimate branch and the in-place mutation of translatedResponse.usage are observable.
// Locks: usage present → buffer+filter; usage absent → estimate from content length; empty content
// (length 2 from JSON.stringify("")) still estimates; the mutation target is translatedResponse.
import { test } from "node:test";
import assert from "node:assert/strict";

const { applyClientUsageBuffer } = await import(
"../../open-sse/handlers/chatCore/clientUsageBuffer.ts"
);

function makeDeps(overrides: Record<string, unknown> = {}) {
const calls = { buffer: [] as unknown[], estimate: [] as unknown[], filter: [] as unknown[] };
const deps = {
addBufferToUsage: (u: unknown) => {
calls.buffer.push(u);
return { ...(u as object), _buffered: true };
},
estimateUsage: (...a: unknown[]) => {
calls.estimate.push(a);
return { _estimated: true };
},
filterUsageForFormat: (u: unknown, _fmt: unknown) => {
calls.filter.push(u);
return { ...(u as object), _filtered: true };
},
...overrides,
} as Parameters<typeof applyClientUsageBuffer>[3];
return { deps, calls };
}

test("usage present → buffer then filter, mutates in place", () => {
const { deps, calls } = makeDeps();
const resp: Record<string, unknown> = { usage: { prompt_tokens: 5 } };
applyClientUsageBuffer(resp, { messages: [] }, "openai", deps);
assert.equal(calls.buffer.length, 1);
assert.equal(calls.estimate.length, 0);
assert.equal((resp.usage as Record<string, unknown>)._buffered, true);
assert.equal((resp.usage as Record<string, unknown>)._filtered, true);
});

test("no usage but content present → estimate then filter", () => {
const { deps, calls } = makeDeps();
const resp: Record<string, unknown> = {
choices: [{ message: { content: "hello world" } }],
};
applyClientUsageBuffer(resp, { messages: [] }, "openai", deps);
assert.equal(calls.buffer.length, 0);
assert.equal(calls.estimate.length, 1);
assert.equal((resp.usage as Record<string, unknown>)._estimated, true);
assert.equal((resp.usage as Record<string, unknown>)._filtered, true);
// estimateUsage receives (body, contentLength, format)
const args = calls.estimate[0] as unknown[];
assert.equal(args[2], "openai");
assert.equal(typeof args[1], "number");
});

test("empty content → JSON.stringify('') length 2 > 0 still estimates", () => {
const { deps, calls } = makeDeps();
const resp: Record<string, unknown> = {};
applyClientUsageBuffer(resp, {}, "claude", deps);
// content "" → JSON.stringify("") = '""' length 2 → contentLength 2 > 0
assert.equal(calls.estimate.length, 1);
const args = calls.estimate[0] as unknown[];
assert.equal(args[1], 2);
});

test("content length is computed from choices[0].message.content", () => {
const { deps, calls } = makeDeps();
const resp: Record<string, unknown> = {
choices: [{ message: { content: "abc" } }],
};
applyClientUsageBuffer(resp, {}, "openai", deps);
// JSON.stringify("abc") = '"abc"' → length 5
const args = calls.estimate[0] as unknown[];
assert.equal(args[1], 5);
});
Loading