Harden MCP schema and wire Claude OAuth/token handling

src/examples/claude.ts

@@ -1,7 +1,12 @@
 import "dotenv/config";
-import { query, type Options } from "@anthropic-ai/claude-agent-sdk";
+import { query, type Options, type SDKMessage } from "@anthropic-ai/claude-agent-sdk";
 import { pathToFileURL } from "node:url";
-import { getConfig } from "../config.js";
+import {
+  buildClaudeAuthEnv,
+  getConfig,
+  resolveAnthropicToken,
+  type AppConfig,
+} from "../config.js";
 import { createSessionContext } from "./session-context.js";
 
 function requiredPrompt(argv: string[]): string {
@@ -22,52 +27,96 @@ function buildOptions(config = getConfig()): Options {
   };
 }
 
-export async function runClaudePrompt(prompt: string): Promise<void> {
-  const config = getConfig();
-  const sessionContext = await createSessionContext("claude", {
+type ClaudeAuthOptionOverrides = {
+  apiKey?: string;
+  authToken?: string;
+};
+
+function buildClaudeAuthOptionOverrides(
+  config: Readonly<AppConfig>,
+): ClaudeAuthOptionOverrides {
+  if (config.provider.anthropicOauthToken) {
+    return { authToken: config.provider.anthropicOauthToken };
+  }
+
+  const token = resolveAnthropicToken(config.provider);
+  return token ? { apiKey: token } : {};
+}
+
+export async function readClaudeResult(
+  session: AsyncIterable<SDKMessage>,
+): Promise<string> {
+  let result = "";
+
+  for await (const message of session) {
+    if (message.type === "result" && message.subtype === "success") {
+      result = message.result.trim();
+    }
+
+    if (message.type === "result" && message.subtype !== "success") {
+      const detail = message.errors.join("; ");
+      throw new Error(
+        `Claude query failed (${message.subtype})${detail ? `: ${detail}` : ""}`,
+      );
+    }
+  }
+
+  if (!result) {
+    throw new Error("Claude run completed without a final result.");
+  }
+
+  return result;
+}
+
+type RunClaudePromptDependencies = {
+  config?: Readonly<AppConfig>;
+  createSessionContextFn?: typeof createSessionContext;
+  queryFn?: typeof query;
+  writeOutput?: (output: string) => void;
+};
+
+export async function runClaudePrompt(
+  prompt: string,
+  dependencies: RunClaudePromptDependencies = {},
+): Promise<void> {
+  const config = dependencies.config ?? getConfig();
+  const createSessionContextFn = dependencies.createSessionContextFn ?? createSessionContext;
+  const queryFn = dependencies.queryFn ?? query;
+  const writeOutput = dependencies.writeOutput ?? ((output: string) => console.log(output));
+  const sessionContext = await createSessionContextFn("claude", {
     prompt,
     config,
   });
 
   try {
+    const runtimeEnv = {
+      ...sessionContext.runtimeInjection.env,
+      ...buildClaudeAuthEnv(config.provider),
+    };
+    const authOptionOverrides = buildClaudeAuthOptionOverrides(config);
+
     const finalResponse = await sessionContext.runInSession(async () => {
-      const session = query({
+      const session = queryFn({
         prompt: sessionContext.promptWithContext,
         options: {
           ...buildOptions(config),
-          ...(sessionContext.mcp.claudeMcpServers ? { mcpServers: sessionContext.mcp.claudeMcpServers } : {}),
+          ...authOptionOverrides,
+          ...(sessionContext.mcp.claudeMcpServers
+            ? { mcpServers: sessionContext.mcp.claudeMcpServers as Options["mcpServers"] }
+            : {}),
           cwd: sessionContext.runtimeInjection.workingDirectory,
-          env: sessionContext.runtimeInjection.env,
+          env: runtimeEnv,
         },
       });
 
-      let result = "";
-
       try {
-        for await (const message of session) {
-          if (message.type === "result" && message.subtype === "success") {
-            result = message.result.trim();
-          }
-
-          if (message.type === "result" && message.subtype !== "success") {
-            const detail = message.errors.join("; ");
-            throw new Error(
-              `Claude query failed (${message.subtype})${detail ? `: ${detail}` : ""}`,
-            );
-          }
-        }
+        return await readClaudeResult(session);
       } finally {
         session.close();
       }
-
-      if (!result) {
-        throw new Error("Claude run completed without a final result.");
-      }
-
-      return result;
     });
 
-    console.log(finalResponse);
+    writeOutput(finalResponse);
   } finally {
     await sessionContext.close();
   }