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/services/interactive-session-service.ts b/apps/api/src/services/interactive-session-service.ts index 6bcb3c21..51080a52 100644 --- a/apps/api/src/services/interactive-session-service.ts +++ b/apps/api/src/services/interactive-session-service.ts @@ -43,11 +43,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/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 69ce5987..b877a1c8 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"); @@ -605,7 +649,6 @@ export function startTaskWorker() { ...(!isGitHubAppConfigured() ? ["GITHUB_TOKEN"] : []), ]), ]; - const taskUserId = task.createdBy ?? null; const resolvedSecrets = await resolveSecretsForTask( secretNames, task.repoUrl, @@ -650,63 +693,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 d998d084..27b9fe3e 100644 --- a/apps/api/src/ws/session-chat.ts +++ b/apps/api/src/ws/session-chat.ts @@ -7,11 +7,13 @@ import { listSessionChatEvents, } 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"; @@ -23,6 +25,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. @@ -106,15 +109,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 }; @@ -126,8 +138,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) { @@ -135,11 +159,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, @@ -211,9 +236,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. @@ -226,20 +259,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 { @@ -254,7 +295,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 }); persistChatEvent(sessionId, entry, log); @@ -292,7 +337,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 }); persistChatEvent(sessionId, entry, log); @@ -369,6 +417,7 @@ export async function sessionChatWs(app: FastifyInstance) { type: "status", status: isProcessing ? "thinking" : "idle", model: currentModel, + agentType, }); } break; @@ -390,50 +439,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 944ccfd6..11ddb758 100644 --- a/apps/web/src/components/session-chat.tsx +++ b/apps/web/src/components/session-chat.tsx @@ -1,36 +1,75 @@ "use client"; - -import { useCallback, useEffect, useRef, useState } from "react"; +import { useState, useEffect, useRef, useCallback, useMemo } from "react"; import { cn } from "@/lib/utils"; -import { Send, Square, Bot, Loader2 } from "lucide-react"; -import { LogViewer } from "@/components/log-viewer"; -import { useSessionLogs } from "@/hooks/use-session-logs"; +import { + Send, + Square, + Bot, + User, + FileText, + Terminal, + Code, + Search, + Globe, + ChevronDown, + ChevronRight, + AlertCircle, + Loader2, + Lightbulb, +} from "lucide-react"; +import { getWsBaseUrl } from "@/lib/ws-client.js"; +import { ANTHROPIC_CATALOG, GEMINI_CATALOG, resolveModelId } from "@optio/shared"; + +interface ChatEvent { + taskId: string; + timestamp: string; + sessionId?: string; + type: "text" | "tool_use" | "tool_result" | "thinking" | "system" | "error" | "info"; + content: string; + metadata?: Record; +} + +interface ChatMessage { + id: string; + role: "user" | "assistant"; + content: string; + timestamp: string; + events: ChatEvent[]; + costUsd?: number; +} + +type ChatStatus = "connecting" | "ready" | "thinking" | "idle" | "error" | "disconnected"; interface SessionChatProps { sessionId: string; onCostUpdate?: (costUsd: number) => void; - /** - * Lets the parent (e.g. terminal) call into the chat composer to inject - * text. The handler appends to the current draft and focuses the input. - */ onSendToAgent?: (handler: (text: string) => void) => void; + onModelUpdate?: ( + model: string, + agentType: string, + availableModels: { id: string; label: string }[], + ) => void; + onModelChange?: (handler: (model: string) => void) => void; } -/** - * Session chat — formerly a 600-line bespoke chat-bubble renderer; now a - * thin shell that: - * 1. owns the message composer (textarea + send button + interrupt), - * 2. owns the connection status / model picker bar, - * 3. delegates ALL message rendering to LogViewer via the - * composer / status / externalLogs / userMessages slots. - * - * Same widget Tasks / Jobs / Reviews / Agents now use, just sourced from - * the session WebSocket and dressed up with a chat composer. - */ -export function SessionChat({ sessionId, onCostUpdate, onSendToAgent }: SessionChatProps) { - const session = useSessionLogs(sessionId, { onCostUpdate }); +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); + + // WebSocket connection + const wsRef = useRef(null); const textareaRef = useRef(null); + const currentAssistantMsgRef = useRef(null); // Terminal can route highlighted text into our composer. const sendToAgent = useCallback((text: string) => { @@ -41,10 +80,122 @@ 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) => { + setModel(newModel); + wsRef.current?.send(JSON.stringify({ type: "set_model", model: newModel })); + }, []); + + useEffect(() => { + onModelChange?.(handleModelChange); + }, [handleModelChange, onModelChange]); + + // Compute model options based on agent type + const modelOptions = useMemo(() => { + const catalog = agentType === "gemini" ? GEMINI_CATALOG : ANTHROPIC_CATALOG; + return catalog.models.map((m) => ({ + id: m.id, + label: m.label, + latest: m.latest, + preview: m.preview, + })); + }, [agentType]); + + // Validate model when agentType changes - ensure model matches agent type + useEffect(() => { + const isValidModel = modelOptions.some((m) => m.id === model); + + if (!isValidModel) { + const defaultModel = agentType === "gemini" ? "gemini-2.5-flash" : "sonnet"; + 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`); + wsRef.current = ws; + + ws.onopen = () => setStatus("ready"); + ws.onmessage = (event) => { + let msg: any; + try { + msg = JSON.parse(event.data); + } catch { + return; + } + + switch (msg.type) { + case "status": + setStatus(msg.status as ChatStatus); + if (msg.model) setModel(msg.model); + if (msg.agentType) setAgentType(msg.agentType); + if (typeof msg.costUsd === "number") { + setCostUsd(msg.costUsd); + onCostUpdate?.(msg.costUsd); + } + break; + + case "chat_event": { + const chatEvent = msg.event as ChatEvent; + setMessages((prev) => { + const msgs = [...prev]; + let currentMsgId = currentAssistantMsgRef.current; + if (!currentMsgId || !msgs.find((m) => m.id === currentMsgId)) { + const newMsg: ChatMessage = { + id: `assistant-${Date.now()}`, + role: "assistant", + content: "", + timestamp: chatEvent.timestamp, + events: [], + }; + msgs.push(newMsg); + currentMsgId = newMsg.id; + currentAssistantMsgRef.current = currentMsgId; + } + const msgIdx = msgs.findIndex((m) => m.id === currentMsgId); + if (msgIdx >= 0) { + const updated = { ...msgs[msgIdx], events: [...msgs[msgIdx].events, chatEvent] }; + if (chatEvent.type === "text") { + updated.content = updated.events + .filter((e) => e.type === "text") + .map((e) => e.content) + .join(""); + } + msgs[msgIdx] = updated; + } + return msgs; + }); + break; + } + + case "cost_update": + setCostUsd(msg.costUsd); + onCostUpdate?.(msg.costUsd); + break; + } + }; + + ws.onclose = () => setStatus("disconnected"); + ws.onerror = () => setStatus("error"); + + return () => { + ws.close(); + }; + }, [sessionId, onCostUpdate]); + const handleSend = () => { const text = input.trim(); - if (!text || session.status === "thinking") return; - session.sendMessage(text); + if (!text || status === "thinking") return; + wsRef.current?.send(JSON.stringify({ type: "message", text })); setInput(""); requestAnimationFrame(() => { if (textareaRef.current) { @@ -53,6 +204,10 @@ export function SessionChat({ sessionId, onCostUpdate, onSendToAgent }: SessionC }); }; + const handleInterrupt = () => { + wsRef.current?.send(JSON.stringify({ type: "interrupt" })); + }; + const handleKeyDown = (e: React.KeyboardEvent) => { if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); @@ -60,110 +215,75 @@ export function SessionChat({ sessionId, onCostUpdate, onSendToAgent }: SessionC } }; - const disabled = session.status === "disconnected" || session.status === "error"; - return (
- - - Agent Chat - - Ask the agent to write code, fix bugs, or explore the repository. It operates in the - same worktree as your terminal. - -
- } - status={ -
-
- - {session.status} - {session.status === "thinking" ? ( - - ) : null} -
-
- {session.model} - {session.costUsd > 0 ? ( - - ${session.costUsd.toFixed(4)} - - ) : null} -
+
+ {messages.map((m) => ( +
+
{m.role}
+
{m.content}
- } - composer={ -
-