From da5c6c2f7951e36e52c0eec2817938a5203a97d6 Mon Sep 17 00:00:00 2001 From: Colby McHenry Date: Wed, 17 Jun 2026 14:54:02 -0500 Subject: [PATCH] =?UTF-8?q?feat(offload):=20managed=20tier=20(CodeGraph=20?= =?UTF-8?q?AI)=20=E2=80=94=20metered=20reasoning=20via=20org=20token=20[WI?= =?UTF-8?q?P]?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the managed offload mode: point codegraph_explore at the CodeGraph AI metered gateway (https://ai.getcodegraph.com) with an org token instead of a BYO provider key. Same synthesis client, pointed at codegraph-ai-proxy (a metered OpenAI-compatible gateway). - credentials.ts — org token in ~/.codegraph/credentials.json (0600); unlike a BYO provider key it's a revocable org-scoped auth token (gh/npm-login style), kept out of config.json - config.ts — managed branch in resolveOffload: default gateway URL + public model id (openai/gpt-oss-120b) + login token as bearer; managed requires a token to be enabled - reasoner.ts — fetchUsage() reads the credit balance from /v1/usage - bin/codegraph.ts — `codegraph offload login --token ` / `logout`; status shows the managed tier + live balance Proven GREEN end-to-end against a local wrangler-dev of the proxy: org token validated, credits prechecked, real Cerebras synthesis returned, and credits metered + charged (250,000 → 248,473). Graceful degrade on upstream failure; balance via /v1/usage. Phase 3 (codegraph login device flow) replaces the manual --token. Co-Authored-By: Claude Opus 4.8 (1M context) --- __tests__/offload.test.ts | 60 +++++++++++++++++++++++++++++++ src/bin/codegraph.ts | 50 +++++++++++++++++++++++--- src/reasoning/config.ts | 68 ++++++++++++++++++++++++------------ src/reasoning/credentials.ts | 43 +++++++++++++++++++++++ src/reasoning/reasoner.ts | 35 +++++++++++++++++++ 5 files changed, 229 insertions(+), 27 deletions(-) create mode 100644 src/reasoning/credentials.ts diff --git a/__tests__/offload.test.ts b/__tests__/offload.test.ts index 886d06d..67c2b92 100644 --- a/__tests__/offload.test.ts +++ b/__tests__/offload.test.ts @@ -18,7 +18,10 @@ import { readOffloadConfig, writeOffloadConfig, resolveOffload, + MANAGED_DEFAULT_URL, + MANAGED_DEFAULT_MODEL, } from '../src/reasoning/config'; +import { readOffloadToken, writeOffloadToken } from '../src/reasoning/credentials'; import { isOffloadEnabled, synthesizeOffload, stripAgentDirectives } from '../src/reasoning/reasoner'; describe('reasoning offload', () => { @@ -183,4 +186,61 @@ describe('reasoning offload', () => { expect(stripped).toContain('code body'); }); }); + + describe('managed tier (CodeGraph AI)', () => { + it('stores the org token at 0600 in credentials.json, not in config.json', () => { + writeOffloadConfig({ managed: true }); + writeOffloadToken('cgai_secrettoken'); + expect(readOffloadToken()).toBe('cgai_secrettoken'); + + // config.json carries the managed flag but NOT the token. + const cfg = fs.readFileSync(path.join(home, '.codegraph', 'config.json'), 'utf8'); + expect(cfg).toContain('managed'); + expect(cfg).not.toContain('cgai_secrettoken'); + + const credPath = path.join(home, '.codegraph', 'credentials.json'); + expect(fs.readFileSync(credPath, 'utf8')).toContain('cgai_secrettoken'); + // POSIX perms must be owner-only (0600). (Windows has no POSIX mode bits.) + if (process.platform !== 'win32') { + expect(fs.statSync(credPath).mode & 0o777).toBe(0o600); + } + }); + + it('resolves managed mode to the gateway URL + public model id + login token', () => { + writeOffloadConfig({ managed: true }); + writeOffloadToken('cgai_live'); + const c = resolveOffload(); + expect(c.enabled).toBe(true); + expect(c.managed).toBe(true); + expect(c.url).toBe(MANAGED_DEFAULT_URL); + expect(c.model).toBe(MANAGED_DEFAULT_MODEL); + expect(c.apiKey).toBe('cgai_live'); + expect(c.keySource).toBe('codegraph login'); + }); + + it('is NOT enabled when managed but signed out (no token)', () => { + writeOffloadConfig({ managed: true }); + const c = resolveOffload(); + expect(c.managed).toBe(true); + expect(c.enabled).toBe(false); // url defaults, but no token → effectively logged out + expect(isOffloadEnabled()).toBe(false); + }); + + it('clears the token on logout', () => { + writeOffloadToken('cgai_live'); + writeOffloadToken(null); + expect(readOffloadToken()).toBeUndefined(); + }); + + it('lets env override the managed endpoint and token (for testing)', () => { + writeOffloadConfig({ managed: true }); + writeOffloadToken('cgai_stored'); + process.env.CODEGRAPH_OFFLOAD_URL = 'http://localhost:8787/v1'; + process.env.CODEGRAPH_OFFLOAD_KEY = 'cgai_env'; + const c = resolveOffload(); + expect(c.url).toBe('http://localhost:8787/v1'); + expect(c.apiKey).toBe('cgai_env'); + expect(c.keySource).toBe('CODEGRAPH_OFFLOAD_KEY'); + }); + }); }); diff --git a/src/bin/codegraph.ts b/src/bin/codegraph.ts index e358a1e..000414e 100644 --- a/src/bin/codegraph.ts +++ b/src/bin/codegraph.ts @@ -37,6 +37,8 @@ import { relaunchWithWasmRuntimeFlagsIfNeeded } from '../extraction/wasm-runtime import { EXTRACTION_VERSION } from '../extraction/extraction-version'; import { getTelemetry, TELEMETRY_DOCS, recordIndexEvent } from '../telemetry'; import { writeOffloadConfig, resolveOffload } from '../reasoning/config'; +import { writeOffloadToken } from '../reasoning/credentials'; +import { fetchUsage } from '../reasoning/reasoner'; // Lazy-load heavy modules (CodeGraph, runInstaller) to keep CLI startup fast. async function loadCodeGraph(): Promise { @@ -1382,12 +1384,52 @@ offloadCmd }); offloadCmd - .command('status') - .description('Show the current reasoning-offload configuration') + .command('login') + .description('Use the managed CodeGraph AI tier (metered) with your account token') + .requiredOption('--token ', 'Your CodeGraph AI org token') + .option('--url ', 'Override the managed gateway URL (advanced/testing)') + .option('--model ', 'Override the model id') + .action((opts: { token: string; url?: string; model?: string }) => { + // Phase 2: the token is pasted in. A future `codegraph login` device flow will + // mint and store it automatically. + writeOffloadConfig({ managed: true, url: opts.url, model: opts.model }); + writeOffloadToken(opts.token); + success('Reasoning offload: signed in to CodeGraph AI (managed).'); + info(' Credits burn from your account. Check the balance with `codegraph offload status`.'); + info(' Restart your editor/agent session for running MCP servers to pick it up.'); + }); + +offloadCmd + .command('logout') + .description('Sign out of CodeGraph AI and clear the stored token') .action(() => { + writeOffloadToken(null); + writeOffloadConfig(null); + success('Signed out of CodeGraph AI; offload turned off.'); + }); + +offloadCmd + .command('status') + .description('Show the current reasoning-offload configuration (and managed balance)') + .action(async () => { const c = resolveOffload(); if (!c.enabled) { - info('Reasoning offload: off. Enable with `codegraph offload set-endpoint `.'); + if (c.managed) info('Reasoning offload: managed, but signed out. Run `codegraph offload login --token `.'); + else info('Reasoning offload: off. Enable with `codegraph offload set-endpoint ` or `codegraph offload login`.'); + return; + } + if (c.managed) { + success(`Reasoning offload: on — CodeGraph AI (managed)`); + info(` endpoint: ${c.url}`); + info(` model: ${c.model}`); + info(` token: present (from ${c.keySource})`); + const usage = await fetchUsage(); + if (usage && typeof usage.remaining === 'number') { + const reset = usage.periodEnd ? ` · allowance resets ${new Date(usage.periodEnd).toISOString().slice(0, 10)}` : ''; + info(` credits: ${usage.remaining.toLocaleString()} remaining (plan ${usage.plan ?? '—'})${reset}`); + } else { + warn(' credits: could not reach CodeGraph AI to read your balance (the offload still degrades gracefully).'); + } return; } success(`Reasoning offload: on (${c.origin === 'env' ? 'from environment' : 'configured'})`); @@ -1400,7 +1442,7 @@ offloadCmd offloadCmd .command('disable') - .description('Turn off the reasoning offload') + .description('Turn off the reasoning offload (keeps any saved login token)') .action(() => { writeOffloadConfig(null); success('Reasoning offload disabled.'); diff --git a/src/reasoning/config.ts b/src/reasoning/config.ts index 0821107..fc04026 100644 --- a/src/reasoning/config.ts +++ b/src/reasoning/config.ts @@ -8,20 +8,30 @@ * Every codegraph MCP server on the machine picks it up, so a user configures it * once. Env vars override the file (CI / ephemeral / advanced use). * - * The API key is NEVER written to disk. The CLI stores the NAME of an env var - * that holds it (`keyEnv`); at call time the key is read from that env var (or - * directly from `CODEGRAPH_OFFLOAD_KEY`). So the config file carries no secret. + * For a BYO endpoint, the API key is NEVER written to disk: the CLI stores the + * NAME of an env var (`keyEnv`) and reads the key from it at call time. The + * MANAGED tier ("CodeGraph AI") instead authenticates with a revocable, org-scoped + * token from `codegraph offload login`, stored separately in `credentials.json` + * (see ./credentials) — so `config.json` itself never carries a secret either way. */ import * as fs from 'fs'; import * as path from 'path'; import * as os from 'os'; +import { readOffloadToken } from './credentials'; + +/** Managed tier ("CodeGraph AI") — the metered gateway used when logged in. */ +export const MANAGED_DEFAULT_URL = 'https://ai.getcodegraph.com/v1'; +/** The gateway's public model id (it translates this to the upstream provider id). */ +export const MANAGED_DEFAULT_MODEL = 'openai/gpt-oss-120b'; export interface OffloadConfig { + /** Managed tier: route through CodeGraph AI (metered) with the logged-in org token. */ + managed?: boolean; /** OpenAI-compatible base URL ending in `/v1` (e.g. https://api.cerebras.ai/v1). */ url?: string; - /** Model id to request (default `gpt-oss-120b`). */ + /** Model id to request (default `gpt-oss-120b` BYO, `openai/gpt-oss-120b` managed). */ model?: string; - /** Name of the env var holding the provider API key (the key itself is never persisted). */ + /** Name of the env var holding the provider API key (never persisted). BYO only. */ keyEnv?: string; /** reasoning_effort: low | medium | high (default `low`). */ effort?: string; @@ -30,13 +40,15 @@ export interface OffloadConfig { } export interface ResolvedOffload { - /** True when a reasoning endpoint is configured (by env or by file). */ + /** True when the offload is usable (endpoint present; for managed, a token too). */ enabled: boolean; + /** Managed tier (CodeGraph AI, metered) vs BYO endpoint. */ + managed: boolean; url?: string; model: string; - /** Resolved API key (from `CODEGRAPH_OFFLOAD_KEY` or the configured `keyEnv`), if any. */ + /** Resolved API key / org token (from env, the configured `keyEnv`, or login), if any. */ apiKey?: string; - /** Which env var the key came from (for `status` display) — never the key itself. */ + /** Where the key/token came from (for `status` display) — never the secret itself. */ keySource?: string; effort: string; style: string; @@ -91,29 +103,39 @@ const trimmed = (v: string | undefined): string | undefined => { /** Merge the persisted config with `CODEGRAPH_OFFLOAD_*` env overrides (env wins). */ export function resolveOffload(env: NodeJS.ProcessEnv = process.env): ResolvedOffload { const c = readOffloadConfig(); - const url = trimmed(env.CODEGRAPH_OFFLOAD_URL) ?? trimmed(c.url); + const managed = !!c.managed; + const envUrl = trimmed(env.CODEGRAPH_OFFLOAD_URL); + const envKey = trimmed(env.CODEGRAPH_OFFLOAD_KEY); - // Key: direct env var first, else the configured env-var name. Never from disk. + let url: string | undefined; let apiKey: string | undefined; let keySource: string | undefined; - if (trimmed(env.CODEGRAPH_OFFLOAD_KEY)) { - apiKey = trimmed(env.CODEGRAPH_OFFLOAD_KEY); - keySource = 'CODEGRAPH_OFFLOAD_KEY'; - } else if (c.keyEnv && trimmed(env[c.keyEnv])) { - apiKey = trimmed(env[c.keyEnv]); - keySource = c.keyEnv; + let model: string; + + if (managed) { + // Managed tier: default to the CodeGraph AI gateway + its public model id; the + // bearer is the org token from `codegraph offload login` (or an env override). + url = envUrl ?? trimmed(c.url) ?? MANAGED_DEFAULT_URL; + model = trimmed(env.CODEGRAPH_OFFLOAD_MODEL) ?? trimmed(c.model) ?? MANAGED_DEFAULT_MODEL; + if (envKey) { apiKey = envKey; keySource = 'CODEGRAPH_OFFLOAD_KEY'; } + else { const t = readOffloadToken(); if (t) { apiKey = t; keySource = 'codegraph login'; } } + } else { + // BYO: endpoint + (optional) provider key resolved from env or the named env var. + url = envUrl ?? trimmed(c.url); + model = trimmed(env.CODEGRAPH_OFFLOAD_MODEL) ?? trimmed(c.model) ?? 'gpt-oss-120b'; + if (envKey) { apiKey = envKey; keySource = 'CODEGRAPH_OFFLOAD_KEY'; } + else if (c.keyEnv && trimmed(env[c.keyEnv])) { apiKey = trimmed(env[c.keyEnv]); keySource = c.keyEnv; } } - const origin: ResolvedOffload['origin'] = trimmed(env.CODEGRAPH_OFFLOAD_URL) - ? 'env' - : trimmed(c.url) - ? 'config' - : 'none'; + const origin: ResolvedOffload['origin'] = envUrl ? 'env' : (managed || trimmed(c.url)) ? 'config' : 'none'; return { - enabled: !!url, + // Managed needs both an endpoint AND a token (no token → effectively logged out); + // BYO needs only an endpoint (some endpoints require no auth). + enabled: managed ? (!!url && !!apiKey) : !!url, + managed, url, - model: trimmed(env.CODEGRAPH_OFFLOAD_MODEL) ?? trimmed(c.model) ?? 'gpt-oss-120b', + model, apiKey, keySource, effort: trimmed(env.CODEGRAPH_OFFLOAD_EFFORT) ?? trimmed(c.effort) ?? 'low', diff --git a/src/reasoning/credentials.ts b/src/reasoning/credentials.ts new file mode 100644 index 0000000..44d1df6 --- /dev/null +++ b/src/reasoning/credentials.ts @@ -0,0 +1,43 @@ +/** + * Managed-offload credentials: the CodeGraph org token that authenticates the + * managed reasoning tier against `codegraph-ai` (the metered gateway). + * + * Unlike a BYO provider key (which is never persisted — the config stores only the + * NAME of an env var), the org token IS a revocable, org-scoped auth token issued + * to this machine — like the token `gh auth` or `npm login` stores. So it lives in + * its own file, `~/.codegraph/credentials.json`, written `0600`, kept out of the + * shareable `config.json`. + */ +import * as fs from 'fs'; +import * as path from 'path'; +import * as os from 'os'; + +function credentialsPath(): string { + return path.join(os.homedir(), '.codegraph', 'credentials.json'); +} + +function read(): Record { + try { + return JSON.parse(fs.readFileSync(credentialsPath(), 'utf8')) as Record; + } catch { + return {}; + } +} + +/** The stored managed-offload org token, if the machine is logged in. */ +export function readOffloadToken(): string | undefined { + const t = read().offloadToken; + return typeof t === 'string' && t.trim() ? t.trim() : undefined; +} + +/** Persist (or, with `null`, clear) the managed-offload org token at `0600`. */ +export function writeOffloadToken(token: string | null): void { + const p = credentialsPath(); + fs.mkdirSync(path.dirname(p), { recursive: true }); + const creds = read(); + if (token === null) delete creds.offloadToken; + else creds.offloadToken = token; + // Write restrictively: create at 0600, and tighten an existing file too. + fs.writeFileSync(p, JSON.stringify(creds, null, 2) + '\n', { mode: 0o600 }); + try { fs.chmodSync(p, 0o600); } catch { /* best-effort on platforms without POSIX modes */ } +} diff --git a/src/reasoning/reasoner.ts b/src/reasoning/reasoner.ts index b8fd160..55be192 100644 --- a/src/reasoning/reasoner.ts +++ b/src/reasoning/reasoner.ts @@ -40,6 +40,41 @@ export function isOffloadEnabled(): boolean { return resolveOffload().enabled; } +export interface OffloadUsage { + plan?: string; + allowance?: number; + used?: number; + overage?: number; + remaining?: number; + periodEnd?: number; + models?: string[]; +} + +/** + * GET `/v1/usage` from the configured (managed) endpoint → the org's credit + * balance/usage, or null on any failure. Drives `codegraph offload status`. + */ +export async function fetchUsage(): Promise { + const cfg = resolveOffload(); + if (!cfg.url || !cfg.apiKey) return null; + const url = cfg.url.replace(/\/+$/, '') + '/usage'; + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), 10000); + try { + const res = await fetch(url, { + headers: { authorization: `Bearer ${cfg.apiKey}` }, + signal: controller.signal, + }); + if (!res.ok) { debug('usage not ok', res.status); return null; } + return (await res.json()) as OffloadUsage; + } catch (err) { + debug('usage error', (err as Error)?.message); + return null; + } finally { + clearTimeout(timer); + } +} + function debug(...args: unknown[]): void { if (process.env.CODEGRAPH_OFFLOAD_DEBUG === '1') { // stderr only — stdout is the MCP JSON-RPC transport.