From 5c0e93b5f097d1e3f19061f3018d3f79e78ab929 Mon Sep 17 00:00:00 2001 From: Ramesh Nethi Date: Fri, 24 Apr 2026 02:50:08 +0530 Subject: [PATCH 1/2] fix: ensure sessions directory exists and improve GitHub token handling - Create /workspace/sessions directory in repo-init.sh for interactive session worktrees - Pass workspaceId to getGitHubToken in interactive-session-service - Add warning log when GitHub token unavailable instead of silent catch --- apps/api/src/services/interactive-session-service.ts | 7 ++++--- scripts/repo-init.sh | 3 ++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/apps/api/src/services/interactive-session-service.ts b/apps/api/src/services/interactive-session-service.ts index fc09fc9f..45ee1ab2 100644 --- a/apps/api/src/services/interactive-session-service.ts +++ b/apps/api/src/services/interactive-session-service.ts @@ -37,11 +37,12 @@ export async function createSession(input: { try { const { getGitHubToken } = await import("./github-token-service.js"); const ghToken = input.userId - ? await getGitHubToken({ userId: input.userId }) + ? await getGitHubToken({ userId: input.userId, workspaceId: input.workspaceId }) : await getGitHubToken({ server: true }); if (ghToken) env.GITHUB_TOKEN = ghToken; - } catch { - // No token, that's fine + } catch (err) { + logger.warn({ err, repoUrl, userId: input.userId }, "No GitHub token available for session"); + // Continue without token - public repos may still work } const imageConfig = repoConfig diff --git a/scripts/repo-init.sh b/scripts/repo-init.sh index c1ebaf8a..cdc7df80 100755 --- a/scripts/repo-init.sh +++ b/scripts/repo-init.sh @@ -103,8 +103,9 @@ echo "[optio] Cloning..." git clone --branch "${OPTIO_REPO_BRANCH}" --recurse-submodules "${OPTIO_REPO_URL}" repo 2>&1 echo "[optio] Repo cloned" -# Create tasks directory for worktrees +# Create directories for worktrees mkdir -p /workspace/tasks +mkdir -p /workspace/sessions # Run repo-level setup if present (.optio/setup.sh) if [ -f /workspace/repo/.optio/setup.sh ]; then From e95a089df5656ab171ddabbe6d20ebf1f2e6e759 Mon Sep 17 00:00:00 2001 From: Ramesh Nethi Date: Fri, 24 Apr 2026 21:45:09 +0530 Subject: [PATCH 2/2] feat(sessions): add multi-agent support and credential injection for interactive sessions Add comprehensive agent credential injection and multi-agent support for interactive sessions (terminal + chat), enabling Gemini, Claude Vertex AI, and all other agent types to work in session environments. ## Features **Agent Credential Service** (new) - Centralized credential retrieval for all agent types - Supports Claude (api-key, oauth-token, max-subscription, vertex-ai) - Supports Gemini, Codex, Copilot, Groq, OpenClaw, OpenCode - Handles service account keys as sensitive files (chmod 600) - 13 test cases covering all auth modes **Multi-Agent Session Support** - Terminal and chat now respect repo's defaultAgentType - Dynamic model switching mid-session - Agent-specific event parsers (Claude vs Gemini) - Integrated with existing header model dropdown **WebSocket Improvements** - Added keepalive ping (20s) to prevent idle disconnects - Filtered harmless warnings (stdin, yolo mode) - Agent stderr logging for debugging **Shared Utilities** - buildSetupFilesScript() utility for file injection ## Changes Backend (apps/api/src): - services/agent-credential-service.ts (NEW): centralized credential retrieval - services/agent-credential-service.test.ts (NEW): 13 test cases - utils/setup-files.ts (NEW): shared file injection utility - ws/session-terminal.ts: inject credentials based on repo agent type - ws/session-chat.ts: multi-agent command building, Gemini parser, keepalive - workers/task-worker.ts: refactored to use agent-credential-service Frontend (apps/web/src): - components/session-chat.tsx: model dropdown integration, agent type tracking - app/sessions/[id]/page.tsx: wired header dropdown to chat component ## Fixes - Terminal was hardcoded to claude-code, now uses repo's defaultAgentType - Chat was missing Vertex AI support, now supports all auth modes - Gemini errors weren't displaying (wrong parser) - 30s WebSocket timeout from lack of keepalive - Duplicate setup file injection code across 3 files ## Testing All new tests passing (13/13). Pre-existing failures unchanged (2). Co-Authored-By: Claude Sonnet 4.5 --- .../services/agent-credential-service.test.ts | 251 ++++++++++++++++ .../src/services/agent-credential-service.ts | 275 ++++++++++++++++++ apps/api/src/utils/setup-files.ts | 41 +++ apps/api/src/workers/task-worker.ts | 115 ++++---- apps/api/src/ws/session-chat.ts | 135 +++++---- apps/api/src/ws/session-terminal.ts | 50 +++- apps/web/src/app/sessions/[id]/page.tsx | 36 ++- apps/web/src/components/session-chat.tsx | 90 +++++- 8 files changed, 869 insertions(+), 124 deletions(-) create mode 100644 apps/api/src/services/agent-credential-service.test.ts create mode 100644 apps/api/src/services/agent-credential-service.ts create mode 100644 apps/api/src/utils/setup-files.ts diff --git a/apps/api/src/services/agent-credential-service.test.ts b/apps/api/src/services/agent-credential-service.test.ts new file mode 100644 index 00000000..d3bbf52f --- /dev/null +++ b/apps/api/src/services/agent-credential-service.test.ts @@ -0,0 +1,251 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { getAgentCredentials } from "./agent-credential-service.js"; + +// Mock dependencies +vi.mock("./secret-service.js", () => ({ + retrieveSecretWithFallback: vi.fn(), +})); + +vi.mock("./auth-service.js", () => ({ + getClaudeAuthToken: vi.fn(), +})); + +vi.mock("../logger.js", () => ({ + logger: { + child: vi.fn().mockReturnValue({ + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }), + }, +})); + +import { retrieveSecretWithFallback } from "./secret-service.js"; +import { getClaudeAuthToken } from "./auth-service.js"; + +describe("agent-credential-service", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe("getAgentCredentials - Claude Code", () => { + it("injects ANTHROPIC_API_KEY in api-key mode", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "CLAUDE_AUTH_MODE") return "api-key"; + if (name === "ANTHROPIC_API_KEY") return "sk-ant-test-key"; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("claude-code", "workspace-1", "user-1"); + + expect(result.env.ANTHROPIC_API_KEY).toBe("sk-ant-test-key"); + expect(result.setupFiles).toBeUndefined(); + }); + + it("injects CLAUDE_CODE_OAUTH_TOKEN in oauth-token mode", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "CLAUDE_AUTH_MODE") return "oauth-token"; + if (name === "CLAUDE_CODE_OAUTH_TOKEN") return "oauth-token-123"; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("claude-code", "workspace-1", "user-1"); + + expect(result.env.CLAUDE_CODE_OAUTH_TOKEN).toBe("oauth-token-123"); + expect(result.setupFiles).toBeUndefined(); + }); + + it("injects CLAUDE_CODE_OAUTH_TOKEN from host in max-subscription mode", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "CLAUDE_AUTH_MODE") return "max-subscription"; + throw new Error("Secret not found"); + }, + ); + + vi.mocked(getClaudeAuthToken).mockReturnValue({ + available: true, + token: "host-oauth-token", + }); + + const result = await getAgentCredentials("claude-code", "workspace-1", "user-1"); + + expect(result.env.CLAUDE_CODE_OAUTH_TOKEN).toBe("host-oauth-token"); + expect(result.setupFiles).toBeUndefined(); + }); + + it("configures Vertex AI with service account key", async () => { + const mockServiceAccountKey = JSON.stringify({ + type: "service_account", + project_id: "test-project", + private_key: "-----BEGIN PRIVATE KEY-----\ntest\n-----END PRIVATE KEY-----\n", + client_email: "test@test.iam.gserviceaccount.com", + }); + + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "CLAUDE_AUTH_MODE") return "vertex-ai"; + if (name === "CLAUDE_VERTEX_PROJECT_ID") return "test-project"; + if (name === "CLAUDE_VERTEX_REGION") return "us-central1"; + if (name === "CLAUDE_VERTEX_SERVICE_ACCOUNT_KEY") return mockServiceAccountKey; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("claude-code", "workspace-1", "user-1"); + + expect(result.env.ANTHROPIC_VERTEX_PROJECT_ID).toBe("test-project"); + expect(result.env.CLOUD_ML_REGION).toBe("us-central1"); + expect(result.env.CLAUDE_CODE_USE_VERTEX).toBe("1"); + expect(result.env.GOOGLE_APPLICATION_CREDENTIALS).toBe( + "/home/agent/.config/gcloud/gsa-key.json", + ); + expect(result.setupFiles).toHaveLength(1); + expect(result.setupFiles![0].path).toBe("/home/agent/.config/gcloud/gsa-key.json"); + expect(result.setupFiles![0].content).toBe(mockServiceAccountKey); + expect(result.setupFiles![0].sensitive).toBe(true); + }); + + it("uses workload identity when no service account key provided (Vertex AI)", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "CLAUDE_AUTH_MODE") return "vertex-ai"; + if (name === "CLAUDE_VERTEX_PROJECT_ID") return "test-project"; + if (name === "CLAUDE_VERTEX_REGION") return "us-central1"; + // No service account key + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("claude-code", "workspace-1", "user-1"); + + expect(result.env.ANTHROPIC_VERTEX_PROJECT_ID).toBe("test-project"); + expect(result.env.CLOUD_ML_REGION).toBe("us-central1"); + expect(result.env.CLAUDE_CODE_USE_VERTEX).toBe("1"); + expect(result.env.GOOGLE_APPLICATION_CREDENTIALS).toBeUndefined(); + expect(result.setupFiles).toBeUndefined(); + }); + }); + + describe("getAgentCredentials - Codex", () => { + it("injects OPENAI_API_KEY in api-key mode", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "CODEX_AUTH_MODE") return "api-key"; + if (name === "OPENAI_API_KEY") return "sk-openai-test"; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("codex", "workspace-1", "user-1"); + + expect(result.env.OPENAI_API_KEY).toBe("sk-openai-test"); + }); + + it("injects CODEX_APP_SERVER_URL in app-server mode", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "CODEX_AUTH_MODE") return "app-server"; + if (name === "CODEX_APP_SERVER_URL") return "https://codex-server.example.com"; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("codex", "workspace-1", "user-1"); + + expect(result.env.CODEX_APP_SERVER_URL).toBe("https://codex-server.example.com"); + }); + }); + + describe("getAgentCredentials - Gemini", () => { + it("injects GEMINI_API_KEY in api-key mode", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "GEMINI_AUTH_MODE") return "api-key"; + if (name === "GEMINI_API_KEY") return "gemini-key-123"; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("gemini", "workspace-1", "user-1"); + + expect(result.env.GEMINI_API_KEY).toBe("gemini-key-123"); + }); + + it("configures Vertex AI for Gemini", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "GEMINI_AUTH_MODE") return "vertex-ai"; + if (name === "GOOGLE_CLOUD_PROJECT") return "gemini-project"; + if (name === "GOOGLE_CLOUD_LOCATION") return "us-west1"; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("gemini", "workspace-1", "user-1"); + + expect(result.env.GOOGLE_CLOUD_PROJECT).toBe("gemini-project"); + expect(result.env.GOOGLE_CLOUD_LOCATION).toBe("us-west1"); + }); + }); + + describe("getAgentCredentials - Other agents", () => { + it("injects GITHUB_TOKEN for Copilot", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "GITHUB_TOKEN") return "ghp_test_token"; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("copilot", "workspace-1", "user-1"); + + expect(result.env.GITHUB_TOKEN).toBe("ghp_test_token"); + }); + + it("injects GROQ_API_KEY for Groq", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "GROQ_API_KEY") return "groq-key-123"; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("groq", "workspace-1", "user-1"); + + expect(result.env.GROQ_API_KEY).toBe("groq-key-123"); + }); + + it("injects OPENCLAW_API_KEY for OpenClaw", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "OPENCLAW_API_KEY") return "openclaw-key-123"; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("openclaw", "workspace-1", "user-1"); + + expect(result.env.OPENCLAW_API_KEY).toBe("openclaw-key-123"); + }); + + it("configures OpenCode defaults", async () => { + vi.mocked(retrieveSecretWithFallback).mockImplementation( + async (name: string): Promise => { + if (name === "OPENCODE_DEFAULT_BASE_URL") return "https://opencode.example.com"; + if (name === "OPENCODE_DEFAULT_MODEL") return "gpt-4"; + throw new Error("Secret not found"); + }, + ); + + const result = await getAgentCredentials("opencode", "workspace-1", "user-1"); + + expect(result.env.OPENCODE_DEFAULT_BASE_URL).toBe("https://opencode.example.com"); + expect(result.env.OPENCODE_DEFAULT_MODEL).toBe("gpt-4"); + }); + }); +}); diff --git a/apps/api/src/services/agent-credential-service.ts b/apps/api/src/services/agent-credential-service.ts new file mode 100644 index 00000000..de54fcc9 --- /dev/null +++ b/apps/api/src/services/agent-credential-service.ts @@ -0,0 +1,275 @@ +import { retrieveSecretWithFallback } from "./secret-service.js"; +import { logger } from "../logger.js"; + +// Agent types supported by Optio +export type AgentType = + | "claude-code" + | "codex" + | "copilot" + | "gemini" + | "groq" + | "opencode" + | "openclaw"; + +export interface AgentCredentials { + env: Record; + setupFiles?: Array<{ path: string; content: string; sensitive?: boolean }>; +} + +/** + * Retrieves agent credentials for interactive sessions or tasks. + * Handles all authentication modes: + * - Claude: api-key, oauth-token, max-subscription, vertex-ai + * - Codex: api-key, app-server + * - Gemini: api-key, vertex-ai + * - Other agents: Groq, OpenClaw, OpenCode, Copilot + * + * @param agentType - The agent type (claude-code, codex, gemini, copilot, groq, opencode, openclaw) + * @param workspaceId - Workspace ID for secret resolution + * @param userId - User ID for user-scoped secrets + * @returns Object with env vars and optional setupFiles (for service account keys) + */ +export async function getAgentCredentials( + agentType: AgentType, + workspaceId?: string | null, + userId?: string | null, +): Promise { + const log = logger.child({ agentType, workspaceId, userId }); + const env: Record = {}; + const setupFiles: Array<{ path: string; content: string; sensitive?: boolean }> = []; + + // ── Claude Code credentials ────────────────────────────────────── + if (agentType === "claude-code") { + const claudeAuthMode = + ((await retrieveSecretWithFallback("CLAUDE_AUTH_MODE", "global", workspaceId).catch( + () => null, + )) as any) ?? "api-key"; + + if (claudeAuthMode === "max-subscription") { + // Max subscription: fetch OAuth token from host keychain via auth proxy + const { getClaudeAuthToken } = await import("./auth-service.js"); + const authResult = getClaudeAuthToken(); + if (authResult.available && authResult.token) { + env.CLAUDE_CODE_OAUTH_TOKEN = authResult.token; + log.info("Injected CLAUDE_CODE_OAUTH_TOKEN from host credentials"); + } else { + log.warn({ error: authResult.error }, "Max subscription auth unavailable"); + } + } else if (claudeAuthMode === "oauth-token") { + // OAuth token mode: retrieve from secrets store + const oauthToken = await retrieveSecretWithFallback( + "CLAUDE_CODE_OAUTH_TOKEN", + "global", + workspaceId, + userId, + ).catch(() => null); + if (oauthToken) { + env.CLAUDE_CODE_OAUTH_TOKEN = oauthToken as string; + log.info("Injected CLAUDE_CODE_OAUTH_TOKEN from secrets store"); + } else { + log.warn("OAuth token mode selected but no CLAUDE_CODE_OAUTH_TOKEN found"); + } + } else if (claudeAuthMode === "vertex-ai") { + // Vertex AI mode: retrieve GCP project config and optional service account key + const projectId = await retrieveSecretWithFallback( + "CLAUDE_VERTEX_PROJECT_ID", + "global", + workspaceId, + ).catch(() => null); + const region = await retrieveSecretWithFallback( + "CLAUDE_VERTEX_REGION", + "global", + workspaceId, + ).catch(() => null); + const serviceAccountKey = await retrieveSecretWithFallback( + "CLAUDE_VERTEX_SERVICE_ACCOUNT_KEY", + "global", + workspaceId, + userId, + ).catch(() => null); + + if (projectId) env.ANTHROPIC_VERTEX_PROJECT_ID = projectId as string; + if (region) env.CLOUD_ML_REGION = region as string; + env.CLAUDE_CODE_USE_VERTEX = "1"; + + // If service account key provided, write it as a sensitive file + // Otherwise, fall back to workload identity + if (serviceAccountKey) { + setupFiles.push({ + path: "/home/agent/.config/gcloud/gsa-key.json", + content: serviceAccountKey as string, + sensitive: true, + }); + env.GOOGLE_APPLICATION_CREDENTIALS = "/home/agent/.config/gcloud/gsa-key.json"; + log.info("Injected CLAUDE_VERTEX_SERVICE_ACCOUNT_KEY as sensitive file"); + } else { + log.info("Using GKE workload identity for Vertex AI (no service account key provided)"); + } + } else if (claudeAuthMode === "api-key") { + // API key mode: retrieve ANTHROPIC_API_KEY + const apiKey = await retrieveSecretWithFallback( + "ANTHROPIC_API_KEY", + "global", + workspaceId, + userId, + ).catch(() => null); + if (apiKey) { + env.ANTHROPIC_API_KEY = apiKey as string; + log.info("Injected ANTHROPIC_API_KEY from secrets store"); + } else { + log.warn("API key mode selected but no ANTHROPIC_API_KEY found"); + } + } + } + + // ── Codex credentials ──────────────────────────────────────────── + if (agentType === "codex") { + const codexAuthMode = + ((await retrieveSecretWithFallback("CODEX_AUTH_MODE", "global", workspaceId).catch( + () => null, + )) as any) ?? "api-key"; + + if (codexAuthMode === "app-server") { + const appServerUrl = await retrieveSecretWithFallback( + "CODEX_APP_SERVER_URL", + "global", + workspaceId, + ).catch(() => null); + if (appServerUrl) { + env.CODEX_APP_SERVER_URL = appServerUrl as string; + log.info("Injected CODEX_APP_SERVER_URL from secrets store"); + } + } else { + // API key mode + const apiKey = await retrieveSecretWithFallback( + "OPENAI_API_KEY", + "global", + workspaceId, + userId, + ).catch(() => null); + if (apiKey) { + env.OPENAI_API_KEY = apiKey as string; + log.info("Injected OPENAI_API_KEY from secrets store"); + } else { + log.warn("Codex API key mode selected but no OPENAI_API_KEY found"); + } + } + } + + // ── Gemini credentials ─────────────────────────────────────────── + if (agentType === "gemini") { + const geminiAuthMode = + ((await retrieveSecretWithFallback("GEMINI_AUTH_MODE", "global", workspaceId).catch( + () => null, + )) as any) ?? "api-key"; + + if (geminiAuthMode === "vertex-ai") { + const projectId = await retrieveSecretWithFallback( + "GOOGLE_CLOUD_PROJECT", + "global", + workspaceId, + ).catch(() => null); + const location = await retrieveSecretWithFallback( + "GOOGLE_CLOUD_LOCATION", + "global", + workspaceId, + ).catch(() => null); + + if (projectId) env.GOOGLE_CLOUD_PROJECT = projectId as string; + if (location) env.GOOGLE_CLOUD_LOCATION = location as string; + log.info("Configured Gemini Vertex AI"); + } else { + // API key mode + const apiKey = await retrieveSecretWithFallback( + "GEMINI_API_KEY", + "global", + workspaceId, + userId, + ).catch(() => null); + if (apiKey) { + env.GEMINI_API_KEY = apiKey as string; + log.info("Injected GEMINI_API_KEY from secrets store"); + } else { + // Try legacy GOOGLE_GENAI_API_KEY + const legacyKey = await retrieveSecretWithFallback( + "GOOGLE_GENAI_API_KEY", + "global", + workspaceId, + userId, + ).catch(() => null); + if (legacyKey) { + env.GOOGLE_GENAI_API_KEY = legacyKey as string; + log.info("Injected GOOGLE_GENAI_API_KEY from secrets store"); + } else { + log.warn("Gemini API key mode selected but no GEMINI_API_KEY found"); + } + } + } + } + + // ── Copilot credentials ────────────────────────────────────────── + if (agentType === "copilot") { + const apiKey = await retrieveSecretWithFallback( + "GITHUB_TOKEN", + "global", + workspaceId, + userId, + ).catch(() => null); + if (apiKey) { + env.GITHUB_TOKEN = apiKey as string; + log.info("Injected GITHUB_TOKEN for Copilot"); + } + } + + // ── Groq credentials ───────────────────────────────────────────── + if (agentType === "groq") { + const apiKey = await retrieveSecretWithFallback( + "GROQ_API_KEY", + "global", + workspaceId, + userId, + ).catch(() => null); + if (apiKey) { + env.GROQ_API_KEY = apiKey as string; + log.info("Injected GROQ_API_KEY from secrets store"); + } else { + log.warn("No GROQ_API_KEY found"); + } + } + + // ── OpenClaw credentials ───────────────────────────────────────── + if (agentType === "openclaw") { + const apiKey = await retrieveSecretWithFallback( + "OPENCLAW_API_KEY", + "global", + workspaceId, + userId, + ).catch(() => null); + if (apiKey) { + env.OPENCLAW_API_KEY = apiKey as string; + log.info("Injected OPENCLAW_API_KEY from secrets store"); + } else { + log.warn("No OPENCLAW_API_KEY found"); + } + } + + // ── OpenCode credentials ───────────────────────────────────────── + if (agentType === "opencode") { + const baseUrl = await retrieveSecretWithFallback( + "OPENCODE_DEFAULT_BASE_URL", + "global", + workspaceId, + ).catch(() => null); + const model = await retrieveSecretWithFallback( + "OPENCODE_DEFAULT_MODEL", + "global", + workspaceId, + ).catch(() => null); + + if (baseUrl) env.OPENCODE_DEFAULT_BASE_URL = baseUrl as string; + if (model) env.OPENCODE_DEFAULT_MODEL = model as string; + log.info("Configured OpenCode defaults"); + } + + return { env, setupFiles: setupFiles.length > 0 ? setupFiles : undefined }; +} diff --git a/apps/api/src/utils/setup-files.ts b/apps/api/src/utils/setup-files.ts new file mode 100644 index 00000000..cb401206 --- /dev/null +++ b/apps/api/src/utils/setup-files.ts @@ -0,0 +1,41 @@ +/** + * Builds a bash script that writes files from a setupFiles array. + * Uses Python3 for safe JSON parsing and file operations. + * + * @param setupFiles Array of files to write + * @param message Optional message to echo before writing files + * @returns Bash script string, or empty string if no files + */ +export function buildSetupFilesScript( + setupFiles: Array<{ + path: string; + content: string; + executable?: boolean; + sensitive?: boolean; + }>, + message = "[optio] Writing setup files...", +): string { + if (!setupFiles || setupFiles.length === 0) { + return ""; + } + + const filesJson = Buffer.from(JSON.stringify(setupFiles)).toString("base64"); + + return [ + `echo "${message}"`, + `echo '${filesJson}' | base64 -d | python3 -c "`, + `import json, sys, os`, + `files = json.load(sys.stdin)`, + `for f in files:`, + ` p = f['path']`, + ` os.makedirs(os.path.dirname(p), exist_ok=True)`, + ` with open(p, 'w') as fh:`, + ` fh.write(f['content'])`, + ` if f.get('executable'):`, + ` os.chmod(p, 0o755)`, + ` elif f.get('sensitive'):`, + ` os.chmod(p, 0o600)`, + ` print(f' wrote {p}')`, + `"`, + ].join("\n"); +} diff --git a/apps/api/src/workers/task-worker.ts b/apps/api/src/workers/task-worker.ts index c7cbe12f..2a933e71 100644 --- a/apps/api/src/workers/task-worker.ts +++ b/apps/api/src/workers/task-worker.ts @@ -36,6 +36,7 @@ import { import { getPromptTemplate } from "../services/prompt-template-service.js"; import { isGitHubAppConfigured } from "../services/github-app-service.js"; import { getCredentialSecret } from "../services/credential-secret-service.js"; +import { getAgentCredentials } from "../services/agent-credential-service.js"; import { subscribeToTaskMessages } from "../services/task-message-bus.js"; import * as messageService from "../services/task-message-service.js"; import { detectAuthFailureInLogs, recordAuthEvent } from "../services/auth-failure-detector.js"; @@ -149,6 +150,7 @@ export function startTaskWorker() { // re-queue the task with a delay until off-peak starts. const { getRepoByUrl } = await import("../services/repo-service.js"); const taskWorkspaceId = currentTask.workspaceId ?? null; + const taskUserId = currentTask.createdBy ?? null; const repoConfig = await getRepoByUrl(currentTask.repoUrl, taskWorkspaceId); if (repoConfig?.offPeakOnly && !currentTask.ignoreOffPeak) { @@ -302,6 +304,40 @@ export function startTaskWorker() { ).catch(() => null)) as any) ?? undefined; const optioApiUrl = `http://${process.env.API_HOST ?? "host.docker.internal"}:${process.env.API_PORT ?? "4000"}`; + // Pre-flight validation for oauth-token mode + // Check cached validation to fail fast before pod provisioning + if (task.agentType === "claude-code" && claudeAuthMode === "oauth-token") { + try { + const { getCachedTokenValidation } = await import("./token-validation-worker.js"); + const cached = await getCachedTokenValidation(); + if (cached?.tokenExists && !cached.valid) { + throw new Error( + "Claude OAuth token is expired (detected by pre-flight validation). " + + "Go to Secrets to update CLAUDE_CODE_OAUTH_TOKEN, or re-run 'claude setup-token'.", + ); + } + } catch (preflight) { + // Re-throw if it's our own validation error; swallow infra errors + if (preflight instanceof Error && preflight.message.includes("pre-flight")) { + throw preflight; + } + } + } + + // Inject agent credentials using centralized service + const agentCredentials = await getAgentCredentials( + task.agentType as any, + taskWorkspaceId, + taskUserId, + ); + log.info( + { agentType: task.agentType, envVarCount: Object.keys(agentCredentials.env).length }, + "Injected agent credentials", + ); + + // Store credential setupFiles to merge after adapter config + const credentialSetupFiles = agentCredentials.setupFiles ?? []; + // Load and render prompt template const promptConfig = await getPromptTemplate(task.repoUrl); @@ -376,9 +412,17 @@ export function startTaskWorker() { maxTurnsReview: repoConfig?.maxTurnsReview ?? undefined, googleCloudProject, googleCloudLocation, - claudeVertexServiceAccountKey, + // Vertex service account key is now handled by agent-credential-service + claudeVertexServiceAccountKey: undefined, }); + // Merge credential setupFiles (e.g., Vertex AI service account keys) + if (credentialSetupFiles.length > 0) { + agentConfig.setupFiles = agentConfig.setupFiles ?? []; + agentConfig.setupFiles.push(...credentialSetupFiles); + log.info({ count: credentialSetupFiles.length }, "Merged credential setup files"); + } + // ── MCP servers & custom skills injection ──────────────────── const { getMcpServersForTask, buildMcpJsonContent } = await import("../services/mcp-server-service.js"); @@ -569,7 +613,6 @@ export function startTaskWorker() { ...(!isGitHubAppConfigured() ? ["GITHUB_TOKEN"] : []), ]), ]; - const taskUserId = task.createdBy ?? null; const resolvedSecrets = await resolveSecretsForTask( secretNames, task.repoUrl, @@ -614,63 +657,31 @@ export function startTaskWorker() { allEnv.OPTIO_RESTART_FROM_BRANCH = "true"; } - // Inject repo-level setup config into pod env - if (repoConfig?.extraPackages) { - allEnv.OPTIO_EXTRA_PACKAGES = repoConfig.extraPackages; - } - if (repoConfig?.setupCommands) { - allEnv.OPTIO_SETUP_COMMANDS = repoConfig.setupCommands; - } + // Inject agent credentials into task env + Object.assign(allEnv, agentCredentials.env); - // For max-subscription mode, fetch the OAuth token from the auth proxy - if (claudeAuthMode === "max-subscription") { - const { getClaudeAuthToken } = await import("../services/auth-service.js"); - const authResult = getClaudeAuthToken(); - if (authResult.available && authResult.token) { - allEnv.CLAUDE_CODE_OAUTH_TOKEN = authResult.token; - log.info("Injected CLAUDE_CODE_OAUTH_TOKEN from host credentials"); - } else { - throw new Error( - `Max subscription auth failed: ${authResult.error ?? "Token not available"}`, - ); - } - } - - // For oauth-token mode, read the token from the secrets store - if (claudeAuthMode === "oauth-token") { - // Pre-flight: check the cached validation from the background worker - // to fail fast before wasting ~10s on pod provisioning + worktree setup - try { - const { getCachedTokenValidation } = await import("./token-validation-worker.js"); - const cached = await getCachedTokenValidation(); - if (cached?.tokenExists && !cached.valid) { - throw new Error( - "Claude OAuth token is expired (detected by pre-flight validation). " + - "Go to Secrets to update CLAUDE_CODE_OAUTH_TOKEN, or re-run 'claude setup-token'.", - ); - } - } catch (preflight) { - // Re-throw if it's our own validation error; swallow infra errors - if (preflight instanceof Error && preflight.message.includes("pre-flight")) { - throw preflight; - } + // Strict validation: throw if required credentials are missing + if (task.agentType === "claude-code") { + if (claudeAuthMode === "max-subscription" && !allEnv.CLAUDE_CODE_OAUTH_TOKEN) { + throw new Error("Max subscription auth failed: Token not available"); } - - const oauthToken = await retrieveSecretWithFallback( - "CLAUDE_CODE_OAUTH_TOKEN", - "global", - taskWorkspaceId, - taskUserId, - ).catch(() => null); - if (oauthToken) { - allEnv.CLAUDE_CODE_OAUTH_TOKEN = oauthToken as string; - log.info("Injected CLAUDE_CODE_OAUTH_TOKEN from secrets store"); - } else { + if (claudeAuthMode === "oauth-token" && !allEnv.CLAUDE_CODE_OAUTH_TOKEN) { throw new Error( "OAuth token mode selected but no CLAUDE_CODE_OAUTH_TOKEN secret found. " + "Run `claude setup-token` and paste the token in the setup wizard.", ); } + if (claudeAuthMode === "api-key" && !allEnv.ANTHROPIC_API_KEY) { + log.warn("API key mode selected but no ANTHROPIC_API_KEY found"); + } + } + + // Inject repo-level setup config into pod env + if (repoConfig?.extraPackages) { + allEnv.OPTIO_EXTRA_PACKAGES = repoConfig.extraPackages; + } + if (repoConfig?.setupCommands) { + allEnv.OPTIO_SETUP_COMMANDS = repoConfig.setupCommands; } // Split env into pod-level (for repo-init.sh) and task-level (for exec). diff --git a/apps/api/src/ws/session-chat.ts b/apps/api/src/ws/session-chat.ts index 98913ed3..53f3ce74 100644 --- a/apps/api/src/ws/session-chat.ts +++ b/apps/api/src/ws/session-chat.ts @@ -3,11 +3,13 @@ import { z } from "zod"; import { getRuntime } from "../services/container-service.js"; import { getSession } from "../services/interactive-session-service.js"; import { getSettings } from "../services/optio-settings-service.js"; +import { getAgentCredentials } from "../services/agent-credential-service.js"; import { db } from "../db/client.js"; import { repoPods, repos, interactiveSessions } from "../db/schema.js"; import { eq } from "drizzle-orm"; import { logger } from "../logger.js"; import { parseClaudeEvent } from "../services/agent-event-parser.js"; +import { parseGeminiEvent } from "../services/gemini-event-parser.js"; import { publishSessionEvent } from "../services/event-bus.js"; import type { ExecSession, OptioSettings } from "@optio/shared"; import { authenticateWs, extractSessionToken } from "./ws-auth.js"; @@ -19,6 +21,7 @@ import { WS_CLOSE_CONNECTION_LIMIT, WS_CLOSE_MESSAGE_TOO_LARGE, } from "./ws-limits.js"; +import { buildSetupFilesScript } from "../utils/setup-files.js"; /** * Session chat WebSocket handler. @@ -102,15 +105,24 @@ export async function sessionChatWs(app: FastifyInstance) { return; } - // Get repo config for model defaults + // Get repo config for model defaults and agent type const [repoConfig] = await db.select().from(repos).where(eq(repos.repoUrl, session.repoUrl)); // Load Optio agent settings (model, system prompt, tool filtering, etc.) const workspaceId = req.user?.workspaceId ?? null; const optioSettings = await getSettings(workspaceId); - // Optio settings take precedence, then repo config, then default - let currentModel = optioSettings.model || repoConfig?.claudeModel || "sonnet"; + // Determine agent type from repo config (defaults to claude-code) + const agentType = (repoConfig?.defaultAgentType || "claude-code") as any; + + // Model selection depends on agent type - use repo config, not optio settings + // (optioSettings.model is for Optio chat interface, not interactive sessions) + let currentModel: string; + if (agentType === "gemini") { + currentModel = repoConfig?.geminiModel || "gemini-2.5-flash"; + } else { + currentModel = repoConfig?.claudeModel || "sonnet"; + } const rt = getRuntime(); const handle = { id: pod.podId ?? pod.podName, name: pod.podName }; @@ -122,8 +134,20 @@ export async function sessionChatWs(app: FastifyInstance) { let outputBuffer = ""; let promptCount = 0; - // Resolve auth env vars for the claude process - const authEnv = await buildAuthEnv(log, user.id); + // Resolve auth env vars and setup files for the agent + let authEnv: Record = {}; + let authSetupFiles: Array<{ path: string; content: string; sensitive?: boolean }> = []; + try { + const credentials = await getAgentCredentials(agentType, session.workspaceId, session.userId); + authEnv = credentials.env; + authSetupFiles = credentials.setupFiles ?? []; + log.info( + { agentType, envVarCount: Object.keys(authEnv).length }, + "Loaded agent credentials for session chat", + ); + } catch (err) { + log.warn({ err, agentType }, "Failed to retrieve agent credentials for session chat"); + } const send = (msg: Record) => { if (socket.readyState === 1) { @@ -131,11 +155,12 @@ export async function sessionChatWs(app: FastifyInstance) { } }; - // Send initial status with model info and settings + // Send initial status with model info, agent type, and settings send({ type: "status", status: "ready", model: currentModel, + agentType, costUsd: cumulativeCost, settings: { maxTurns: optioSettings.maxTurns, @@ -144,6 +169,13 @@ export async function sessionChatWs(app: FastifyInstance) { }, }); + // WebSocket keepalive: send ping every 20s to prevent idle timeout + const pingInterval = setInterval(() => { + if (socket.readyState === 1) { + socket.ping(); + } + }, 20000); + /** * Execute a single claude prompt in the pod worktree. * Uses `claude -p` in one-shot mode with stream-json output. @@ -174,9 +206,17 @@ export async function sessionChatWs(app: FastifyInstance) { fullPrompt = `${prompt}\n\n[Additional instructions: ${optioSettings.systemPrompt}]`; } - // Build the claude command + // Build the agent command const escapedPrompt = fullPrompt.replace(/'/g, "'\\''"); - const modelFlag = currentModel ? `--model ${currentModel}` : ""; + let agentCommand: string; + if (agentType === "gemini") { + const modelFlag = currentModel ? `-m ${currentModel}` : ""; + agentCommand = `gemini -p '${escapedPrompt}' ${modelFlag} --output-format stream-json --approval-mode yolo < /dev/null || true`; + } else { + // claude-code, codex, copilot, etc. + const modelFlag = currentModel ? `--model ${currentModel}` : ""; + agentCommand = `claude -p '${escapedPrompt}' ${modelFlag} --output-format stream-json --verbose --dangerously-skip-permissions < /dev/null || true`; + } // Build auth passthrough env vars so the agent can make // authenticated API calls on behalf of the requesting user. @@ -189,20 +229,28 @@ export async function sessionChatWs(app: FastifyInstance) { passthroughEnv.OPTIO_API_URL = apiUrl; } + // Build script to write auth setup files (e.g., Vertex AI service account keys) + const setupFilesScript = buildSetupFilesScript( + authSetupFiles, + "[optio] Writing agent credential files...", + ); + const script = [ "set -e", // Wait for repo to be ready "for i in $(seq 1 30); do [ -f /workspace/.ready ] && break; sleep 1; done", '[ -f /workspace/.ready ] || { echo "Repo not ready"; exit 1; }', + // Write auth setup files if present + ...(setupFilesScript ? [setupFilesScript] : []), `cd "${worktreePath}"`, - // Set auth env vars for the Claude process + // Set auth env vars for the agent process ...Object.entries(authEnv).map(([k, v]) => `export ${k}='${v.replace(/'/g, "'\\''")}'`), // Set auth passthrough env vars for Optio API calls ...Object.entries(passthroughEnv).map( ([k, v]) => `export ${k}='${v.replace(/'/g, "'\\''")}'`, ), - // Run claude in one-shot prompt mode with streaming JSON output - `claude -p '${escapedPrompt}' ${modelFlag} --output-format stream-json --verbose --dangerously-skip-permissions 2>&1 || true`, + // Run agent in one-shot prompt mode with streaming JSON output + agentCommand, ].join("\n"); try { @@ -217,7 +265,11 @@ export async function sessionChatWs(app: FastifyInstance) { for (const line of lines) { if (!line.trim()) continue; - const { entries } = parseClaudeEvent(line, sessionId); + // Use agent-specific parser + const { entries } = + agentType === "gemini" + ? parseGeminiEvent(line, sessionId) + : parseClaudeEvent(line, sessionId); for (const entry of entries) { send({ type: "chat_event", event: entry }); @@ -237,7 +289,13 @@ export async function sessionChatWs(app: FastifyInstance) { execSession.stderr.on("data", (chunk: Buffer) => { const text = chunk.toString("utf-8").trim(); - if (text) { + // Skip harmless warnings + if ( + text && + !text.includes("no stdin data received") && + !text.includes("approval mode: yolo") + ) { + log.warn({ agentType, text }, "Agent stderr output"); send({ type: "chat_event", event: { @@ -255,7 +313,10 @@ export async function sessionChatWs(app: FastifyInstance) { execSession!.stdout.on("end", () => { // Process any remaining buffer if (outputBuffer.trim()) { - const { entries } = parseClaudeEvent(outputBuffer, sessionId); + const { entries } = + agentType === "gemini" + ? parseGeminiEvent(outputBuffer, sessionId) + : parseClaudeEvent(outputBuffer, sessionId); for (const entry of entries) { send({ type: "chat_event", event: entry }); } @@ -322,6 +383,7 @@ export async function sessionChatWs(app: FastifyInstance) { type: "status", status: isProcessing ? "thinking" : "idle", model: currentModel, + agentType, }); } break; @@ -333,6 +395,7 @@ export async function sessionChatWs(app: FastifyInstance) { socket.on("close", () => { log.info("Session chat disconnected"); + clearInterval(pingInterval); releaseConnection(clientIp); if (execSession) { execSession.close(); @@ -343,50 +406,6 @@ export async function sessionChatWs(app: FastifyInstance) { } /** Build auth environment variables for the claude process in the pod. */ -async function buildAuthEnv( - log: { warn: (obj: any, msg: string) => void }, - userId?: string | null, -): Promise> { - const env: Record = {}; - - try { - const { retrieveSecret, retrieveSecretWithFallback } = - await import("../services/secret-service.js"); - const authMode = (await retrieveSecret("CLAUDE_AUTH_MODE").catch(() => null)) as string | null; - - if (authMode === "api-key") { - const apiKey = await retrieveSecretWithFallback( - "ANTHROPIC_API_KEY", - "global", - undefined, - userId, - ).catch(() => null); - if (apiKey) { - env.ANTHROPIC_API_KEY = apiKey as string; - } - } else if (authMode === "max-subscription") { - const { getClaudeAuthToken } = await import("../services/auth-service.js"); - const result = getClaudeAuthToken(); - if (result.available && result.token) { - env.CLAUDE_CODE_OAUTH_TOKEN = result.token; - } - } else if (authMode === "oauth-token") { - const token = await retrieveSecretWithFallback( - "CLAUDE_CODE_OAUTH_TOKEN", - "global", - undefined, - userId, - ).catch(() => null); - if (token) { - env.CLAUDE_CODE_OAUTH_TOKEN = token as string; - } - } - } catch (err) { - log.warn({ err }, "Failed to build auth env for session chat"); - } - - return env; -} /** Update the cumulative cost on the session record. */ async function updateSessionCost(sessionId: string, costUsd: number) { diff --git a/apps/api/src/ws/session-terminal.ts b/apps/api/src/ws/session-terminal.ts index 64a77601..a9b50524 100644 --- a/apps/api/src/ws/session-terminal.ts +++ b/apps/api/src/ws/session-terminal.ts @@ -2,8 +2,9 @@ import type { FastifyInstance } from "fastify"; import { z } from "zod"; import { getRuntime } from "../services/container-service.js"; import { getSession, addSessionPr } from "../services/interactive-session-service.js"; +import { getAgentCredentials } from "../services/agent-credential-service.js"; import { db } from "../db/client.js"; -import { repoPods } from "../db/schema.js"; +import { repoPods, repos } from "../db/schema.js"; import { eq } from "drizzle-orm"; import { logger } from "../logger.js"; import type { ContainerHandle, ExecSession } from "@optio/shared"; @@ -16,6 +17,7 @@ import { WS_CLOSE_CONNECTION_LIMIT, WS_CLOSE_MESSAGE_TOO_LARGE, } from "./ws-limits.js"; +import { buildSetupFilesScript } from "../utils/setup-files.js"; const PR_URL_REGEX = /https:\/\/github\.com\/[^/]+\/[^/]+\/pull\/(\d+)/g; @@ -86,11 +88,48 @@ export async function sessionTerminalWs(app: FastifyInstance) { const branch = session.branch; const repoUrl = session.repoUrl; + // Get repo config to determine agent type + const [repoConfig] = await db.select().from(repos).where(eq(repos.repoUrl, session.repoUrl)); + const agentType = (repoConfig?.defaultAgentType || "claude-code") as any; + + // Get agent credentials for the user based on repo's agent type + let credentials: { env: Record; setupFiles?: any[] } = { env: {} }; + try { + credentials = await getAgentCredentials(agentType, session.workspaceId, session.userId); + log.info( + { agentType, envVarCount: Object.keys(credentials.env).length }, + "Injected agent credentials", + ); + } catch (err) { + log.warn( + { err, agentType }, + "Failed to retrieve agent credentials, terminal will launch without them", + ); + } + + // Build env vars export statements + const envExports = Object.entries(credentials.env) + .map(([key, value]) => { + // Escape single quotes in the value + const escaped = value.replace(/'/g, "'\\''"); + return `export ${key}='${escaped}'`; + }) + .join("\n"); + + // Build script to write credential files if present + const setupFilesScript = buildSetupFilesScript( + credentials.setupFiles ?? [], + "[optio] Writing agent credential files...", + ); + const setupScript = [ "set -e", // Wait for repo to be ready "for i in $(seq 1 60); do [ -f /workspace/.ready ] && break; sleep 1; done", '[ -f /workspace/.ready ] || { echo "Repo not ready"; exit 1; }', + // Inject agent credentials (don't fail if credentials are missing) + ...(envExports ? [envExports] : []), + ...(setupFilesScript ? [setupFilesScript] : []), // Acquire repo lock for worktree setup "exec 9>/workspace/.repo-lock", "flock 9", @@ -130,6 +169,13 @@ export async function sessionTerminalWs(app: FastifyInstance) { try { execSession = await rt.exec(handle, ["bash", "-c", setupScript], { tty: true }); + // WebSocket keepalive: send ping every 20s to prevent idle timeout + const pingInterval = setInterval(() => { + if (socket.readyState === 1) { + socket.ping(); + } + }, 20000); + // Pipe exec stdout → WebSocket + scan for PR URLs execSession.stdout.on("data", (chunk: Buffer) => { if (socket.readyState === 1) { @@ -170,6 +216,7 @@ export async function sessionTerminalWs(app: FastifyInstance) { // Handle exec session end execSession.stdout.on("end", () => { + clearInterval(pingInterval); if (socket.readyState === 1) { socket.close(); } @@ -178,6 +225,7 @@ export async function sessionTerminalWs(app: FastifyInstance) { // Handle WebSocket close socket.on("close", () => { log.info("Session terminal disconnected"); + clearInterval(pingInterval); releaseConnection(clientIp); execSession?.close(); }); diff --git a/apps/web/src/app/sessions/[id]/page.tsx b/apps/web/src/app/sessions/[id]/page.tsx index 51862b2a..8c21179c 100644 --- a/apps/web/src/app/sessions/[id]/page.tsx +++ b/apps/web/src/app/sessions/[id]/page.tsx @@ -51,10 +51,13 @@ export default function SessionDetailPage({ params }: { params: Promise<{ id: st const [showEndWarning, setShowEndWarning] = useState(false); const [liveCost, setLiveCost] = useState(0); const [selectedModel, setSelectedModel] = useState(""); + const [agentType, setAgentType] = useState("claude-code"); + const [availableModels, setAvailableModels] = useState<{ id: string; label: string }[]>([]); const [showModelDropdown, setShowModelDropdown] = useState(false); - // Ref for "send to agent" handler + // Refs for agent chat handlers const sendToAgentRef = useRef<((text: string) => void) | null>(null); + const modelChangeRef = useRef<((model: string) => void) | null>(null); const fetchSession = async () => { try { @@ -115,6 +118,19 @@ export default function SessionDetailPage({ params }: { params: Promise<{ id: st sendToAgentRef.current = handler; }, []); + const handleModelUpdate = useCallback( + (model: string, agent: string, models: { id: string; label: string }[]) => { + setSelectedModel(model); + setAgentType(agent); + setAvailableModels(models); + }, + [], + ); + + const handleModelChangeRegister = useCallback((handler: (model: string) => void) => { + modelChangeRef.current = handler; + }, []); + if (loading) { return (
@@ -189,14 +205,14 @@ export default function SessionDetailPage({ params }: { params: Promise<{ id: st )} {/* Model selector */} - {isActive && modelConfig && ( + {isActive && selectedModel && availableModels.length > 0 && (
{showModelDropdown && ( @@ -205,20 +221,20 @@ export default function SessionDetailPage({ params }: { params: Promise<{ id: st className="fixed inset-0 z-40" onClick={() => setShowModelDropdown(false)} /> -
- {modelConfig.availableModels.map((m) => ( +
+ {availableModels.map((m) => ( ))}
@@ -294,6 +310,8 @@ export default function SessionDetailPage({ params }: { params: Promise<{ id: st sessionId={id} onCostUpdate={handleCostUpdate} onSendToAgent={handleSendToAgentRegister} + onModelUpdate={handleModelUpdate} + onModelChange={handleModelChangeRegister} /> } diff --git a/apps/web/src/components/session-chat.tsx b/apps/web/src/components/session-chat.tsx index 00c9fcc9..e50eb836 100644 --- a/apps/web/src/components/session-chat.tsx +++ b/apps/web/src/components/session-chat.tsx @@ -1,6 +1,6 @@ "use client"; -import { useState, useEffect, useRef, useCallback } from "react"; +import { useState, useEffect, useRef, useCallback, useMemo } from "react"; import { cn } from "@/lib/utils"; import { Send, @@ -19,6 +19,7 @@ import { Lightbulb, } from "lucide-react"; import { getWsBaseUrl } from "@/lib/ws-client.js"; +import { ANTHROPIC_CATALOG, GEMINI_CATALOG, resolveModelId } from "@optio/shared"; interface ChatEvent { taskId: string; @@ -44,13 +45,26 @@ interface SessionChatProps { sessionId: string; onCostUpdate?: (costUsd: number) => void; onSendToAgent?: (handler: (text: string) => void) => void; + onModelUpdate?: ( + model: string, + agentType: string, + availableModels: { id: string; label: string }[], + ) => void; + onModelChange?: (handler: (model: string) => void) => void; } -export function SessionChat({ sessionId, onCostUpdate, onSendToAgent }: SessionChatProps) { +export function SessionChat({ + sessionId, + onCostUpdate, + onSendToAgent, + onModelUpdate, + onModelChange, +}: SessionChatProps) { const [messages, setMessages] = useState([]); const [input, setInput] = useState(""); const [status, setStatus] = useState("connecting"); const [model, setModel] = useState("sonnet"); + const [agentType, setAgentType] = useState("claude-code"); const [costUsd, setCostUsd] = useState(0); const [expandedTools, setExpandedTools] = useState>(new Set()); @@ -73,6 +87,62 @@ export function SessionChat({ sessionId, onCostUpdate, onSendToAgent }: SessionC onSendToAgent?.(sendToAgent); }, [sendToAgent, onSendToAgent]); + // Expose a handler for external model changes (from header dropdown) + const handleModelChange = useCallback((newModel: string) => { + console.log("[SessionChat] External model change to:", newModel); + setModel(newModel); + wsRef.current?.send(JSON.stringify({ type: "set_model", model: newModel })); + }, []); + + useEffect(() => { + onModelChange?.(handleModelChange); + }, [handleModelChange, onModelChange]); + + // Debug: log agentType changes + useEffect(() => { + console.log("[SessionChat] agentType state changed to:", agentType); + }, [agentType]); + + // Compute model options based on agent type + const modelOptions = useMemo(() => { + const catalog = agentType === "gemini" ? GEMINI_CATALOG : ANTHROPIC_CATALOG; + const options = catalog.models.map((m) => ({ + id: m.id, + label: m.label, + latest: m.latest, + preview: m.preview, + })); + console.log( + "[SessionChat] Model options updated for agentType:", + agentType, + "count:", + options.length, + ); + return options; + }, [agentType]); + + // Validate model when agentType changes - ensure model matches agent type + useEffect(() => { + const isValidModel = modelOptions.some((m) => m.id === model); + + if (!isValidModel) { + // Model doesn't exist for this agent type, reset to default + const defaultModel = agentType === "gemini" ? "gemini-2.5-flash" : "sonnet"; + console.log( + `[SessionChat] Model "${model}" invalid for agent "${agentType}", resetting to "${defaultModel}"`, + ); + setModel(defaultModel); + wsRef.current?.send(JSON.stringify({ type: "set_model", model: defaultModel })); + } + }, [agentType, model, modelOptions]); + + // Notify parent when model/agentType/modelOptions change + useEffect(() => { + if (model && agentType && modelOptions.length > 0) { + onModelUpdate?.(model, agentType, modelOptions); + } + }, [model, agentType, modelOptions, onModelUpdate]); + // WebSocket connection useEffect(() => { const ws = new WebSocket(`${getWsBaseUrl()}/ws/sessions/${sessionId}/chat`); @@ -93,7 +163,19 @@ export function SessionChat({ sessionId, onCostUpdate, onSendToAgent }: SessionC switch (msg.type) { case "status": setStatus(msg.status as ChatStatus); - if (msg.model) setModel(msg.model); + if (msg.model) { + console.log("[SessionChat] Setting model to:", msg.model); + setModel(msg.model); + } + if (msg.agentType) { + console.log( + "[SessionChat] Setting agentType to:", + msg.agentType, + "current:", + agentType, + ); + setAgentType(msg.agentType); + } if (typeof msg.costUsd === "number") { setCostUsd(msg.costUsd); onCostUpdate?.(msg.costUsd); @@ -369,7 +451,7 @@ export function SessionChat({ sessionId, onCostUpdate, onSendToAgent }: SessionC ? "Agent is working... Press Esc or click Stop to interrupt" : "Enter to send, Shift+Enter for new line"} - {model} + {agentType}