Переглянути джерело

feat(offload): managed tier (CodeGraph AI) — metered reasoning via org token [WIP]

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 <t>` / `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) <noreply@anthropic.com>
Colby McHenry 5 днів тому
батько
коміт
da5c6c2f79

+ 60 - 0
__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');
+    });
+  });
 });

+ 46 - 4
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<typeof import('../index')> {
@@ -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 <token>', 'Your CodeGraph AI org token')
+  .option('--url <url>', 'Override the managed gateway URL (advanced/testing)')
+  .option('--model <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 <url>`.');
+      if (c.managed) info('Reasoning offload: managed, but signed out.  Run `codegraph offload login --token <token>`.');
+      else info('Reasoning offload: off.  Enable with `codegraph offload set-endpoint <url>` 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.');

+ 45 - 23
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',

+ 43 - 0
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<string, unknown> {
+  try {
+    return JSON.parse(fs.readFileSync(credentialsPath(), 'utf8')) as Record<string, unknown>;
+  } 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 */ }
+}

+ 35 - 0
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<OffloadUsage | null> {
+  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.