feat(reasoning): add CODEGRAPH_OFFLOAD_DISABLE kill-switch and per-call usage log

`CODEGRAPH_OFFLOAD_DISABLE=1` immediately disables the offload for the current
process without touching the persisted config or stored login — useful for A/B
arms or sessions where raw source is preferred.

`CODEGRAPH_OFFLOAD_USAGE_LOG=` appends one JSONL entry per call with token
counts, charged credits, and derived cost (`creditsCharged / 100_000`) so a
harness can attribute CodeGraph AI spend to a single run independently of the
server's cumulative totals. Both features are best-effort and never disrupt the
degradable offload path.

Also fixes the `login` credit display to check `unlimited` before the numeric
balance, so comped/internal accounts don't incorrectly show "0 remaining".
This commit is contained in:
Colby McHenry
2026-06-18 21:10:10 -05:00
parent c9e207a0f2
commit 6d5cb6b25c
4 changed files with 122 additions and 4 deletions
+59 -1
View File
@@ -34,7 +34,7 @@ describe('reasoning offload', () => {
'CODEGRAPH_OFFLOAD_URL', 'CODEGRAPH_OFFLOAD_MODEL', 'CODEGRAPH_OFFLOAD_KEY',
'CODEGRAPH_OFFLOAD_EFFORT', 'CODEGRAPH_OFFLOAD_STYLE', 'CODEGRAPH_OFFLOAD_TIMEOUT_MS',
'CODEGRAPH_OFFLOAD_MAXTOKENS', 'CODEGRAPH_OFFLOAD_STRIP', 'CODEGRAPH_OFFLOAD_DEBUG',
'CEREBRAS_API_KEY',
'CODEGRAPH_OFFLOAD_DISABLE', 'CODEGRAPH_OFFLOAD_USAGE_LOG', 'CEREBRAS_API_KEY',
];
let saved: Record<string, string | undefined>;
@@ -118,6 +118,64 @@ describe('reasoning offload', () => {
});
});
describe('CODEGRAPH_OFFLOAD_DISABLE kill-switch', () => {
it('forces the offload off even when managed + signed in', () => {
writeOffloadConfig({ managed: true });
writeOffloadToken('cgai_live');
expect(resolveOffload().enabled).toBe(true); // sanity: on without the flag
process.env.CODEGRAPH_OFFLOAD_DISABLE = '1';
const c = resolveOffload();
expect(c.enabled).toBe(false);
expect(c.managed).toBe(false);
expect(c.origin).toBe('none');
expect(isOffloadEnabled()).toBe(false);
});
it('forces the offload off even with a BYO endpoint + key', () => {
process.env.CODEGRAPH_OFFLOAD_URL = 'https://env.example/v1';
process.env.CODEGRAPH_OFFLOAD_KEY = 'sk-direct';
expect(resolveOffload().enabled).toBe(true);
process.env.CODEGRAPH_OFFLOAD_DISABLE = '1';
expect(resolveOffload().enabled).toBe(false);
});
});
describe('per-call usage log (CODEGRAPH_OFFLOAD_USAGE_LOG)', () => {
const okResponse = () => ({
ok: true, status: 200,
headers: { get: (h: string) => (h === 'x-cg-credits-charged' ? '127' : null) },
json: async () => ({
choices: [{ message: { content: 'Coverage: full.\nThe answer.' }, finish_reason: 'stop' }],
usage: { prompt_tokens: 700, completion_tokens: 80, total_tokens: 780 },
}),
});
it('appends one JSON line with tokens + charged credits when the log path is set', async () => {
writeOffloadConfig({ url: 'https://api.cerebras.ai/v1', keyEnv: 'CEREBRAS_API_KEY' });
process.env.CEREBRAS_API_KEY = 'sk-live';
vi.stubGlobal('fetch', vi.fn().mockResolvedValue(okResponse()));
const logPath = path.join(home, 'usage.jsonl');
process.env.CODEGRAPH_OFFLOAD_USAGE_LOG = logPath;
await synthesizeOffload({ query: 'q', context: 'src' });
const line = JSON.parse(fs.readFileSync(logPath, 'utf8').trim());
expect(line.totalTokens).toBe(780);
expect(line.promptTokens).toBe(700);
expect(line.creditsCharged).toBe(127);
expect(line.costUsd).toBeCloseTo(0.00127, 6); // 100k credits = $1
expect(line.answerLen).toBeGreaterThan(0);
});
it('is a no-op (and never throws) when the log path is unset', async () => {
writeOffloadConfig({ url: 'https://api.cerebras.ai/v1', keyEnv: 'CEREBRAS_API_KEY' });
process.env.CEREBRAS_API_KEY = 'sk-live';
vi.stubGlobal('fetch', vi.fn().mockResolvedValue(okResponse()));
// no CODEGRAPH_OFFLOAD_USAGE_LOG set → answer still returns fine
const out = await synthesizeOffload({ query: 'q', context: 'src' });
expect(out).toContain('Coverage: full.');
});
});
describe('strict degradation (never throws, returns null to fall back)', () => {
it('returns null when no endpoint is configured', async () => {
expect(await synthesizeOffload({ query: 'q', context: 'ctx' })).toBeNull();