diff --git a/.github/workflows/opencode.yml b/.github/workflows/opencode.yml index 8d87261e..842bd730 100644 --- a/.github/workflows/opencode.yml +++ b/.github/workflows/opencode.yml @@ -27,41 +27,156 @@ jobs: with: persist-credentials: false - - name: Probe primary model availability + - name: Probe model availability id: probe shell: bash env: + COMMENT: ${{ github.event.comment.body }} ALIBABA_TOKEN_PLAN_API_KEY: ${{ secrets.ALIBABA_TOKEN_PLAN_API_KEY }} DASHSCOPE_API_KEY: ${{ secrets.DASHSCOPE_API_KEY }} OPENCODE_API_KEY: ${{ secrets.OPENCODE_API_KEY }} run: | + probe_zen() { + local id="$1" + local url="https://opencode.ai/zen/v1/chat/completions" + local body="{\"model\":\"${id}\",\"messages\":[{\"role\":\"user\",\"content\":\"ping\"}],\"max_tokens\":1}" + local auth="Authorization: Bearer ${OPENCODE_API_KEY}" + local extra=() + if [[ "${id}" == "gpt-5.6-luna" ]]; then + url="https://opencode.ai/zen/v1/responses" + body="{\"model\":\"${id}\",\"input\":[{\"role\":\"user\",\"content\":[{\"type\":\"input_text\",\"text\":\"ping\"}]}],\"max_output_tokens\":1}" + fi + curl -sf --max-time 15 -X POST "${url}" \ + -H "${auth}" "${extra[@]}" -H "Content-Type: application/json" \ + -d "${body}" >/dev/null 2>&1 + } + MODEL="opencode/big-pickle" + VARIANT="" + if [[ "${COMMENT}" =~ (^|[[:space:]])/(oc|opencode)[[:space:]]+review ]]; then + echo "::notice::/oc review detected - using review model (opencode/gpt-5.6-luna, max reasoning)" + MODEL="opencode/gpt-5.6-luna" + VARIANT="max" + if probe_zen gpt-5.6-luna; then + echo "model=${MODEL}" >> "$GITHUB_OUTPUT" + echo "variant=${VARIANT}" >> "$GITHUB_OUTPUT" + exit 0 + fi + echo "::warning::Review model (opencode/gpt-5.6-luna) unavailable - falling back to big-pickle" + MODEL="opencode/big-pickle" + VARIANT="" + fi + + if probe_zen big-pickle; then + echo "model=${MODEL}" >> "$GITHUB_OUTPUT" + echo "variant=${VARIANT}" >> "$GITHUB_OUTPUT" + exit 0 + fi + echo "::warning::Primary model (${MODEL}) unavailable - falling back to alibaba-token-plan" + MODEL="alibaba-token-plan/qwen3.8-max" + VARIANT="" if ! curl -sf --max-time 15 \ - -X POST "https://opencode.ai/zen/v1/chat/completions" \ - -H "Authorization: Bearer ${OPENCODE_API_KEY}" \ + -X POST "https://token-plan.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1/chat/completions" \ + -H "Authorization: Bearer ${ALIBABA_TOKEN_PLAN_API_KEY}" \ -H "Content-Type: application/json" \ - -d '{"model":"big-pickle","messages":[{"role":"user","content":"ping"}],"max_tokens":1}' \ + -d '{"model":"qwen3.8-max","messages":[{"role":"user","content":"ping"}],"max_tokens":1}' \ >/dev/null 2>&1; then - echo "::warning::Primary model (opencode/big-pickle) unavailable - falling back to alibaba-token-plan" - MODEL="alibaba-token-plan/qwen3.8-max" - if ! curl -sf --max-time 15 \ - -X POST "https://token-plan.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1/chat/completions" \ - -H "Authorization: Bearer ${ALIBABA_TOKEN_PLAN_API_KEY}" \ - -H "Content-Type: application/json" \ - -d '{"model":"qwen3.8-max","messages":[{"role":"user","content":"ping"}],"max_tokens":1}' \ - >/dev/null 2>&1; then - echo "::warning::alibaba-token-plan unavailable - falling back to DashScope" - MODEL="alibaba/qwen3.8-max" - fi + echo "::warning::alibaba-token-plan unavailable - falling back to DashScope" + MODEL="alibaba/qwen3.8-max" fi echo "model=${MODEL}" >> "$GITHUB_OUTPUT" + echo "variant=${VARIANT}" >> "$GITHUB_OUTPUT" + + # Replicates anomalyco/opencode/github@v1.18.18 (composite action) inline so + # the run step can retry. The composite action cannot wrap itself: classic + # Actions has no step retry and the CLI commits/pushes from this checkout, + # so on a rejected push (remote advanced mid-run) the work would be lost. + - name: Get opencode version + id: version + shell: bash + run: | + VERSION=$(curl -sf https://api.github.com/repos/anomalyco/opencode/releases/latest | grep -o '"tag_name": *"[^"]*"' | cut -d'"' -f4) + echo "version=${VERSION:-latest}" >> "$GITHUB_OUTPUT" + + - name: Cache opencode + id: cache + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: ~/.opencode/bin + key: opencode-${{ runner.os }}-${{ runner.arch }}-${{ steps.version.outputs.version }} + + - name: Install opencode + if: steps.cache.outputs.cache-hit != 'true' + shell: bash + run: curl -fsSL https://opencode.ai/install | bash - name: Run opencode - uses: anomalyco/opencode/github@14b37df39168eaf6a6faf862ec4a7bbe9c825bbd # v1.18.18 + id: run_opencode + shell: bash env: ALIBABA_TOKEN_PLAN_API_KEY: ${{ secrets.ALIBABA_TOKEN_PLAN_API_KEY }} DASHSCOPE_API_KEY: ${{ secrets.DASHSCOPE_API_KEY }} OPENCODE_API_KEY: ${{ secrets.OPENCODE_API_KEY }} GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - with: - model: ${{ steps.probe.outputs.model }} + MODEL: ${{ steps.probe.outputs.model }} + VARIANT: ${{ steps.probe.outputs.variant }} + run: | + set -u + echo "$HOME/.opencode/bin" >> "$GITHUB_PATH" + export PATH="$HOME/.opencode/bin:$PATH" + + run_opencode() { + opencode github run + } + + if run_opencode; then + echo "::notice::opencode github run succeeded (attempt 1)" + exit 0 + fi + echo "::warning::opencode github run failed on attempt 1" + + BRANCH="$(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo '')" + + # Salvage: the CLI committed locally before a rejected push (remote + # advanced mid-run). Rebase those commits onto the updated remote and + # push. This matters because review threads may already be resolved, + # so a plain re-run would find nothing to do and silently lose work. + if [[ -n "${BRANCH}" && "${BRANCH}" != "HEAD" ]]; then + echo "==> salvaging agent commits onto updated remote (branch: ${BRANCH})" + git config --local http.https://github.com/.extraheader \ + "AUTHORIZATION: basic $(printf "x-access-token:${GH_TOKEN}" | base64 -w 0)" 2>/dev/null || true + git fetch --prune origin 2>/dev/null || true + if git rev-parse --verify "origin/${BRANCH}" >/dev/null 2>&1; then + if [[ "$(git rev-list --count "origin/${BRANCH}"..HEAD 2>/dev/null || echo 0)" -gt 0 ]]; then + if git rebase "origin/${BRANCH}"; then + if git push 2>/dev/null; then + echo "::notice::salvaged agent work: rebased onto origin/${BRANCH} and pushed" + exit 0 + fi + echo "::warning::salvage push failed; falling back to synced re-run" + fi + git rebase --abort 2>/dev/null || true + else + echo "::notice::no local-only commits; failure was not a rejected push" + fi + fi + fi + + # Fallback: sync this checkout's branch to the remote (discarding the + # failed run's local state) and run a fresh session against the + # current base, so a subsequent push is a fast-forward. + echo "==> syncing branch to remote and retrying once" + if [[ -n "${BRANCH}" && "${BRANCH}" != "HEAD" ]]; then + git fetch --prune origin 2>/dev/null || true + if git rev-parse --verify "origin/${BRANCH}" >/dev/null 2>&1; then + git checkout -B "${BRANCH}" "origin/${BRANCH}" >/dev/null 2>&1 || true + git reset --hard "origin/${BRANCH}" >/dev/null 2>&1 || true + fi + fi + + if run_opencode; then + echo "::notice::opencode github run succeeded on retry" + exit 0 + fi + echo "::error::opencode github run failed on both attempts" + exit 1 diff --git a/.opencode/github-commands.md b/.opencode/github-commands.md index 9abf102d..c21bbaa4 100644 --- a/.opencode/github-commands.md +++ b/.opencode/github-commands.md @@ -5,6 +5,25 @@ GitHub Action and begins with `/oc` (or `/opencode`). The message usually carrie `` context block (title, body, changed files, comments, reviews) — read it carefully before answering. +## Model routing + +The workflow probes and selects the agent model per command (see `.github/workflows/opencode.yml`): + +| Command | Primary model | Fallback chain | +| ------------ | -------------------------------------- | ---------------------------------------------------------------------- | +| `/oc review` | `opencode/gpt-5.6-luna` (`variant: max`) | `opencode/big-pickle` → `alibaba-token-plan/qwen3.8-max` → `alibaba/qwen3.8-max` | +| `/oc fix` | `opencode/big-pickle` | `alibaba-token-plan/qwen3.8-max` → `alibaba/qwen3.8-max` | + +Each model is probed with a minimal request before the run; a disabled or unavailable model +falls through to the next in the chain. `/oc review` runs are short and judgment-heavy, so +the cost-efficient `gpt-5.6-luna` runs with `max` reasoning effort to maximize finding +quality while keeping per-run cost in the tens of cents; `/oc fix` runs are long agentic +edit loops, where the free big-pickle keeps cost at $0. Note that `big-pickle` advertises no +reasoning-effort variants, so `variant: max` is only applied when `gpt-5.6-luna` is actually +selected — the probe clears it on any fallback. Review runs send code snippets to an +OpenAI-hosted model — acceptable for public repos; keep in mind OpenAI may retain requests +for evaluation purposes. + ## `/oc review` When a user message is exactly `/oc review` or begins with `/oc review`, treat it as a @@ -17,24 +36,60 @@ request to review the current pull request. Extra text after the shortcut, e.g. 1. Identify the actionable findings. An actionable finding is one where you can point at a concrete problem in the code and, when feasible, propose a specific change. -2. Post each actionable finding as its **own comment** on the PR via the `gh` CLI - (preinstalled in GitHub Actions; the `GITHUB_TOKEN` env var is available, no login needed): - - ```bash - gh api repos/{owner}/{repo}/issues/{pr_number}/comments -F body=@finding.md - ``` +2. Post each actionable finding as its **own resolvable review thread** via the `gh` CLI + (preinstalled in GitHub Actions; the `GITHUB_TOKEN` env var is available, no login + needed). Fall back down this ladder until the finding is posted: + + a. **Inline line comment** (preferred) — pins the finding to a line in the PR diff and + creates a resolvable thread. Use the PR head SHA (`Head: { Sha: ... }` in the + `` context) as `commit_id`, plus the file and line the finding is + about: + + ```bash + gh api repos/{owner}/{repo}/pulls/{pr_number}/comments \ + -f body=@finding.md \ + -f path="src/example.ts" \ + -F line=42 \ + -f commit_id="$HEAD_SHA" + ``` + + For a finding spanning a line range, add `-F start_line=` (and, for a + deletion, `-f start_side=LEFT`). + + b. **File-level comment** — if the line is not part of the diff (the call above returns a + 422), retry against the file without a line number: + + ```bash + gh api repos/{owner}/{repo}/pulls/{pr_number}/comments \ + -f body=@finding.md \ + -f path="src/example.ts" \ + -f subject_type=file + ``` + + c. **Issue comment** (last resort) — if the file is not in the PR diff either, post to + the timeline (not a resolvable thread) and flag it in the "Out of diff" section of + the summary: + + ```bash + gh api repos/{owner}/{repo}/issues/{pr_number}/comments -F body=@finding.md + ``` Derive `owner`/`repo` from `baseRepository.nameWithOwner` in the `` - context (split on `/`), and `pr_number` from `Number:`. Write the finding body to a - temp file (`finding.md`) rather than passing a giant `-f body=` string, so multiline - Markdown and code blocks survive intact. Post comments one at a time and keep a list of - the posted comment IDs/URLs. If a `gh` call fails, do not stop the review — fall back to - including that finding in the final summary comment instead. + context (split on `/`), `pr_number` from `Number:`, and `HEAD_SHA` from + `Head: { Sha: ... }`. Write the finding body to a temp file (`finding.md`) rather than + passing a giant `-f body=` string, so multiline Markdown and code blocks survive intact. + Post threads one at a time — this endpoint is secondary-rate-limited if you post too + fast — and keep a list of the posted comment IDs/URLs and of which findings fell back to + an issue comment. If a `gh` call fails at every level, do not stop the review — record + the finding in the "Out of diff" section of the summary instead. 3. **Your final reply text** (what the action posts as the single reply comment) must be a - **short summary index**: overall assessment, plus one line per posted finding with its - file:line, severity, and a link to that finding's comment (`gh api .../issues/{n}/comments` - responses include the `html_url`). Keep it tight — the detail lives in the per-finding - comments. + **short summary index**: overall assessment; one line per threaded finding with its + file:line, severity, and a link to that finding's comment (both endpoint responses + include the `html_url`); and an **"Out of diff"** section listing every finding that + could not be posted as a review thread — fallback issue comments and any finding with no + diff location (e.g. missing tests, missing docs, cross-file concerns) — with its + severity, the file name(s) and line(s) it covers, and the issue found. Keep the rest + tight — the detail lives in the per-finding comments. 4. Group low-severity nits and non-actionable observations into the final summary comment instead of posting more comments. @@ -45,9 +100,20 @@ You are reviewing, not editing: - **Do NOT modify any files and do NOT leave the working tree dirty.** The action auto-commits and pushes any uncommitted changes to the PR branch — that is not wanted here. - Include a **committable suggestion** in each finding comment when it is feasible to write - one for that specific finding: a concrete diff (lines with `+`/`-`) or exact replacement - snippet the author can apply. If a finding does not have a cut-and-dried fix, say so and - describe the change needed instead of inventing code. + one for that specific finding. Wrap the exact replacement in a GitHub `suggestion` + fenced block so GitHub renders a one-click **Commit suggestion** button right in the + comment: + + ```` + ```suggestion + + ``` + ```` + + Use one contiguous block per finding, matching the existing lines it replaces; GitHub + applies it to the file on commit. If a finding does not have a cut-and-dried fix — no + contiguous single-file replacement — say so and describe the change needed instead of + inventing code. ### Finding comment format @@ -56,7 +122,8 @@ Each finding comment should contain: 1. **Severity** — `high` / `medium` / `low` (or `critical`). 2. **Location** — `file:line` (or a line range). 3. **Problem** — why it is wrong, grounded in the actual code. -4. **Suggested fix** — a committable diff or snippet when feasible. +4. **Suggested fix** — a GitHub `suggestion` fenced block (see "Committing behavior") + when the fix is a contiguous replacement, otherwise a description of the change needed. ### Review scope diff --git a/server.ts b/server.ts index 9c4d2830..82222caa 100644 --- a/server.ts +++ b/server.ts @@ -7,6 +7,8 @@ import { ANY_DOT_VENV_DIR } from "./src/server/shared/anyDotVenvDir.ts"; import { loadStudioEnv } from "./src/server/loadStudioEnv.ts"; import { shutdownMcpClient } from "./src/server/services/mcp/client.ts"; +import { shutdownSidecar } from "./src/server/services/genai/venv.ts"; +import { createSingleFlightShutdown } from "./src/server/gracefulShutdown.ts"; import { mountSystemRoutes, type SystemProbeOptions } from "./src/server/routes/system.ts"; import { mountGithubRoutes } from "./src/server/routes/github.ts"; import { mountAiRoutes } from "./src/server/routes/ai/index.ts"; @@ -32,24 +34,26 @@ const app = express(); app.use(express.json({ limit: "10mb", strict: false })); // Security headers with a CSP that allows the app's own assets + CDN for ORT WASM. -app.use(helmet({ - contentSecurityPolicy: { - directives: { - defaultSrc: ["'self'"], - scriptSrc: ["'self'", "'unsafe-inline'", "https://esm.sh"], - styleSrc: ["'self'", "'unsafe-inline'"], - imgSrc: ["'self'", "data:", "blob:"], - connectSrc: ["'self'", "https://cdn.jsdelivr.net", "https://huggingface.co", "https://*.hf.co"], - workerSrc: ["'self'", "blob:"], - childSrc: ["'self'", "blob:"], - fontSrc: ["'self'"], - objectSrc: ["'none'"], - baseUri: ["'self'"], - formAction: ["'self'"], - frameAncestors: ["'self'"], +app.use( + helmet({ + contentSecurityPolicy: { + directives: { + defaultSrc: ["'self'"], + scriptSrc: ["'self'", "'unsafe-inline'", "https://esm.sh"], + styleSrc: ["'self'", "'unsafe-inline'"], + imgSrc: ["'self'", "data:", "blob:"], + connectSrc: ["'self'", "https://cdn.jsdelivr.net", "https://huggingface.co", "https://*.hf.co"], + workerSrc: ["'self'", "blob:"], + childSrc: ["'self'", "blob:"], + fontSrc: ["'self'"], + objectSrc: ["'none'"], + baseUri: ["'self'"], + formAction: ["'self'"], + frameAncestors: ["'self'"], + }, }, - }, -})); + }), +); // Prevent indexing if accidentally exposed to the public internet app.use((_req, res, next) => { @@ -184,13 +188,7 @@ async function startServer() { server: { middlewareMode: true, watch: { - ignored: [ - "**/.venv/**", - ANY_DOT_VENV_DIR, - "**/node_modules/**", - "**/models/**", - "**/.cache/**", - ], + ignored: ["**/.venv/**", ANY_DOT_VENV_DIR, "**/node_modules/**", "**/models/**", "**/.cache/**"], }, }, appType: "spa", @@ -265,7 +263,7 @@ async function startServer() { // end of startServer() so errors from Vite, static middleware, and all // application middleware are handled here. app.use((err: unknown, _req: Request, res: Response, _next: NextFunction) => { - console.error("[express] unhandled error:", err instanceof Error ? err.stack ?? err.message : err); + console.error("[express] unhandled error:", err instanceof Error ? (err.stack ?? err.message) : err); if (res.headersSent) return; res.status(500).json({ error: "Internal server error" }); }); @@ -305,16 +303,20 @@ async function startServer() { }); } +/** + * Single-flight shutdown runner. A repeated SIGINT/SIGTERM while cleanup is + * running returns the in-flight promise instead of starting a second cleanup + * (which could exit before the first sidecar/MCP shutdown completes); only the + * shared flow calls process.exit. + */ +const gracefulShutdown = createSingleFlightShutdown([shutdownSidecar, shutdownMcpClient]); + process.on("SIGINT", () => { - // eslint-disable-next-line no-console -- intentional shutdown logging - console.log("\n[SIGINT] Shutting down."); - void shutdownMcpClient().finally(() => process.exit(0)); + void gracefulShutdown("SIGINT"); }); process.on("SIGTERM", () => { - // eslint-disable-next-line no-console -- intentional shutdown logging - console.log("\n[SIGTERM] Shutting down."); - void shutdownMcpClient().finally(() => process.exit(0)); + void gracefulShutdown("SIGTERM"); }); process.on("exit", () => { diff --git a/src/components/features/assistant/GenaiEnginePanel.tsx b/src/components/features/assistant/GenaiEnginePanel.tsx new file mode 100644 index 00000000..8ff8e507 --- /dev/null +++ b/src/components/features/assistant/GenaiEnginePanel.tsx @@ -0,0 +1,154 @@ +import { Download, HardDriveDownload, PackageCheck, RefreshCw } from "lucide-react"; +import { useEffect, useState } from "react"; +import { cn } from "@/lib/utils"; + +interface GenaiModelStatus { + ready: boolean; + filesPresent: number; + filesRequired: number; + localSizeBytes: number; +} + +interface GenaiStatus { + venvReady: boolean; + model: GenaiModelStatus; +} + +function StatusBadge({ ok, label }: { ok: boolean; label: string }) { + return ( + + {label} + + ); +} + +/** + * Engine setup controls for the built-in GenAI provider: installs the Python + * runtime and downloads the model the provider needs before first use. + */ +export function GenaiEnginePanel() { + const [status, setStatus] = useState(null); + const [busy, setBusy] = useState<"setup" | "download" | null>(null); + const [error, setError] = useState(""); + + // Fetch resolves asynchronously, so the state update happens in the fetch + // callback (external system sync), not in the effect body itself. + const refresh = (signal?: AbortSignal) => { + void fetch("/api/ai/genai/status", { signal }) + .then(async (r) => { + if (!r.ok || signal?.aborted) return; + setStatus((await r.json()) as GenaiStatus); + }) + .catch(() => { + // Server unreachable; the panel keeps its last known state. + }); + }; + + useEffect(() => { + const controller = new AbortController(); + refresh(controller.signal); + return () => controller.abort(); + }, []); + + const run = async (action: "setup" | "download", url: string) => { + setBusy(action); + setError(""); + try { + const r = await fetch(url, { method: "POST" }); + const data = (await r.json().catch(() => ({}))) as { error?: string }; + if (!r.ok) throw new Error(data.error || `HTTP ${r.status}`); + refresh(); + } catch (err: unknown) { + setError(err instanceof Error ? err.message : String(err)); + } finally { + setBusy(null); + } + }; + + const venvReady = status?.venvReady ?? false; + const model = status?.model; + const modelReady = model?.ready ?? false; + + return ( +
+
+

Local engine setup

+ +
+ +
+
+ + Runtime engine + +
+ +
+ +
+
+ + Model files + 0 + ? `${model.filesPresent}/${model.filesRequired} files` + : "Not downloaded" + } + /> +
+ +
+ +

+ Install the engine and download the model before activating this provider. Both run once and are + cached locally. +

+ + {error &&

{error}

} +
+ ); +} diff --git a/src/components/features/assistant/ManualProviderSetup.tsx b/src/components/features/assistant/ManualProviderSetup.tsx index 2e9f398d..b3d85e42 100644 --- a/src/components/features/assistant/ManualProviderSetup.tsx +++ b/src/components/features/assistant/ManualProviderSetup.tsx @@ -5,6 +5,7 @@ import { cn } from "@/lib/utils"; import { CATEGORY_LABELS, PROVIDER_OPTIONS, type ProviderId } from "./aiProviderCatalog"; import { CodexAccountPanel } from "./CodexAccountPanel"; import { DevinAccountPanel } from "./DevinAccountPanel"; +import { GenaiEnginePanel } from "./GenaiEnginePanel"; import { ModelCombobox } from "./ModelCombobox"; import type { AiProviderSettings } from "./useAiProviderSettings"; @@ -18,6 +19,9 @@ interface ProvidersProp { * @param providers - Provider settings and selection handler used to control the dropdown */ function ProviderSelect({ providers }: ProvidersProp) { + // Category headers render once even when a category is non-contiguous + // (e.g. "custom" holds both the built-in engine and openai-compat). + const emittedCategories = new Set(); return (
)} -
- - providers.setSettingsApiKey(e.target.value)} - onBlur={() => providers.refreshModelsForTypedApiKey()} - onKeyDown={(e) => e.key === "Enter" && void providers.saveProvider()} - className="w-full bg-slate-900 border border-slate-700 rounded-lg px-3 py-2 text-sm text-slate-200 focus:outline-none focus:border-electric-blue" - /> -
+ {settingsProvider === "bedrock" && ( +
+ + providers.setSettingsBaseUrl(e.target.value)} + onKeyDown={(e) => e.key === "Enter" && void providers.saveProvider()} + className="w-full bg-slate-900 border border-slate-700 rounded-lg px-3 py-2 text-sm text-slate-200 focus:outline-none focus:border-electric-blue" + /> +

+ Optional — defaults to AWS_REGION, then us-east-1. Keys use the format accessKeyId:secretAccessKey + (add :sessionToken for assumed-role credentials), or leave blank for the default AWS credential + chain. +

+
+ )} + + {settingsProvider !== "genai" && ( +
+ + providers.setSettingsApiKey(e.target.value)} + onBlur={() => providers.refreshModelsForTypedApiKey()} + onKeyDown={(e) => e.key === "Enter" && void providers.saveProvider()} + className="w-full bg-slate-900 border border-slate-700 rounded-lg px-3 py-2 text-sm text-slate-200 focus:outline-none focus:border-electric-blue" + /> +
+ )} + + {settingsProvider === "genai" && } {settingsProvider === "cloudflare" && (
diff --git a/src/components/features/assistant/aiProviderCatalog.ts b/src/components/features/assistant/aiProviderCatalog.ts index aff32430..60951fc2 100644 --- a/src/components/features/assistant/aiProviderCatalog.ts +++ b/src/components/features/assistant/aiProviderCatalog.ts @@ -64,6 +64,16 @@ export const PROVIDER_OPTIONS: readonly ProviderOption[] = [ baseUrl: "https://api.x.ai/v1", category: "direct", }, + { + id: "bedrock", + name: "AWS Bedrock", + models: ["anthropic.claude-3-5-haiku-20241022-v1:0"], + keyEnvVar: "AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY", + docsUrl: "aws.amazon.com/bedrock", + category: "direct", + description: + "Converse API. Paste accessKeyId:secretAccessKey (optionally :sessionToken), or leave blank for the default AWS credential chain.", + }, // ── API Routers & Aggregators ──────────────────────────────────────── { id: "openrouter", diff --git a/src/components/features/assistant/useAiProviderSettings.ts b/src/components/features/assistant/useAiProviderSettings.ts index 0f22462e..cec63f06 100644 --- a/src/components/features/assistant/useAiProviderSettings.ts +++ b/src/components/features/assistant/useAiProviderSettings.ts @@ -41,6 +41,8 @@ function validateApiKeyProviderForm(input: { input.settingsProvider === "openai-compat" || // Bedrock can activate on the default AWS credential chain alone. input.settingsProvider === "bedrock" || + // Built-in GenAI runs a local engine and never uses an API key. + input.settingsProvider === "genai" || isLocalAllowEmptyKey(input.resolvedBaseUrl); if (input.settingsProvider === "cloudflare") { // Env-only activation requires both manual fields empty. A partial paste is rejected. @@ -155,8 +157,7 @@ export function useAiProviderSettings({ */ const userModelOverrideRef = useRef(false); - const providerOption = - PROVIDER_OPTIONS.find((p) => p.id === settingsProvider) ?? PROVIDER_OPTIONS[0]!; + const providerOption = PROVIDER_OPTIONS.find((p) => p.id === settingsProvider) ?? PROVIDER_OPTIONS[0]!; const isCompatMode = settingsProvider === "openai-compat" || !!providerOption.baseUrl; const isStaleRefresh = (sequence: number) => sequence !== refreshSequenceRef.current; diff --git a/src/server/__tests__/gracefulShutdown.test.ts b/src/server/__tests__/gracefulShutdown.test.ts new file mode 100644 index 00000000..172654aa --- /dev/null +++ b/src/server/__tests__/gracefulShutdown.test.ts @@ -0,0 +1,58 @@ +import { describe, it, expect, vi, afterEach } from "vitest"; +import { createSingleFlightShutdown } from "../gracefulShutdown.ts"; + +describe("createSingleFlightShutdown", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("runs each cleanup once, then exits", async () => { + const sidecar = vi.fn(async () => {}); + const mcp = vi.fn(async () => {}); + const exit = vi.spyOn(process, "exit").mockImplementation(() => undefined as never); + + const shutdown = createSingleFlightShutdown([sidecar, mcp]); + await shutdown("SIGINT"); + + expect(sidecar).toHaveBeenCalledTimes(1); + expect(mcp).toHaveBeenCalledTimes(1); + expect(exit).toHaveBeenCalledWith(0); + }); + + it("is single-flight: a second signal reuses the in-flight shutdown", async () => { + const deferred = () => { + let resolve!: (value: T) => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; + }; + + const sidecarGate = deferred(); + const mcpGate = deferred(); + const sidecar = vi.fn(() => sidecarGate.promise); + const mcp = vi.fn(() => mcpGate.promise); + const exit = vi.spyOn(process, "exit").mockImplementation(() => undefined as never); + + const shutdown = createSingleFlightShutdown([sidecar, mcp]); + + const first = shutdown("SIGINT"); + const second = shutdown("SIGTERM"); + + // The repeated signal must not start a second cleanup pass. + expect(sidecar).toHaveBeenCalledTimes(1); + expect(mcp).toHaveBeenCalledTimes(1); + expect(exit).not.toHaveBeenCalled(); + + // Let cleanup finish, then assert a single process.exit(0). + sidecarGate.resolve(); + mcpGate.resolve(); + await first; + await second; + + expect(sidecar).toHaveBeenCalledTimes(1); + expect(mcp).toHaveBeenCalledTimes(1); + expect(exit).toHaveBeenCalledTimes(1); + expect(exit).toHaveBeenCalledWith(0); + }); +}); diff --git a/src/server/gracefulShutdown.ts b/src/server/gracefulShutdown.ts new file mode 100644 index 00000000..2743565a --- /dev/null +++ b/src/server/gracefulShutdown.ts @@ -0,0 +1,31 @@ +/** + * Single-flight graceful shutdown helper. + * + * SIGINT/SIGTERM can arrive repeatedly while cleanup is running (a second + * Ctrl+C, a service-manager SIGTERM, etc.). Each new signal must not start a + * second cleanup — that could exit the parent before the first shutdown of + * child processes finishes. Only the first invocation runs cleanup and calls + * `process.exit`; subsequent calls return the same in-flight promise. + */ + +/** + * Creates a graceful-shutdown runner that is idempotent while in flight. + * + * @param cleanup - Callbacks that stop child processes (e.g. GenAI sidecar, MCP client). + * @returns An async shutdown runner keyed by the signal that triggered it. + */ +export function createSingleFlightShutdown(cleanup: Array<() => Promise>) { + let shutdownPromise: Promise | null = null; + return (signal: string): Promise => { + if (shutdownPromise) return shutdownPromise; + + shutdownPromise = (async () => { + // eslint-disable-next-line no-console -- intentional shutdown logging + console.log(`\n[${signal}] Shutting down.`); + await Promise.allSettled(cleanup.map((stop) => stop())); + process.exit(0); + })(); + + return shutdownPromise; + }; +} diff --git a/src/server/routes/ai/modelCatalog.ts b/src/server/routes/ai/modelCatalog.ts index 773a95d1..4e05692b 100644 --- a/src/server/routes/ai/modelCatalog.ts +++ b/src/server/routes/ai/modelCatalog.ts @@ -2,10 +2,7 @@ * Live model-catalog fetching for AI providers. * `baseUrl` arguments must already be sanitized before reaching these helpers. */ -import { - sanitizeProviderBaseUrl, - stripTrailingSlashes, -} from "../../services/ai/security.ts"; +import { sanitizeProviderBaseUrl, stripTrailingSlashes } from "../../services/ai/security.ts"; import { fetchWithTimeout } from "../../services/shared/http.ts"; import { catalogModelsFromOpenAiCompatRows, @@ -25,6 +22,25 @@ export async function fetchLiveModelCatalog(provider: string, apiKey: string, ba if (provider === "copilot") { return await fetchCopilotModelCatalog(apiKey, baseUrl); } + if (provider === "bedrock") { + // Bedrock model availability is region/account-scoped and has no + // OpenAI-style catalog endpoint; the static default list stands. + return { + models: [], + source: "fallback" as const, + error: "Bedrock models are region-scoped; use the default model ID or one enabled for your account.", + }; + } + if (provider === "genai") { + // The built-in engine serves a single locally downloaded model; it has + // no remote catalog endpoint. Return the static default list instead of + // attempting an OpenAI-style fetch. + return { + models: [], + source: "fallback" as const, + error: "Built-in GenAI runs one locally downloaded model; the default is preconfigured in Settings.", + }; + } const base = stripTrailingSlashes(baseUrl || defaultBaseUrl(provider)); const headers: Record = { diff --git a/src/server/routes/ai/providerRoutes.test.ts b/src/server/routes/ai/providerRoutes.test.ts new file mode 100644 index 00000000..91628501 --- /dev/null +++ b/src/server/routes/ai/providerRoutes.test.ts @@ -0,0 +1,202 @@ +import { describe, it, expect, vi, beforeAll, beforeEach, afterAll } from "vitest"; +import express, { type Express } from "express"; +import type { Server } from "node:http"; + +// Deterministic GenAI service stubs: never touch the real venv, disk cache, +// or network during route tests. +vi.mock("../../services/genai/venv.ts", () => ({ + isGenaiVenvReady: vi.fn(() => true), + ensureGenaiVenv: vi.fn(async () => ({ ok: true })), +})); + +vi.mock("../../services/genai/modelDownload.ts", () => ({ + DEFAULT_GENAI_MODEL: "qwen2.5-coder-1.5b-instruct-onnx", + getModelStatus: vi.fn(() => ({ + ready: true, + localPath: "", + filesPresent: 6, + filesRequired: 6, + localSizeBytes: 0, + })), + // If the loopback gate ever fails to block, this surfaces as a test error. + downloadModel: vi.fn(async () => { + throw new Error("downloadModel must not run in this test"); + }), +})); + +// In-memory provider state so activation tests never write the studio config. +const state = vi.hoisted(() => { + let runtime: { provider: string; apiKey: string; model: string; baseUrl?: string } | null = null; + return { + get: () => runtime, + set: (cfg: typeof runtime) => { + runtime = cfg; + }, + }; +}); + +vi.mock("../../services/ai/state.ts", () => ({ + getRuntimeAiProvider: () => state.get(), + setRuntimeAiProvider: (cfg: unknown) => state.set(cfg as never), + readAiPreference: () => null, + restoreProviderFromPreference: () => null, +})); + +// Machine-local env keys must not leak into activation expectations. +vi.mock("../../services/ai/env.ts", () => ({ + readEnvApiKey: vi.fn(() => undefined), + matchedEnvApiKeyName: vi.fn(() => undefined), +})); + +import { mountProviderRoutes } from "./providerRoutes.ts"; +import { downloadModel, getModelStatus } from "../../services/genai/modelDownload.ts"; +import { isGenaiVenvReady } from "../../services/genai/venv.ts"; + +let server: Server; +let baseUrl: string; + +beforeAll(async () => { + const app: Express = express(); + app.use(express.json()); + const router = express.Router(); + mountProviderRoutes(router); + app.use("/api", router); + await new Promise((resolve, reject) => { + server = app.listen(0, "127.0.0.1", () => resolve()); + server.on("error", reject); + }); + const addr = server.address(); + const port = typeof addr === "object" && addr ? addr.port : 0; + baseUrl = `http://127.0.0.1:${port}`; +}); + +beforeEach(() => { + vi.mocked(isGenaiVenvReady).mockReturnValue(true); + vi.mocked(getModelStatus).mockReturnValue({ + ready: true, + localPath: "", + filesPresent: 6, + filesRequired: 6, + localSizeBytes: 0, + }); +}); + +afterAll(async () => { + await new Promise((resolve) => server.close(() => resolve())); +}); + +describe("POST /api/ai/provider keyless activation", () => { + it("activates genai without an API key", async () => { + const res = await fetch(`${baseUrl}/api/ai/provider`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ provider: "genai", model: "qwen2.5-coder-1.5b-instruct-onnx" }), + }); + expect(res.status).toBe(200); + expect(await res.json()).toMatchObject({ ok: true, provider: "genai" }); + expect(state.get()?.provider).toBe("genai"); + }); + + it("still rejects key-required providers without credentials", async () => { + const res = await fetch(`${baseUrl}/api/ai/provider`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ provider: "gemini", model: "gemini-2.5-flash" }), + }); + expect(res.status).toBe(400); + }); + + it("rejects genai activation before the engine is installed", async () => { + vi.mocked(isGenaiVenvReady).mockReturnValue(false); + const res = await fetch(`${baseUrl}/api/ai/provider`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ provider: "genai", model: "qwen2.5-coder-1.5b-instruct-onnx" }), + }); + expect(res.status).toBe(400); + expect(await res.json()).toMatchObject({ error: expect.stringContaining("Install the GenAI engine") }); + }); + + it("rejects genai activation before the model is downloaded", async () => { + vi.mocked(getModelStatus).mockReturnValue({ + ready: false, + localPath: "", + filesPresent: 2, + filesRequired: 6, + localSizeBytes: 0, + }); + const res = await fetch(`${baseUrl}/api/ai/provider`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ provider: "genai", model: "qwen2.5-coder-1.5b-instruct-onnx" }), + }); + expect(res.status).toBe(400); + expect(await res.json()).toMatchObject({ error: expect.stringContaining("Download the GenAI model") }); + }); +}); + +describe("POST /api/ai/models keyless catalog", () => { + it("returns the region-scoped Bedrock fallback for keyless requests", async () => { + const res = await fetch(`${baseUrl}/api/ai/models`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ provider: "bedrock", baseUrl: "us-west-2" }), + }); + expect(res.status).toBe(200); + const body = (await res.json()) as { models: unknown[]; source: string; error: string }; + expect(body.models).toEqual([]); + expect(body.source).toBe("fallback"); + expect(body.error).toContain("region-scoped"); + }); + + it("returns the GenAI fallback for keyless requests instead of an API-key error", async () => { + const res = await fetch(`${baseUrl}/api/ai/models`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ provider: "genai" }), + }); + expect(res.status).toBe(200); + const body = (await res.json()) as { models: unknown[]; source: string; error: string }; + expect(body.models).toEqual([]); + expect(body.source).toBe("fallback"); + expect(body.error).toContain("locally downloaded model"); + }); + + it("still returns the API-key fallback for key-required providers", async () => { + const res = await fetch(`${baseUrl}/api/ai/models`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ provider: "gemini" }), + }); + expect(res.status).toBe(200); + const body = (await res.json()) as { models: unknown[]; source: string; error: string }; + expect(body.models).toEqual([]); + expect(body.source).toBe("fallback"); + expect(body.error).toContain("No API key available"); + }); +}); + +describe("genai engine endpoints", () => { + it("reports engine and model status", async () => { + const res = await fetch(`${baseUrl}/api/ai/genai/status`); + expect(res.status).toBe(200); + expect(await res.json()).toMatchObject({ venvReady: true, model: { ready: true } }); + }); + + it("blocks downloads that arrive via a reverse proxy hop", async () => { + const res = await fetch(`${baseUrl}/api/ai/genai/download`, { + method: "POST", + headers: { "x-forwarded-for": "203.0.113.9" }, + }); + expect(res.status).toBe(403); + expect(downloadModel).not.toHaveBeenCalled(); + }); + + it("blocks engine setup that arrives via a reverse proxy hop", async () => { + const res = await fetch(`${baseUrl}/api/ai/genai/setup`, { + method: "POST", + headers: { "x-forwarded-for": "203.0.113.9" }, + }); + expect(res.status).toBe(403); + }); +}); diff --git a/src/server/routes/ai/providerRoutes.ts b/src/server/routes/ai/providerRoutes.ts index 3d8d3aeb..381a8004 100644 --- a/src/server/routes/ai/providerRoutes.ts +++ b/src/server/routes/ai/providerRoutes.ts @@ -19,7 +19,8 @@ import { cloudflareAiBaseUrl, isValidCloudflareAccountId } from "../../../lib/cl import { getCodexAppServer } from "../../../lib/codex/CodexAppServerClient.ts"; import { listDevinModels } from "../../../lib/devin/client.ts"; import type { ProviderConfig } from "../../types.ts"; -import { authActionRateLimit } from "../../middleware/rateLimit.ts"; +import { authActionRateLimit, heavyCommandRateLimit } from "../../middleware/rateLimit.ts"; +import { studioLocalOnly } from "../../middleware/localOnly.ts"; import { isParseBodyError, parseBody } from "../../middleware/bodyGuard.ts"; import { fetchLiveModelCatalog } from "./modelCatalog.ts"; import { ensureGenaiVenv, isGenaiVenvReady } from "../../services/genai/venv.ts"; @@ -35,6 +36,17 @@ function isLocalOpenaiCompat(provider: string, normalizedBaseUrl?: string): bool } } +/** + * Whether a provider may list models without an API key. Bedrock authenticates + * through the default AWS credential chain and GenAI is a local engine — both + * have keyless catalog paths, so the key-missing fallback must not preempt + * their provider-specific catalog responses. + */ +function canListModelsWithoutKey(provider: string, normalizedBaseUrl?: string): boolean { + if (provider === "bedrock" || provider === "genai") return true; + return isLocalOpenaiCompat(provider, normalizedBaseUrl); +} + /** Resolve the API key for a live model catalog request: body > runtime > env. */ function resolveCatalogApiKey(provider: ProviderConfig["provider"], apiKey?: string): string { const explicit = typeof apiKey === "string" ? apiKey.trim() : ""; @@ -48,7 +60,10 @@ function resolveCatalogApiKey(provider: ProviderConfig["provider"], apiKey?: str } /** Resolve the base URL candidate for a live model catalog request. */ -function resolveCatalogBaseUrlCandidate(provider: ProviderConfig["provider"], baseUrl?: string): string | undefined { +function resolveCatalogBaseUrlCandidate( + provider: ProviderConfig["provider"], + baseUrl?: string, +): string | undefined { if (baseUrl) return baseUrl; const runtime = getRuntimeAiProvider(); if (runtime && runtime.provider === provider && runtime.baseUrl) return runtime.baseUrl; @@ -64,7 +79,9 @@ function allowsEmptyApiKey(provider: string, normalizedBaseUrl?: string): boolea provider === "cloudflare" || // Bedrock can authenticate through the default AWS chain (profile, IAM // role, ~/.aws/credentials) without an explicit key. - provider === "bedrock" + provider === "bedrock" || + // Built-in GenAI runs a local ONNX Runtime engine; it has no API key. + provider === "genai" ) { return true; } @@ -104,7 +121,11 @@ async function fetchSpecialProviderCatalog( await server.start(); const models = await server.listModels(); if (models.length > 0) return { models, source: "live" }; - return { models: [], source: "fallback", error: "Codex returned an empty model catalog. Sign in, then Refresh." }; + return { + models: [], + source: "fallback", + error: "Codex returned an empty model catalog. Sign in, then Refresh.", + }; } catch (err: unknown) { return fallback(err); } @@ -137,8 +158,7 @@ async function fetchSpecialProviderCatalog( } type CredentialResult = - | { ok: true; apiKey: string; baseUrl: string | undefined } - | { ok: false; error: string }; + { ok: true; apiKey: string; baseUrl: string | undefined } | { ok: false; error: string }; /** * Resolve the effective API key and base URL for a provider activation. @@ -165,7 +185,11 @@ function resolveProviderCredentials( finalBaseUrl = finalBaseUrl || cloudflareAiBaseUrl(auth.accountId); } if (!finalKey || !finalBaseUrl) { - return { ok: false, error: "Cloudflare is not signed in. Use Sign in + Sync credentials, or set CLOUDFLARE_API_TOKEN and CLOUDFLARE_ACCOUNT_ID." }; + return { + ok: false, + error: + "Cloudflare is not signed in. Use Sign in + Sync credentials, or set CLOUDFLARE_API_TOKEN and CLOUDFLARE_ACCOUNT_ID.", + }; } } return { ok: true, apiKey: finalKey, baseUrl: finalBaseUrl }; @@ -176,12 +200,14 @@ export function mountProviderRoutes(router: Router): void { return res.json({ venvReady: isGenaiVenvReady(), model: getModelStatus(DEFAULT_GENAI_MODEL) }); }); - router.post("/ai/genai/setup", authActionRateLimit, async (_req, res) => { + // Loopback-gated: these heavy endpoints trigger multi-GB disk/network work + // and must not be reachable when Studio is bound to a LAN address. + router.post("/ai/genai/setup", studioLocalOnly, heavyCommandRateLimit, async (_req, res) => { const result = await ensureGenaiVenv((line) => console.warn(line)); return res.status(result.ok ? 200 : 500).json(result); }); - router.post("/ai/genai/download", authActionRateLimit, async (_req, res) => { + router.post("/ai/genai/download", studioLocalOnly, heavyCommandRateLimit, async (_req, res) => { const result = await downloadModel(DEFAULT_GENAI_MODEL); return res.status(result.ok ? 200 : 500).json(result); }); @@ -268,6 +294,19 @@ export function mountProviderRoutes(router: Router): void { const plugin = getProvider(provider); const creds = resolveProviderCredentials(provider, apiKey, normalizedBaseUrl); if (!creds.ok) return res.status(400).json({ error: creds.error }); + // The built-in GenAI engine cannot serve requests before its Python venv + // and model are installed, so refuse to "activate" a broken runtime. + if (provider === "genai") { + if (!isGenaiVenvReady()) { + return res.status(400).json({ error: "Install the GenAI engine first, then activate the provider." }); + } + const modelId = model || DEFAULT_GENAI_MODEL; + if (!getModelStatus(modelId).ready) { + return res + .status(400) + .json({ error: `Download the GenAI model (${modelId}) first, then activate the provider.` }); + } + } setRuntimeAiProvider({ provider, apiKey: creds.apiKey, @@ -322,7 +361,7 @@ export function mountProviderRoutes(router: Router): void { }); } - const allowEmptyKey = isLocalOpenaiCompat(provider, safeBaseUrl); + const allowEmptyKey = canListModelsWithoutKey(provider, safeBaseUrl); if (!key && !allowEmptyKey) { return res.json({ models: [], diff --git a/src/server/services/ai/bedrock.test.ts b/src/server/services/ai/bedrock.test.ts new file mode 100644 index 00000000..f9a39691 --- /dev/null +++ b/src/server/services/ai/bedrock.test.ts @@ -0,0 +1,108 @@ +import { describe, it, expect, beforeEach, afterEach } from "vitest"; + +// Side-effect import (registration) + named import in one statement. +import { parsePackedCredentials } from "./bedrock.ts"; +import { getProvider } from "./registry.ts"; + +describe("parsePackedCredentials", () => { + it("parses a two-segment static credential", () => { + expect(parsePackedCredentials("AKIAIOSFODNN7EXAMPLE:wJalrXUtnFEMI")).toEqual({ + accessKeyId: "AKIAIOSFODNN7EXAMPLE", + secretAccessKey: "wJalrXUtnFEMI", + }); + }); + + it("parses a three-segment temporary credential with session token", () => { + expect(parsePackedCredentials("ASIAEXAMPLE:secret:sessionToken123")).toEqual({ + accessKeyId: "ASIAEXAMPLE", + secretAccessKey: "secret", + sessionToken: "sessionToken123", + }); + }); + + it("drops a blank trailing session token", () => { + expect(parsePackedCredentials("AKIAEXAMPLE:secret: ")).toEqual({ + accessKeyId: "AKIAEXAMPLE", + secretAccessKey: "secret", + }); + }); + + it("throws on a single-segment string instead of silently corrupting", () => { + expect(() => parsePackedCredentials("AKIAONLYONE")).toThrow(/accessKeyId:secretAccessKey/); + }); + + it("throws when the secret segment is empty", () => { + expect(() => parsePackedCredentials("AKIAEXAMPLE: ")).toThrow(/missing a secret access key/); + }); + + it("throws when the access key id segment is empty", () => { + expect(() => parsePackedCredentials(":secret")).toThrow(/missing an access key id/); + }); +}); + +describe("bedrock buildConfig", () => { + const saved = { + secret: process.env.AWS_SECRET_ACCESS_KEY, + token: process.env.AWS_SESSION_TOKEN, + region: process.env.AWS_REGION, + }; + + beforeEach(() => { + delete process.env.AWS_SECRET_ACCESS_KEY; + delete process.env.AWS_SESSION_TOKEN; + delete process.env.AWS_REGION; + }); + + afterEach(() => { + if (saved.secret === undefined) delete process.env.AWS_SECRET_ACCESS_KEY; + else process.env.AWS_SECRET_ACCESS_KEY = saved.secret; + if (saved.token === undefined) delete process.env.AWS_SESSION_TOKEN; + else process.env.AWS_SESSION_TOKEN = saved.token; + if (saved.region === undefined) delete process.env.AWS_REGION; + else process.env.AWS_REGION = saved.region; + }); + + it("packs access key and secret for static credentials", () => { + process.env.AWS_SECRET_ACCESS_KEY = "staticSecret"; + const cfg = getProvider("bedrock")!.buildConfig("AKIAEXAMPLE"); + expect(cfg.apiKey).toBe("AKIAEXAMPLE:staticSecret"); + }); + + it("packs AWS_SESSION_TOKEN as a third segment for assumed-role credentials", () => { + process.env.AWS_SECRET_ACCESS_KEY = "tempSecret"; + process.env.AWS_SESSION_TOKEN = "tempToken"; + const cfg = getProvider("bedrock")!.buildConfig("ASIAEXAMPLE"); + expect(cfg.apiKey).toBe("ASIAEXAMPLE:tempSecret:tempToken"); + }); + + it("round-trips packed temporary credentials through the parser", () => { + process.env.AWS_SECRET_ACCESS_KEY = "tempSecret"; + process.env.AWS_SESSION_TOKEN = "tempToken"; + const cfg = getProvider("bedrock")!.buildConfig("ASIAEXAMPLE"); + expect(parsePackedCredentials(cfg.apiKey)).toEqual({ + accessKeyId: "ASIAEXAMPLE", + secretAccessKey: "tempSecret", + sessionToken: "tempToken", + }); + }); + + it("never mixes an env session token into a long-term (AKIA) key", () => { + process.env.AWS_SECRET_ACCESS_KEY = "staticSecret"; + process.env.AWS_SESSION_TOKEN = "staleToken"; + const cfg = getProvider("bedrock")!.buildConfig("AKIAEXAMPLE"); + expect(cfg.apiKey).toBe("AKIAEXAMPLE:staticSecret"); + }); + + it("does not pack a profile name from AWS_PROFILE detection into a credential", () => { + delete process.env.AWS_SECRET_ACCESS_KEY; + const cfg = getProvider("bedrock")!.buildConfig("my-production-profile"); + expect(cfg.apiKey).toBe(""); + }); + + it("does not pack when the secret key is missing", () => { + process.env.AWS_SESSION_TOKEN = "tempToken"; + const cfg = getProvider("bedrock")!.buildConfig("ASIAEXAMPLE"); + // No secret key available → fall back to the default credential chain. + expect(cfg.apiKey).toBe(""); + }); +}); diff --git a/src/server/services/ai/bedrock.ts b/src/server/services/ai/bedrock.ts index ff479199..f25f98ab 100644 --- a/src/server/services/ai/bedrock.ts +++ b/src/server/services/ai/bedrock.ts @@ -5,6 +5,8 @@ * 1. UI-entered: user pastes AWS access key + secret key + region in the settings panel. * These are passed via ProviderConfig.apiKey as a packed string: * `:` (region comes from baseUrl field as the region string). + * An optional third segment carries the STS session token for assumed-role / + * temporary credentials: `::`. * 2. Auto-detect: if no explicit keys are provided, falls back to the default AWS credential * chain (env vars, ~/.aws/credentials, IAM role, etc.) via @aws-sdk/credential-providers. * @@ -66,6 +68,41 @@ function hasHalfPackedCredentials(cfg: ProviderConfig): boolean { return cfg.apiKey.slice(cfg.apiKey.indexOf(":") + 1).trim().length === 0; } +/** + * Splits a packed `accessKeyId:secretAccessKey[:sessionToken]` string. + * AWS key material never contains ":", so segment boundaries are unambiguous. + * Invalid packed formats fail loudly rather than silently producing corrupted + * credentials that would authenticate as a different account. + */ +export function parsePackedCredentials(packed: string): { + accessKeyId: string; + secretAccessKey: string; + sessionToken?: string; +} { + const colonIdx = packed.indexOf(":"); + if (colonIdx === -1) { + throw new Error( + "Invalid packed Bedrock credentials: expected accessKeyId:secretAccessKey[:sessionToken].", + ); + } + const accessKeyId = packed.slice(0, colonIdx).trim(); + if (!accessKeyId) { + throw new Error("Packed Bedrock credentials are missing an access key id."); + } + const rest = packed.slice(colonIdx + 1); + const secondColon = rest.indexOf(":"); + const secretAccessKey = (secondColon === -1 ? rest : rest.slice(0, secondColon)).trim(); + if (!secretAccessKey) { + throw new Error("Packed Bedrock credentials are missing a secret access key."); + } + const sessionToken = secondColon === -1 ? undefined : rest.slice(secondColon + 1).trim(); + return { + accessKeyId, + secretAccessKey, + ...(sessionToken ? { sessionToken } : {}), + }; +} + /** * Creates a BedrockRuntimeClient with appropriate credentials. * @@ -84,14 +121,15 @@ function createBedrockClient(cfg: ProviderConfig): BedrockRuntimeClient { const region = resolveRegion(cfg); if (hasExplicitCredentials(cfg)) { - const colonIdx = cfg.apiKey.indexOf(":"); - const accessKeyId = cfg.apiKey.slice(0, colonIdx); - const secretAccessKey = cfg.apiKey.slice(colonIdx + 1); + // Temporary / assumed-role credentials carry a third packed segment; + // static keys omit it and no env token is mixed in. + const { accessKeyId, secretAccessKey, sessionToken } = parsePackedCredentials(cfg.apiKey); return new BedrockRuntimeClient({ region, credentials: { accessKeyId, secretAccessKey, + ...(sessionToken ? { sessionToken } : {}), }, }); } @@ -185,13 +223,24 @@ registerProvider({ // AWS_PROFILE alone is enough for the default chain, so detect it too. envVarNames: ["AWS_ACCESS_KEY_ID", "AWS_PROFILE"], buildConfig: (apiKey) => { - const secretKey = process.env.AWS_SECRET_ACCESS_KEY?.trim() ?? ""; - return { - provider: "bedrock", - apiKey: `${apiKey}:${secretKey}`, - model: "anthropic.claude-3-5-haiku-20241022-v1:0", - baseUrl: process.env.AWS_REGION?.trim() || "us-east-1", - }; + const defaultModel = "anthropic.claude-3-5-haiku-20241022-v1:0"; + const region = process.env.AWS_REGION?.trim() || "us-east-1"; + const accessKeyId = apiKey.trim(); + // Env detection may match AWS_PROFILE, which yields a profile name rather + // than an access key id — packing it would produce an invalid credential. + // Also bail when the secret key is absent so we never send a half key; + // createBedrockClient falls back to the default AWS credential chain. + if (!/^A[KS]IA/.test(accessKeyId) || !process.env.AWS_SECRET_ACCESS_KEY?.trim()) { + return { provider: "bedrock", apiKey: "", model: defaultModel, baseUrl: region }; + } + const secretKey = process.env.AWS_SECRET_ACCESS_KEY!.trim(); + // Only temporary (ASIA) credentials carry a session token; never mix an + // env token into a long-term (AKIA) key. + const sessionToken = /^ASIA/.test(accessKeyId) ? process.env.AWS_SESSION_TOKEN?.trim() : undefined; + const packedKey = sessionToken + ? `${accessKeyId}:${secretKey}:${sessionToken}` + : `${accessKeyId}:${secretKey}`; + return { provider: "bedrock", apiKey: packedKey, model: defaultModel, baseUrl: region }; }, call, supportsJsonResponseFormat: false, diff --git a/src/server/services/ai/state.test.ts b/src/server/services/ai/state.test.ts new file mode 100644 index 00000000..75cdb62a --- /dev/null +++ b/src/server/services/ai/state.test.ts @@ -0,0 +1,59 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +// Mock env + config BEFORE state/provider imports so nothing touches disk or +// machine-local credentials. +vi.mock("./env.ts", () => ({ + readEnvApiKey: vi.fn(), + matchedEnvApiKeyName: vi.fn(), +})); + +vi.mock("../../config.ts", () => ({ + readStudioConfig: vi.fn(() => ({})), + writeStudioConfig: vi.fn(), +})); + +// Cloudflare auth reads a local credential file and bypasses readEnvApiKey — +// stub it so machine-local credentials cannot leak into restore tests. +vi.mock("../../../lib/cloudflare/client.ts", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + resolveCloudflareAuth: vi.fn(() => null), + }; +}); + +// Side-effect import: triggers all providers to register themselves. +import "./index.ts"; + +import { readEnvApiKey } from "./env.ts"; +import { restoreProviderFromPreference, type AiPreference } from "./state.ts"; + +describe("restoreProviderFromPreference", () => { + beforeEach(() => { + vi.mocked(readEnvApiKey).mockReturnValue(undefined); + }); + + it("restores genai without any API key (local engine)", () => { + const pref: AiPreference = { provider: "genai", model: "qwen2.5-coder-1.5b-instruct-onnx" }; + const cfg = restoreProviderFromPreference(pref); + expect(cfg?.provider).toBe("genai"); + expect(cfg?.model).toBe("qwen2.5-coder-1.5b-instruct-onnx"); + }); + + it("restores bedrock on the default AWS credential chain alone", () => { + const pref: AiPreference = { + provider: "bedrock", + model: "anthropic.claude-3-5-haiku-20241022-v1:0", + baseUrl: "us-west-2", + }; + const cfg = restoreProviderFromPreference(pref); + expect(cfg?.provider).toBe("bedrock"); + // The baseUrl field carries the AWS region for bedrock. + expect(cfg?.baseUrl).toBe("us-west-2"); + }); + + it("still rejects key-required providers without env credentials", () => { + const pref: AiPreference = { provider: "gemini", model: "gemini-2.5-flash" }; + expect(restoreProviderFromPreference(pref)).toBeNull(); + }); +}); diff --git a/src/server/services/ai/state.ts b/src/server/services/ai/state.ts index 70f33938..9fa63b97 100644 --- a/src/server/services/ai/state.ts +++ b/src/server/services/ai/state.ts @@ -75,6 +75,8 @@ export function restoreProviderFromPreference(pref: AiPreference): ProviderConfi pref.provider === "devin" || // Bedrock can restore on the default AWS credential chain alone. pref.provider === "bedrock" || + // Built-in GenAI runs a local engine and never stores an API key. + pref.provider === "genai" || Boolean(pref.baseUrl && /localhost|127\.0\.0\.1/i.test(pref.baseUrl)); if (pref.provider === "cloudflare") { diff --git a/src/server/services/genai/venv.test.ts b/src/server/services/genai/venv.test.ts new file mode 100644 index 00000000..3bbf5edd --- /dev/null +++ b/src/server/services/genai/venv.test.ts @@ -0,0 +1,155 @@ +/** + * Unit tests for the GenAI venv service: + * - ensureGenaiVenv serializes concurrent setup into a single operation + * (remounts / multiple callers must not create+install the same venv) + * - sidecar exitPromise settles on spawn 'error' as well as 'exit' + * - shutdownSidecar escalates to a hard kill when graceful shutdown times out + * + * `child_process` is mocked via `src/server/__tests__/childProcessTestMocks.ts`. + */ +import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; +import { EventEmitter } from "events"; +import type { ChildProcess } from "child_process"; +import fs from "fs"; + +const mocks = vi.hoisted(() => ({ + execFileImpl: null as null | ((...args: unknown[]) => unknown), + spawnImpl: null as null | ((...args: unknown[]) => unknown), + execFileCalls: [] as unknown[][], +})); + +vi.mock("child_process", async (importOriginal) => { + const { childProcessVitestMockFactory } = await import("../../__tests__/childProcessTestMocks.ts"); + return childProcessVitestMockFactory(mocks, { includeSpawn: true })(importOriginal); +}); + +import { ensureGenaiVenv, shutdownSidecar, spawnSidecar } from "./venv.ts"; + +/** + * Makes the GenAI venv look "not ready" while letting the mocked venv + * creation flip it to ready afterwards. mkdirSync is stubbed so tests never + * create a real `.venvs/` directory in the repo. + */ +function stubVenvNotReady(): void { + let venvCreated = false; + vi.spyOn(fs, "existsSync").mockImplementation((p) => { + const path = String(p); + if (path.includes(".venvs") && path.includes("genai") && /python(\.exe)?$/.test(path)) { + return venvCreated; + } + return true; + }); + vi.spyOn(fs, "mkdirSync").mockImplementation(() => "" as never); + mocks.execFileImpl = async (file: unknown, args: unknown) => { + const a = args as string[]; + if (a.includes("--version")) return { stdout: "Python 3.12.0", stderr: "" }; + if (a.includes("venv")) venvCreated = true; + return { stdout: "", stderr: "" }; + }; +} + +/** A fake child process that never exits on its own (caller emits events). */ +function inertChildProcess(): ChildProcess { + const proc = new EventEmitter() as unknown as ChildProcess; + proc.stdout = new EventEmitter() as ChildProcess["stdout"]; + proc.stderr = new EventEmitter() as ChildProcess["stderr"]; + proc.stdin = new EventEmitter() as ChildProcess["stdin"]; + const record = proc as unknown as Record; + record.pid = 4242; + record.exitCode = null; + record.killed = false; + record.kill = () => true; + return proc; +} + +describe("ensureGenaiVenv", () => { + beforeEach(() => { + mocks.execFileImpl = null; + mocks.spawnImpl = null; + mocks.execFileCalls.length = 0; + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("shares one setup operation across concurrent callers", async () => { + stubVenvNotReady(); + let execCalls = 0; + const realImpl = mocks.execFileImpl!; + mocks.execFileImpl = async (...args) => { + execCalls += 1; + await new Promise((resolve) => setTimeout(resolve, 10)); + return realImpl(...args); + }; + + const progress: string[][] = [[], []]; + const [a, b] = await Promise.all([ + ensureGenaiVenv((line) => progress[0].push(line)), + ensureGenaiVenv((line) => progress[1].push(line)), + ]); + + expect(a.ok).toBe(true); + expect(b.ok).toBe(true); + // Two callers must not run the python probe / venv create / pip / verify + // sequence twice — that is exactly one full setup. + expect(execCalls).toBe(4); + // The late joiner still receives progress lines from the shared setup. + expect(progress[1].length).toBeGreaterThan(0); + expect(progress[1]).toContain("[genai] Setup complete. onnxruntime-genai ready."); + }); + + it("returns an error result when no system python is found", async () => { + stubVenvNotReady(); + mocks.execFileImpl = async () => { + throw new Error("ENOENT"); + }; + const result = await ensureGenaiVenv(() => {}); + expect(result.ok).toBe(false); + expect(result.error).toContain("Python >=3.10 not found"); + }); +}); + +describe("spawnSidecar lifecycle", () => { + beforeEach(() => { + mocks.execFileImpl = null; + mocks.spawnImpl = null; + }); + + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it("settles exitPromise when the sidecar fails to spawn (error, not exit)", async () => { + const proc = inertChildProcess(); + mocks.spawnImpl = () => proc; + + const sidecar = spawnSidecar("/tmp/whatever", "cpu"); + proc.emit("error", new Error("ENOENT")); + + await expect(sidecar.exitPromise).resolves.toBeNull(); + expect(sidecar.alive()).toBe(false); + }); + + it("hard-kills a sidecar that ignores graceful shutdown after the wait cap", async () => { + vi.useFakeTimers(); + const proc = inertChildProcess(); + const kills: string[] = []; + (proc as unknown as Record).kill = (sig?: string) => { + kills.push(sig ?? ""); + return true; + }; + mocks.spawnImpl = () => proc; + + spawnSidecar("/tmp/whatever", "cpu"); + const pending = shutdownSidecar(); + + await vi.advanceTimersByTimeAsync(2000); // graceful SIGTERM escalator + await vi.advanceTimersByTimeAsync(3000); // past the 5s cap → SIGKILL + await pending; + + expect(kills).toContain("SIGTERM"); + expect(kills).toContain("SIGKILL"); + }); +}); diff --git a/src/server/services/genai/venv.ts b/src/server/services/genai/venv.ts index cb5aeb5d..3d58f695 100644 --- a/src/server/services/genai/venv.ts +++ b/src/server/services/genai/venv.ts @@ -64,8 +64,14 @@ export type SetupListener = (line: string) => void; async function findSystemPython(): Promise { const candidates: Array<[string, string[]]> = process.platform === "win32" - ? [["py", ["-3", "--version"]], ["python", ["--version"]]] - : [["python3", ["--version"]], ["python", ["--version"]]]; + ? [ + ["py", ["-3", "--version"]], + ["python", ["--version"]], + ] + : [ + ["python3", ["--version"]], + ["python", ["--version"]], + ]; for (const [cmd, args] of candidates) { try { @@ -86,16 +92,40 @@ async function findSystemPython(): Promise { return null; } +/** + * In-flight setup operation shared by concurrent callers so two requests can + * never create/install the same venv at the same time (remounts, multiple + * panels, or direct loopback calls all serialize on this promise). + */ +let activeVenvSetup: Promise<{ ok: boolean; error?: string }> | null = null; +/** Progress listeners for every caller waiting on the shared setup operation. */ +const setupListeners = new Set(); + /** * Ensures the GenAI venv is set up with onnxruntime-genai installed. - * Idempotent — skips if already ready. + * Idempotent — skips if already ready. Concurrent callers await the same + * in-flight setup and receive its progress lines and result. * * @param onLine - Progress callback for UI streaming. * @returns Success/failure result. */ -export async function ensureGenaiVenv( - onLine: SetupListener, -): Promise<{ ok: boolean; error?: string }> { +export function ensureGenaiVenv(onLine: SetupListener): Promise<{ ok: boolean; error?: string }> { + setupListeners.add(onLine); + if (!activeVenvSetup) { + activeVenvSetup = (async () => { + const operation = runGenaiVenvSetup((line) => { + for (const listener of setupListeners) listener(line); + }); + return operation; + })().finally(() => { + activeVenvSetup = null; + setupListeners.clear(); + }); + } + return activeVenvSetup; +} + +async function runGenaiVenvSetup(onLine: SetupListener): Promise<{ ok: boolean; error?: string }> { if (isGenaiVenvReady()) { // Quick check: is onnxruntime-genai importable? try { @@ -183,10 +213,14 @@ export interface SidecarProcess { onTerminate: (settle: () => void) => () => void; /** Settle all in-flight requests immediately (replacement/shutdown). */ settlePending: () => void; - /** Kill the sidecar process. */ + /** Kill the sidecar process (graceful shutdown command, then SIGTERM). */ kill: () => void; + /** Force-kill the sidecar (SIGKILL) after a graceful shutdown timeout. */ + killHard: () => void; /** Whether the process is still alive. */ alive: () => boolean; + /** Resolves with the exit code once the child process has exited. */ + exitPromise: Promise; } let activeSidecar: SidecarProcess | null = null; @@ -226,6 +260,26 @@ export function spawnSidecar(modelPath: string, ep: string = "cpu"): SidecarProc }, }); + // Lets shutdownSidecar await actual process exit instead of racing it. A + // failed spawn emits 'error' (never 'exit'), so both events settle the same + // idempotent resolver — otherwise shutdown would wait out its full timeout. + let resolveExit: (code: number | null) => void = () => {}; + let exitSettled = false; + const exitPromise = new Promise((resolve) => { + resolveExit = (code) => { + if (exitSettled) return; + exitSettled = true; + resolve(code); + }; + }); + child.once("exit", (code) => resolveExit(code)); + child.once("error", () => resolveExit(null)); + + // A failed spawn ('error' event) leaves exitCode/killed untouched, so the + // sidecar must be marked dead explicitly or shutdown would try to talk to a + // process that was never started. + let spawnFailed = false; + const responseHandlers: Array<(data: Record) => void> = []; // Settlers for requests that are waiting on this process. Invoked when the // process dies or is replaced so callers reject instead of timing out. @@ -236,7 +290,9 @@ export function spawnSidecar(modelPath: string, ep: string = "cpu"): SidecarProc for (const settle of settlers) { try { settle(); - } catch { /* a failing settler must not block teardown */ } + } catch { + /* a failing settler must not block teardown */ + } } }; let buffer = ""; @@ -267,6 +323,7 @@ export function spawnSidecar(modelPath: string, ep: string = "cpu"): SidecarProc // that would crash the whole server. child.on("error", (err) => { console.warn("[genai-sidecar] spawn error:", err.message); + spawnFailed = true; settlePending(); if (activeSidecar === sidecar) { activeSidecar = null; @@ -309,12 +366,24 @@ export function spawnSidecar(modelPath: string, ep: string = "cpu"): SidecarProc settlePending(); try { sidecar.send({ command: "shutdown" }); - } catch { /* ignore */ } + } catch { + /* ignore */ + } setTimeout(() => { if (child.exitCode === null) child.kill("SIGTERM"); }, 2000); }, - alive: () => child.exitCode === null && !child.killed, + killHard: () => { + if (child.exitCode !== null || child.killed) return; + try { + // Windows maps SIGKILL to TerminateProcess; POSIX sends SIGKILL. + child.kill("SIGKILL"); + } catch { + /* already gone */ + } + }, + alive: () => !spawnFailed && child.exitCode === null && !child.killed, + exitPromise, }; activeSidecar = sidecar; @@ -339,9 +408,35 @@ export function getActiveSidecar(modelPath?: string, ep?: string): SidecarProces return null; } -/** Shutdown the active sidecar cleanly. */ -export function shutdownSidecar(): void { - activeSidecar?.kill(); +/** + * Shutdown the active sidecar cleanly. Waits (bounded) for the child to exit + * so server shutdown does not orphan the inference process with its loaded + * model still holding CPU/GPU memory. + */ +export async function shutdownSidecar(): Promise { + const sidecar = activeSidecar; activeSidecar = null; activeSidecarKey = null; + if (!sidecar?.alive()) return; + sidecar.kill(); + // kill() sends the shutdown command now and SIGTERM after 2s; cap the wait + // so a wedged sidecar can never hang server shutdown. + let timedOut = false; + await Promise.race([ + sidecar.exitPromise.then(() => { + timedOut = false; + }), + new Promise((resolve) => + setTimeout(() => { + timedOut = true; + resolve(); + }, 5000), + ), + ]); + // A sidecar that ignores the graceful shutdown (command + SIGTERM) would + // keep its loaded model in memory after Studio exits — escalate to SIGKILL. + if (timedOut) { + console.warn("[genai-sidecar] shutdown timed out; sending SIGKILL."); + sidecar.killHard(); + } }