Merge pull request #109 from AxmeAI/feat/auth-mode-selection-20260416

feat(auth): user-selectable auth mode (subscription vs api_key)

Author: George-iam

src/agents/memory-extractor.ts

@@ -12,7 +12,7 @@
 import type { Memory } from "../types.js";
 import { extractCostFromResult, zeroCost, type CostInfo } from "../utils/cost-extractor.js";
 import { toMemorySlug } from "../storage/memory.js";
-import { findClaudePath } from "../utils/agent-options.js";
+import { buildAgentEnv, findClaudePath } from "../utils/agent-options.js";
 
 export interface MemoryExtractionResult {
   memories: Memory[];
@@ -78,6 +78,7 @@ export async function runMemoryExtraction(opts: {
     allowDangerouslySkipPermissions: true,
     allowedTools: [] as string[],
     disallowedTools: ["Write", "Edit", "Bash", "Glob", "Grep", "Read", "Agent", "NotebookEdit", "Skill", "TodoWrite"],
+    env: buildAgentEnv(),
   };
 
   const prompt = `${EXTRACTION_PROMPT}\n\nSession ID: ${opts.sessionId}\n\nSession transcript:\n${opts.sessionEvents}`;

src/agents/session-auditor.ts

@@ -21,7 +21,7 @@ import { basename, relative } from "node:path";
 import type { Memory, Decision, SessionHandoff, WorkspaceInfo } from "../types.js";
 import { DEFAULT_AUDITOR_MODEL } from "../types.js";
 import { extractCostFromResult, zeroCost, type CostInfo } from "../utils/cost-extractor.js";
-import { findClaudePath } from "../utils/agent-options.js";
+import { buildAgentEnv, findClaudePath } from "../utils/agent-options.js";
 import { toMemorySlug } from "../storage/memory.js";
 import { toSlug, listDecisions } from "../storage/decisions.js";
 import { listMemories } from "../storage/memory.js";
@@ -647,7 +647,7 @@ async function runSingleAuditCall(opts: {
     // auto-loading the project's .claude/settings.json, but users or CI may
     // register hooks via environment or other means, so the belt-and-braces
     // env check in every hook handler is what actually stops the recursion.
-    env: { ...process.env, AXME_SKIP_HOOKS: "1", AXME_TELEMETRY_DISABLED: "1" },
+    env: buildAgentEnv(),
   };
 
   const isMultiChunk = opts.totalChunks > 1;
@@ -901,7 +901,7 @@ ${freeTextAnalysis}`;
       "Read", "Grep", "Glob", "Write", "Edit", "NotebookEdit", "Agent",
       "Skill", "TodoWrite", "WebFetch", "WebSearch", "Bash", "ToolSearch",
     ],
-    env: { ...process.env, AXME_SKIP_HOOKS: "1", AXME_TELEMETRY_DISABLED: "1" },
+    env: buildAgentEnv(),
   };
 
   const q = sdk.query({ prompt: formatPrompt, options: queryOpts });

src/cli.ts

@@ -16,7 +16,10 @@ import { statusTool } from "./tools/status.js";
 import { detectWorkspace } from "./utils/workspace-detector.js";
 import { atomicWrite, ensureDir } from "./storage/engine.js";
 import { saveMemory, toMemorySlug } from "./storage/memory.js";
-import type { WorkspaceInfo } from "./types.js";
+import { detectAuthOptions } from "./utils/auth-detect.js";
+import { authConfigPath, loadAuthConfig, saveAuthConfig } from "./utils/auth-config.js";
+import { formatDetectionBlock, hasAnyAuth, promptAuthChoice } from "./utils/auth-prompt.js";
+import type { AuthMode, WorkspaceInfo } from "./types.js";
 import { AXME_CODE_DIR } from "./types.js";
 
 const args = process.argv.slice(2);
@@ -150,6 +153,49 @@ function hasAuth(): boolean {
   return false;
 }
 
+/**
+ * Print detection block + saved choice (if any).
+ */
+function printAuthStatus(): void {
+  const options = detectAuthOptions();
+  console.log(formatDetectionBlock(options));
+  const saved = loadAuthConfig();
+  if (saved) {
+    console.log(`\nCurrent mode: ${saved.mode} (saved ${saved.chosenAt})`);
+    console.log(`Config file:  ${authConfigPath()}`);
+  } else {
+    console.log("\nCurrent mode: not configured (using heuristic fallback)");
+    console.log(`Config file:  ${authConfigPath()} (will be created on first choice)`);
+  }
+}
+
+/**
+ * Called from `axme-code setup` before LLM scanners launch. Persists a user
+ * choice to ~/.config/axme-code/auth.yaml exactly once on first setup, so
+ * subsequent runs (and all MCP server scanner/auditor subprocesses) use the
+ * same mode without re-prompting. In non-TTY contexts (CI, scripts) we skip
+ * the prompt and let `resolveAuthMode()` fall back to its heuristic — we do
+ * NOT persist a guessed choice silently.
+ */
+async function ensureAuthConfiguredForSetup(): Promise<void> {
+  if (loadAuthConfig()) return;
+  if (!process.stdin.isTTY) return;
+
+  const options = detectAuthOptions();
+  if (!hasAnyAuth(options)) return; // hasAuth preflight will fail anyway
+
+  console.log("\nAuthentication setup for LLM scanners");
+  console.log(formatDetectionBlock(options));
+  console.log("");
+  const choice = await promptAuthChoice(options);
+  if (!choice) {
+    console.log("  Auth selection cancelled. Heuristic fallback will be used.");
+    return;
+  }
+  saveAuthConfig(choice);
+  console.log(`  Saved auth mode: ${choice} (${authConfigPath()})`);
+}
+
 function generateWorkspaceYaml(workspacePath: string, ws: WorkspaceInfo): void {
   const wsYaml = yaml.dump({
     name: workspacePath.split("/").pop(),
@@ -256,6 +302,9 @@ Usage:
   axme-code setup [path] [--force]         Initialize project (LLM scan + .mcp.json + CLAUDE.md)
   axme-code serve                         Start MCP server (stdio transport)
   axme-code status [path]                 Show project status
+  axme-code auth                          Re-detect and choose auth mode (subscription/api_key)
+  axme-code auth status                   Show current auth mode + detected options
+  axme-code auth use <subscription|api_key>  Set auth mode non-interactively
   axme-code cleanup legacy-artifacts [--dry-run]  Remove pre-PR#7 sessions/logs
   axme-code cleanup decisions-normalize [--dry-run]  Add status:active to decisions
   axme-code audit-kb [path] [--all-repos]             KB audit: dedup, conflicts, compaction
@@ -339,6 +388,11 @@ async function main() {
         process.exit(1);
       }
 
+      // Auth mode selection — prompt once on first interactive setup, persist
+      // to ~/.config/axme-code/auth.yaml. Later runs and scanner subprocesses
+      // read the saved mode via resolveAuthMode() in buildAgentEnv().
+      await ensureAuthConfiguredForSetup();
+
       // Init with LLM scanners (parallel)
       try {
         if (isWorkspace) {
@@ -649,6 +703,51 @@ Do NOT skip — without context you will miss critical project rules.
       break;
     }
 
+    case "auth": {
+      const sub = args[1];
+      if (sub === "status" || sub === "show") {
+        printAuthStatus();
+        break;
+      }
+      if (sub === "use" || sub === "set") {
+        const mode = args[2];
+        if (mode !== "subscription" && mode !== "api_key") {
+          console.error("Usage: axme-code auth use <subscription|api_key>");
+          process.exit(1);
+        }
+        saveAuthConfig(mode as AuthMode);
+        console.log(`Saved auth mode: ${mode} (${authConfigPath()})`);
+        break;
+      }
+      if (sub === undefined || sub === "choose") {
+        const options = detectAuthOptions();
+        console.log("Authentication setup for LLM scanners");
+        console.log(formatDetectionBlock(options));
+        const saved = loadAuthConfig();
+        if (saved) console.log(`\nCurrent mode: ${saved.mode} (saved ${saved.chosenAt})`);
+        console.log("");
+        if (!hasAnyAuth(options)) {
+          console.error("No authentication detected. Set ANTHROPIC_API_KEY or run `claude /login`, then re-run `axme-code auth`.");
+          process.exit(1);
+        }
+        if (!process.stdin.isTTY) {
+          console.error("`axme-code auth` requires an interactive terminal. Use `axme-code auth use <subscription|api_key>` non-interactively.");
+          process.exit(1);
+        }
+        const choice = await promptAuthChoice(options);
+        if (!choice) {
+          console.log("Cancelled. No change.");
+          break;
+        }
+        saveAuthConfig(choice);
+        console.log(`Saved auth mode: ${choice} (${authConfigPath()})`);
+        break;
+      }
+      console.error(`Unknown 'auth' subcommand: ${sub}`);
+      console.error("Available: (none)|choose, status|show, use|set <subscription|api_key>");
+      process.exit(1);
+    }
+
     case "help":
     case "--help":
     case "-h":

src/types.ts

@@ -370,3 +370,12 @@ export const DEFAULT_AGENT_PERMISSIONS: AgentPermissions = {
 };
 
 export type E2EMode = "after-task" | "after-stage" | "manual";
+
+// --- Auth ---
+
+export type AuthMode = "subscription" | "api_key";
+
+export interface AuthConfig {
+  mode: AuthMode;
+  chosenAt: string;
+}

src/utils/agent-options.ts

@@ -3,6 +3,7 @@
  */
 
 import { execSync } from "node:child_process";
+import { resolveAuthMode } from "./auth-config.js";
 
 type Options = import("@anthropic-ai/claude-agent-sdk").Options;
 
@@ -54,6 +55,35 @@ const ROLE_TOOLS: Record<AgentRole, { allowed: string[]; disallowed: string[] }>
   },
 };
 
+/**
+ * Build the env passed to every Claude Code subprocess we spawn for LLM work.
+ *
+ * Two things happen here:
+ *   1. `AXME_TELEMETRY_DISABLED` and `AXME_SKIP_HOOKS` are set to suppress
+ *      recursive startup events and ghost AXME sessions when the sub-claude
+ *      inadvertently launches axme-code as its own MCP server.
+ *   2. If the user has selected `subscription` as the auth mode (either via
+ *      `axme-code auth` / `axme-code setup`, or by heuristic when only the
+ *      subscription is detected), we delete `ANTHROPIC_API_KEY` before
+ *      handing env to the subprocess. Claude Code checks the env var before
+ *      its OAuth credentials, so leaving an empty-balance key in env would
+ *      surface as "Credit balance is too low" or 401 auth errors even when
+ *      the user has an active subscription. Delete, not empty string: Claude
+ *      Code treats an empty-string value as "set" and still prefers it over
+ *      OAuth.
+ */
+export function buildAgentEnv(): NodeJS.ProcessEnv {
+  const env: NodeJS.ProcessEnv = {
+    ...process.env,
+    AXME_TELEMETRY_DISABLED: "1",
+    AXME_SKIP_HOOKS: "1",
+  };
+  if (resolveAuthMode() === "subscription") {
+    delete env.ANTHROPIC_API_KEY;
+  }
+  return env;
+}
+
 export function buildAgentQueryOptions(base: {
   cwd: string;
   model: string;
@@ -78,11 +108,6 @@ export function buildAgentQueryOptions(base: {
     allowedTools: tools.allowed,
     disallowedTools: tools.disallowed,
     includePartialMessages: true,
-    // Disable telemetry in spawned subprocesses. Sub-claude sessions started
-    // by scanners/auditors may pick up the parent's .mcp.json and re-launch
-    // axme-code as an MCP server. Each re-launch would otherwise fire its
-    // own startup event, inflating DAU and skewing scanner cost metrics.
-    // The parent process owns the lifecycle event for this user action.
-    env: { ...process.env, AXME_TELEMETRY_DISABLED: "1", AXME_SKIP_HOOKS: "1" },
+    env: buildAgentEnv(),
   };
 }

src/utils/auth-config.ts

@@ -0,0 +1,77 @@
+/**
+ * User-level authentication mode config at ~/.config/axme-code/auth.yaml.
+ *
+ * Auth mode is a per-machine concern (which credential should Claude Code
+ * subprocesses use), not a per-project one, so it lives outside the repo's
+ * .axme-code/ storage. One choice applies to every project on this machine.
+ *
+ * The file stores only the selected mode and the timestamp of the choice.
+ * `resolveAuthMode()` returns the persisted mode when present, otherwise
+ * falls back to a detection-based heuristic without writing anything — so
+ * non-interactive callers (scanner subprocesses, auditor) never surprise
+ * the user by persisting a guessed choice.
+ */
+
+import { homedir } from "node:os";
+import { join } from "node:path";
+import yaml from "js-yaml";
+import { atomicWrite, ensureDir, readSafe, pathExists } from "../storage/engine.js";
+import type { AuthConfig, AuthMode } from "../types.js";
+import { detectAuthOptions, type AuthOptions } from "./auth-detect.js";
+
+/**
+ * Resolve paths lazily (not at module load) so tests can swap $HOME between
+ * cases without needing to bust the ESM module cache.
+ */
+function configDir(): string {
+  return join(homedir(), ".config", "axme-code");
+}
+
+export function authConfigPath(): string {
+  return join(configDir(), "auth.yaml");
+}
+
+export function loadAuthConfig(): AuthConfig | null {
+  const file = authConfigPath();
+  if (!pathExists(file)) return null;
+  const raw = readSafe(file);
+  if (!raw) return null;
+  try {
+    const parsed = yaml.load(raw) as Partial<AuthConfig> | null;
+    if (!parsed || typeof parsed !== "object") return null;
+    if (parsed.mode !== "subscription" && parsed.mode !== "api_key") return null;
+    const chosenAt = typeof parsed.chosenAt === "string" ? parsed.chosenAt : new Date().toISOString();
+    return { mode: parsed.mode, chosenAt };
+  } catch {
+    return null;
+  }
+}
+
+export function saveAuthConfig(mode: AuthMode): AuthConfig {
+  ensureDir(configDir());
+  const config: AuthConfig = { mode, chosenAt: new Date().toISOString() };
+  atomicWrite(authConfigPath(), yaml.dump(config));
+  return config;
+}
+
+/**
+ * Choose the sensible default when no saved choice exists and we can't ask
+ * the user. If an API key is set (regardless of subscription state) we keep
+ * the existing behavior: pass env through to Claude Code and let it decide.
+ * If only subscription is available, prefer it. If neither, return api_key
+ * so we fail the same way Claude Code would fail on its own.
+ */
+function heuristicMode(options: AuthOptions): AuthMode {
+  if (options.subscription.present && !options.apiKey.present) return "subscription";
+  return "api_key";
+}
+
+/**
+ * Effective auth mode for a scanner call. Reads the saved config if present,
+ * otherwise returns a heuristic without persisting anything.
+ */
+export function resolveAuthMode(): AuthMode {
+  const saved = loadAuthConfig();
+  if (saved) return saved.mode;
+  return heuristicMode(detectAuthOptions());
+}