diff --git a/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-branch-name.ts b/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-branch-name.ts index 7ad12334bab..65c1c7e5ca6 100644 --- a/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-branch-name.ts +++ b/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-branch-name.ts @@ -3,7 +3,7 @@ import { getSmallModel } from "@superset/chat/server/shared"; import { sanitizeBranchNameWithMaxLength } from "shared/utils/branch"; const BRANCH_NAME_INSTRUCTIONS = - "Generate a concise git branch name (2-4 words, kebab-case, descriptive). Return ONLY the branch name, nothing else."; + "Generate a concise git branch name (2-4 words, kebab-case, descriptive, 20 characters or less). Return ONLY the branch name, nothing else."; const MAX_CONFLICT_RESOLUTION_ATTEMPTS = 1000; const INITIAL_CONFLICT_SUFFIX = 2; @@ -58,7 +58,7 @@ export async function generateBranchNameFromPrompt( existingBranches: string[], branchPrefix?: string, ): Promise { - const model = getSmallModel(); + const model = await getSmallModel(); if (!model) return null; let generated: string | null; diff --git a/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-name.test.ts b/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-name.test.ts index 5fa3df6dff0..7926fa98075 100644 --- a/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-name.test.ts +++ b/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-name.test.ts @@ -1,7 +1,7 @@ import { afterAll, beforeEach, describe, expect, it, mock } from "bun:test"; const getSmallModelMock = mock( - (() => null) as (...args: unknown[]) => unknown | null, + (async () => null) as (...args: unknown[]) => Promise, ); const generateTitleFromMessageMock = mock( (async () => null) as (...args: unknown[]) => Promise, @@ -79,7 +79,7 @@ const { describe("generateWorkspaceNameFromPrompt", () => { beforeEach(() => { getSmallModelMock.mockClear(); - getSmallModelMock.mockReturnValue(null); + getSmallModelMock.mockResolvedValue(null); generateTitleFromMessageMock.mockClear(); generateTitleFromMessageMock.mockResolvedValue(null); selectGetMock.mockReset(); @@ -102,7 +102,7 @@ describe("generateWorkspaceNameFromPrompt", () => { }); it("returns the model-generated title when a model is available", async () => { - getSmallModelMock.mockReturnValueOnce({ id: "test-model" }); + getSmallModelMock.mockResolvedValueOnce({ id: "test-model" }); generateTitleFromMessageMock.mockResolvedValueOnce("Checking In"); await expect( @@ -116,13 +116,14 @@ describe("generateWorkspaceNameFromPrompt", () => { agentModel: { id: "test-model" }, agentId: "workspace-namer", agentName: "Workspace Namer", - instructions: "You generate concise workspace titles.", + instructions: + "You generate concise workspace titles. 20 characters or less. Return ONLY the title, nothing else.", tracingContext: { surface: "workspace-auto-name" }, }); }); it("preserves empty-string model results instead of forcing fallback", async () => { - getSmallModelMock.mockReturnValueOnce({ id: "test-model" }); + getSmallModelMock.mockResolvedValueOnce({ id: "test-model" }); generateTitleFromMessageMock.mockResolvedValueOnce(""); await expect( @@ -134,7 +135,7 @@ describe("generateWorkspaceNameFromPrompt", () => { }); it("falls back when generation throws", async () => { - getSmallModelMock.mockReturnValueOnce({ id: "test-model" }); + getSmallModelMock.mockResolvedValueOnce({ id: "test-model" }); generateTitleFromMessageMock.mockRejectedValueOnce(new Error("boom")); await expect( @@ -155,7 +156,7 @@ afterAll(() => { describe("attemptWorkspaceAutoRenameFromPrompt", () => { beforeEach(() => { getSmallModelMock.mockClear(); - getSmallModelMock.mockReturnValue(null); + getSmallModelMock.mockResolvedValue(null); generateTitleFromMessageMock.mockClear(); generateTitleFromMessageMock.mockResolvedValue(null); selectGetMock.mockReset(); @@ -196,7 +197,7 @@ describe("attemptWorkspaceAutoRenameFromPrompt", () => { isUnnamed: true, deletingAt: null, }); - getSmallModelMock.mockReturnValueOnce({ id: "test-model" }); + getSmallModelMock.mockResolvedValueOnce({ id: "test-model" }); generateTitleFromMessageMock.mockResolvedValueOnce(""); await expect( diff --git a/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-name.ts b/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-name.ts index 1bb06606be5..0df5e43de28 100644 --- a/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-name.ts +++ b/apps/desktop/src/lib/trpc/routers/workspaces/utils/ai-name.ts @@ -34,7 +34,7 @@ export async function generateWorkspaceNameFromPrompt(prompt: string): Promise<{ usedPromptFallback: boolean; warning?: string; }> { - const model = getSmallModel(); + const model = await getSmallModel(); if (model) { try { const generated = await generateTitleFromMessage({ @@ -42,7 +42,8 @@ export async function generateWorkspaceNameFromPrompt(prompt: string): Promise<{ agentModel: model, agentId: "workspace-namer", agentName: "Workspace Namer", - instructions: "You generate concise workspace titles.", + instructions: + "You generate concise workspace titles. 20 characters or less. Return ONLY the title, nothing else.", tracingContext: { surface: "workspace-auto-name" }, }); if (generated !== null && generated !== undefined) { diff --git a/packages/chat/src/server/desktop/auth/provider-ids.ts b/packages/chat/src/server/desktop/auth/provider-ids.ts index 4c2e3736cac..dca854f3d2e 100644 --- a/packages/chat/src/server/desktop/auth/provider-ids.ts +++ b/packages/chat/src/server/desktop/auth/provider-ids.ts @@ -1,6 +1,5 @@ -export const ANTHROPIC_AUTH_PROVIDER_ID = "anthropic"; -export const OPENAI_AUTH_PROVIDER_ID = "openai-codex"; -export const OPENAI_AUTH_PROVIDER_IDS = [ +export { + ANTHROPIC_AUTH_PROVIDER_ID, OPENAI_AUTH_PROVIDER_ID, - "openai", -] as const; + OPENAI_AUTH_PROVIDER_IDS, +} from "../../shared/auth-provider-ids"; diff --git a/packages/chat/src/server/shared/auth-provider-ids.ts b/packages/chat/src/server/shared/auth-provider-ids.ts new file mode 100644 index 00000000000..373408bd24e --- /dev/null +++ b/packages/chat/src/server/shared/auth-provider-ids.ts @@ -0,0 +1,8 @@ +export const ANTHROPIC_AUTH_PROVIDER_ID = "anthropic"; +export const OPENAI_AUTH_PROVIDER_ID = "openai-codex"; +// Mastracode historically wrote OpenAI under "openai" before the Codex split. +// Read both when resolving credentials. +export const OPENAI_AUTH_PROVIDER_IDS = [ + OPENAI_AUTH_PROVIDER_ID, + "openai", +] as const; diff --git a/packages/chat/src/server/shared/small-model/get-small-model.test.ts b/packages/chat/src/server/shared/small-model/get-small-model.test.ts new file mode 100644 index 00000000000..f6c278fa64e --- /dev/null +++ b/packages/chat/src/server/shared/small-model/get-small-model.test.ts @@ -0,0 +1,58 @@ +import { describe, expect, it } from "bun:test"; +import { isAnthropicApiKey, isOpenAIApiKey } from "./get-small-model"; + +describe("isAnthropicApiKey", () => { + it("accepts a real-shaped key", () => { + expect( + isAnthropicApiKey( + "sk-ant-api03-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + ), + ).toBe(true); + }); + + it("rejects dev placeholders", () => { + expect(isAnthropicApiKey("dummy")).toBe(false); + expect(isAnthropicApiKey("placeholder")).toBe(false); + expect(isAnthropicApiKey("")).toBe(false); + }); + + it("rejects OAuth access tokens (sk-ant-oat…) sent as api keys", () => { + // OAuth tokens fail when sent via x-api-key. Filter them so we fall + // through to the OAuth path which sends them via Authorization Bearer. + expect( + isAnthropicApiKey("sk-ant-oat-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), + ).toBe(false); + }); + + it("rejects keys with the prefix but absurd lengths", () => { + expect(isAnthropicApiKey("sk-ant-api")).toBe(false); + }); + + it("rejects unrelated provider keys", () => { + expect(isAnthropicApiKey("sk-proj-foo")).toBe(false); + }); +}); + +describe("isOpenAIApiKey", () => { + it("accepts legacy, project, and service-account key shapes", () => { + expect( + isOpenAIApiKey("sk-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), + ).toBe(true); + expect( + isOpenAIApiKey("sk-proj-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), + ).toBe(true); + expect( + isOpenAIApiKey("sk-svcacct-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), + ).toBe(true); + }); + + it("rejects dev placeholders and obviously-fake values", () => { + expect(isOpenAIApiKey("dummy")).toBe(false); + expect(isOpenAIApiKey("sk-")).toBe(false); + expect(isOpenAIApiKey("")).toBe(false); + }); + + it("rejects values without the sk- prefix", () => { + expect(isOpenAIApiKey("api-key-aaaaaaaaaaaaaaaaaaaaaaaaaaaaa")).toBe(false); + }); +}); diff --git a/packages/chat/src/server/shared/small-model/get-small-model.ts b/packages/chat/src/server/shared/small-model/get-small-model.ts index df489d56270..58cb4837fa1 100644 --- a/packages/chat/src/server/shared/small-model/get-small-model.ts +++ b/packages/chat/src/server/shared/small-model/get-small-model.ts @@ -1,101 +1,177 @@ -import { existsSync, readFileSync } from "node:fs"; -import { homedir, platform } from "node:os"; -import { join } from "node:path"; import { createAnthropic } from "@ai-sdk/anthropic"; import { createOpenAI } from "@ai-sdk/openai"; +import { createAuthStorage } from "mastracode"; +import { + ANTHROPIC_AUTH_PROVIDER_ID, + OPENAI_AUTH_PROVIDER_IDS, +} from "../auth-provider-ids"; const ANTHROPIC_SMALL_MODEL_ID = "claude-haiku-4-5-20251001"; const OPENAI_SMALL_MODEL_ID = "gpt-4o-mini"; +const MIN_API_KEY_LENGTH = 30; + +// OAuth tokens issued through the Claude Code flow are accepted by the +// Anthropic API only when these companion headers are sent alongside the +// `Authorization: Bearer` header. Mastracode hands us the token; we own +// the wiring into createAnthropic and the request-time headers. +const ANTHROPIC_OAUTH_HEADERS = { + "anthropic-beta": "claude-code-20250219,oauth-2025-04-20", + "user-agent": "claude-cli/2.1.2 (external, cli)", + "x-app": "cli", +} as const; + +type AuthStorage = ReturnType; + +let cachedAuthStorage: AuthStorage | null = null; + +function getAuthStorage(): AuthStorage { + if (!cachedAuthStorage) { + cachedAuthStorage = createAuthStorage(); + } + cachedAuthStorage.reload(); + return cachedAuthStorage; +} + /** - * Resolves the mastracode auth.json path (same logic as mastracode's - * `getAppDataDir`). We read it directly to avoid importing mastracode, - * which eagerly loads @mastra/fastembed → onnxruntime-node (208 MB native - * binary) and breaks electron-vite bundling. + * Anthropic API keys are issued in the form `sk-ant-api…` (currently + * `sk-ant-api03-…`). Reject anything else — most importantly OAuth access + * tokens (`sk-ant-oat…`), which Anthropic rejects when sent as `x-api-key`, + * and dev placeholders like `dummy`. */ -function getAuthJsonPath(): string { - const p = platform(); - let base: string; - if (p === "darwin") { - base = join(homedir(), "Library", "Application Support"); - } else if (p === "win32") { - base = process.env.APPDATA ?? join(homedir(), "AppData", "Roaming"); - } else { - base = process.env.XDG_DATA_HOME ?? join(homedir(), ".local", "share"); - } - return join(base, "mastracode", "auth.json"); +export function isAnthropicApiKey(key: string): boolean { + return key.startsWith("sk-ant-api") && key.length >= MIN_API_KEY_LENGTH; } -type AuthData = Record; +/** + * OpenAI keys all start with `sk-` (legacy `sk-…`, project `sk-proj-…`, + * service-account `sk-svcacct-…`). The length floor catches placeholders. + */ +export function isOpenAIApiKey(key: string): boolean { + return key.startsWith("sk-") && key.length >= MIN_API_KEY_LENGTH; +} -function readAuthData(): AuthData | null { - const path = getAuthJsonPath(); - if (!existsSync(path)) return null; - try { - return JSON.parse(readFileSync(path, "utf-8")) as AuthData; - } catch { - return null; - } +function isObjectRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; } -function getStoredApiKey( - authData: AuthData | null, - providerId: string, -): string | null { - if (!authData) return null; - const entry = authData[`apikey:${providerId}`]; - if ( - typeof entry === "object" && - entry !== null && - "type" in entry && - entry.type === "api_key" && - "key" in entry && - typeof entry.key === "string" && - entry.key.trim().length > 0 - ) { - return entry.key.trim(); +type AnthropicResolved = + | { kind: "apiKey"; key: string } + | { kind: "oauth"; accessToken: string }; + +async function resolveAnthropic(): Promise { + const env = process.env.ANTHROPIC_API_KEY?.trim(); + if (env && isAnthropicApiKey(env)) { + return { kind: "apiKey", key: env }; + } + + try { + const authStorage = getAuthStorage(); + + // Settings-saved API keys are stored at `apikey:`. Prefer + // these over whatever sits in the main slot — otherwise an OAuth + // login (which writes to the main slot) would mask a stored API key + // the user explicitly added. + const storedApiKey = authStorage + .getStoredApiKey(ANTHROPIC_AUTH_PROVIDER_ID) + ?.trim(); + if (storedApiKey && isAnthropicApiKey(storedApiKey)) { + return { kind: "apiKey", key: storedApiKey }; + } + + const credential = authStorage.get(ANTHROPIC_AUTH_PROVIDER_ID); + if (!isObjectRecord(credential)) return null; + + if ( + credential.type === "api_key" && + typeof credential.key === "string" && + isAnthropicApiKey(credential.key.trim()) + ) { + return { kind: "apiKey", key: credential.key.trim() }; + } + + if (credential.type === "oauth") { + // Mastracode's getApiKey returns a fresh access token, refreshing + // via the Claude Code OAuth flow when expired and persisting the + // new credential back to auth.json. This replaces the custom + // refresh dance we used to maintain in this package. + const accessToken = await authStorage.getApiKey( + ANTHROPIC_AUTH_PROVIDER_ID, + ); + if (typeof accessToken === "string" && accessToken.trim().length > 0) { + return { kind: "oauth", accessToken: accessToken.trim() }; + } + } + } catch (error) { + console.warn("[get-small-model] anthropic auth resolution failed:", error); } + return null; } -function resolveApiKey( - envVar: string | undefined, - authData: AuthData | null, - providerId: string, -): string | null { - const env = envVar?.trim(); - if (env) return env; - return getStoredApiKey(authData, providerId); +async function resolveOpenAIApiKey(): Promise { + const env = process.env.OPENAI_API_KEY?.trim(); + if (env && isOpenAIApiKey(env)) return env; + + try { + const authStorage = getAuthStorage(); + for (const providerId of OPENAI_AUTH_PROVIDER_IDS) { + // Same precedence reasoning as Anthropic: dedicated apikey: slot + // before the main slot. + const stored = authStorage.getStoredApiKey(providerId)?.trim(); + if (stored && isOpenAIApiKey(stored)) return stored; + + const credential = authStorage.get(providerId); + if ( + isObjectRecord(credential) && + credential.type === "api_key" && + typeof credential.key === "string" && + isOpenAIApiKey(credential.key.trim()) + ) { + return credential.key.trim(); + } + } + } catch (error) { + console.warn("[get-small-model] openai auth resolution failed:", error); + } + + return null; } /** * Returns an AI-SDK `LanguageModel` for small-model tasks (branch naming, - * title generation). Tries Anthropic first, falls back to OpenAI. Returns - * `null` if no credentials are available. + * title generation). Returns `null` if no usable credentials are available. * - * Reads credentials from env vars and mastracode's auth.json directly - * (API keys only). OAuth-only users fall back to `null`. + * Resolution order: + * 1. ANTHROPIC_API_KEY env var (validated) + * 2. mastracode auth storage — Anthropic api key + * 3. mastracode auth storage — Anthropic OAuth (refreshed on the fly) + * 4. OPENAI_API_KEY env var (validated) + * 5. mastracode auth storage — OpenAI api key (`openai-codex` / `openai`) + * + * API keys are validated by prefix + minimum length so dev placeholders + * (e.g. `ANTHROPIC_API_KEY=dummy` from a sample .env) fall through to the + * next path instead of being sent to the API and failing 401. */ -export function getSmallModel(): unknown | null { - const authData = readAuthData(); - - const anthropicKey = resolveApiKey( - process.env.ANTHROPIC_API_KEY, - authData, - "anthropic", - ); - if (anthropicKey) { - return createAnthropic({ apiKey: anthropicKey })(ANTHROPIC_SMALL_MODEL_ID); +export async function getSmallModel(): Promise { + const anthropic = await resolveAnthropic(); + if (anthropic?.kind === "apiKey") { + return createAnthropic({ apiKey: anthropic.key })(ANTHROPIC_SMALL_MODEL_ID); + } + if (anthropic?.kind === "oauth") { + return createAnthropic({ + authToken: anthropic.accessToken, + headers: ANTHROPIC_OAUTH_HEADERS, + })(ANTHROPIC_SMALL_MODEL_ID); } - const openaiKey = resolveApiKey( - process.env.OPENAI_API_KEY, - authData, - "openai", - ); + const openaiKey = await resolveOpenAIApiKey(); if (openaiKey) { return createOpenAI({ apiKey: openaiKey }).chat(OPENAI_SMALL_MODEL_ID); } + console.warn( + "[get-small-model] no credentials found — naming will fall back", + ); return null; } diff --git a/packages/host-service/src/trpc/router/workspace-creation/utils/ai-branch-name.ts b/packages/host-service/src/trpc/router/workspace-creation/utils/ai-branch-name.ts index b2bbec615af..7a236e8d037 100644 --- a/packages/host-service/src/trpc/router/workspace-creation/utils/ai-branch-name.ts +++ b/packages/host-service/src/trpc/router/workspace-creation/utils/ai-branch-name.ts @@ -3,7 +3,7 @@ import { getSmallModel } from "@superset/chat/server/shared"; import { deduplicateBranchName } from "./sanitize-branch"; const BRANCH_NAME_INSTRUCTIONS = - "Generate a concise git branch name (2-4 words, kebab-case, descriptive). Return ONLY the branch name, nothing else."; + "Generate a concise git branch name (2-4 words, kebab-case, descriptive, 20 characters or less). Return ONLY the branch name, nothing else."; const MAX_BRANCH_LENGTH = 100; @@ -30,7 +30,7 @@ export async function generateBranchNameFromPrompt( prompt: string, existingBranches: string[], ): Promise { - const model = getSmallModel(); + const model = await getSmallModel(); if (!model) return null; let generated: string | null;