From d6c47f666448a68cb47e7158b8e254e1b8a7fdae Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 14:46:03 +0000 Subject: [PATCH] Studio: wire Codex provider through the UI end-to-end Post-review pass driven by reviewer.py. The original PR shipped the backend codex provider, the registry entry (with `hidden:true`), the status API, and the `CodexParallelTabs` component, but the chat UI never surfaced the row, required an API key for the connection, and never sent `parallel_calls` over the wire. Also fixes a CodeQL leak in the parallel fan-out error path and adds the canonical streaming hook upstream actually exposes. Frontend * chat-providers-dialog.tsx now calls `/api/codex/status` alongside `/api/providers/registry`. When the host has Codex installed the Add connection dialog gains a synthetic Codex row (curated model list comes from `supported_models`) so the picker is reachable. * The Add / Edit connection guards now skip the API-key requirement for Codex the same way they do for the custom OpenAI-compat presets; the field itself is also hidden so the user is not asked for a key Studio will not use. * chat-adapter.ts now also exempts Codex from the "Missing API key" pre-flight, and emits `parallel_calls` on the outgoing request when the selected connection is Codex (clamped to [1, 20] by the shared helper, defaults to 1). * external-providers.ts adds `codexParallelCalls` to ExternalProviderConfig so future composer UI can persist the user's pick per connection. Backend * `_stream_thread_run` now tries `thread.turn(prompt).stream()` first, mirroring the canonical openai_codex API (`openai/codex/sdk/python/src/openai_codex/api.py`). The legacy `thread.run_streaming(prompt)` path is kept as a fallback and the buffered `await thread.run(prompt)` stays as the last resort. * `_stream_codex_parallel` no longer echoes `str(exc)` in the `codex_tab_error` SSE event. Per-tab failures now surface a generic "Codex tab failed" message plus an `exception_type` discriminator; `CodexUnavailableError` is the only exception whose text is forwarded verbatim because it is a user-actionable install hint with no sensitive content (CodeQL `py/information-exposure-through-exception`). Tests * New `TestCodexHardenedRegressions::test_parallel_tab_error_sanitised` injects a fake SDK that raises with a path-like message and asserts the SSE frames do not echo it. * New `TestCodexHardenedRegressions::test_thread_turn_stream_path_taken` verifies the canonical `thread.turn(prompt).stream()` hook is preferred over the legacy helper. All 26 codex_provider tests pass. Frontend `tsc --noEmit` clean. --- .../backend/core/inference/codex_provider.py | 70 ++++++++++++--- studio/backend/tests/test_codex_provider.py | 89 +++++++++++++++++++ .../src/features/chat/api/chat-adapter.ts | 26 +++++- .../features/chat/chat-providers-dialog.tsx | 60 ++++++++++--- .../src/features/chat/external-providers.ts | 7 ++ 5 files changed, 227 insertions(+), 25 deletions(-) diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 437722d6a2..6461dfbdf0 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -307,20 +307,51 @@ async def _stream_thread_run( ) -> AsyncGenerator[str, None]: """Yield raw text chunks from a Codex thread. - Prefers ``thread.run_streaming(prompt)`` because that's what the - docs surface for token-by-token delivery. When the installed SDK - doesn't have that helper, fall back to ``await thread.run(prompt)`` - and yield the full text once -- this still works end-to-end, just - without streaming feedback in the UI. + Tries three SDK surfaces in order: + + 1. ``thread.turn(prompt).stream()`` -- the canonical streaming path + on the upstream ``openai_codex`` SDK (see + ``openai/codex/sdk/python/src/openai_codex/api.py``: AsyncThread.turn + returns an AsyncTurnHandle whose ``.stream()`` yields events). + 2. ``thread.run_streaming(prompt)`` -- a legacy helper exposed by + some earlier SDK pre-releases. Kept for forward-compat. + 3. ``await thread.run(prompt)`` -- the always-supported buffered + path. Used when neither streaming helper resolves and as the + final fallback. + + On any streaming exception we log and fall through to the buffered + path so a partially-broken streaming helper does not take the + whole turn down. """ + # 1. Canonical: thread.turn(prompt).stream() + turn_factory = getattr(thread, "turn", None) + if turn_factory is not None: + try: + turn_handle = turn_factory(prompt) + if asyncio.iscoroutine(turn_handle): + turn_handle = await turn_handle + stream_fn = getattr(turn_handle, "stream", None) + if stream_fn is not None: + stream_obj = stream_fn() + if asyncio.iscoroutine(stream_obj): + stream_obj = await stream_obj + async for event in stream_obj: + text = _coerce_text(getattr(event, "payload", event)) + if text: + yield text + return + except Exception as exc: + logger.warning( + "codex_provider.turn_stream_failed_fallback", + exc_type = type(exc).__name__, + error = str(exc), + ) + + # 2. Legacy: thread.run_streaming(prompt) run_streaming = getattr(thread, "run_streaming", None) if run_streaming is not None: try: stream_obj = run_streaming(prompt) - # The SDK may return either an async iterator directly or a - # coroutine that resolves to one. Handle both shapes so a - # future SDK rev doesn't silently fall off the streaming - # path. if asyncio.iscoroutine(stream_obj): stream_obj = await stream_obj async for event in stream_obj: @@ -331,12 +362,11 @@ async def _stream_thread_run( except Exception as exc: logger.warning( "codex_provider.run_streaming_failed_fallback", + exc_type = type(exc).__name__, error = str(exc), ) - # Intentional fallthrough to the non-streaming path so a - # broken streaming helper doesn't take the whole turn down. - # Non-streaming fallback: await the full TurnResult, emit one chunk. + # 3. Buffered fallback: await the full TurnResult, emit one chunk. result = await thread.run(prompt) text = _coerce_text(result) or getattr(result, "final_response", "") or str(result) if text: @@ -516,18 +546,32 @@ async def _stream_codex_parallel( ) ) except Exception as exc: + # CodeQL: never echo str(exc) in client-facing SSE events. + # Log full reason server-side; surface a generic message plus + # an exception_type discriminator so the UI can still group + # failures without leaking file paths / env vars from the + # SDK traceback. CodexUnavailableError is the one exception + # we DO surface verbatim because it's a user-actionable + # install hint with no sensitive content. logger.warning( "codex_provider.parallel_tab_failed", tab_id = tab_id, + exc_type = type(exc).__name__, error = str(exc), ) + public_error = ( + str(exc) + if isinstance(exc, CodexUnavailableError) + else "Codex tab failed" + ) await queue.put( _chunk_tool_event( completion_id, { "type": "codex_tab_error", "tab_id": tab_id, - "error": str(exc), + "error": public_error, + "exception_type": type(exc).__name__, }, ) ) diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index b6a255a785..e41e207685 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -643,6 +643,95 @@ class TestCodexHardenedRegressions: if ls.startswith("yield ") and re.search(r"\{exc\}|\{e\}", ls): assert False, f"raw exception leaked: {ls}" + def test_parallel_tab_error_sanitised(self, monkeypatch): + """A worker that raises with a path-leaking message must NOT + send that text to the client; the SSE codex_tab_error event + must carry a generic message + exception_type. + """ + fake = _FakeAsyncCodex(raise_on_start=RuntimeError( + "secret /home/alice/.codex/config.json token=abc" + )) + _install_fake_codex_sdk(monkeypatch, lambda: fake) + from core.inference.codex_provider import stream_codex + + chunks: list[str] = [] + + async def _collect(): + async for c in stream_codex( + messages=[{"role": "user", "content": "hi"}], + model="gpt-5.5", + parallel_calls=2, + ): + chunks.append(c) + + asyncio.run(_collect()) + body = "".join(chunks) + assert "secret /home/alice" not in body, ( + "raw exception text leaked into codex_tab_error SSE frame" + ) + assert "Codex tab failed" in body or "exception_type" in body + + def test_thread_turn_stream_path_taken(self, monkeypatch): + """The canonical openai_codex API uses thread.turn(prompt).stream(); + the provider must prefer that over the legacy run_streaming hook. + """ + events_seen = {"turn_called": False, "run_streaming_called": False} + + class _TurnEvent: + def __init__(self, txt): + self.payload = {"text": txt} + + class _TurnHandle: + def __init__(self, prompt): + self.prompt = prompt + + async def stream(self): + yield _TurnEvent("hello ") + yield _TurnEvent("from turn.stream") + + class _ThreadWithTurn: + def turn(self, prompt): + events_seen["turn_called"] = True + return _TurnHandle(prompt) + + def run_streaming(self, prompt): + events_seen["run_streaming_called"] = True + raise AssertionError("should not be called when turn().stream() works") + + async def run(self, prompt): + raise AssertionError("should not fall through to buffered run()") + + class _Async: + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def thread_start(self, **kw): + return _ThreadWithTurn() + + _install_fake_codex_sdk(monkeypatch, _Async) + from core.inference.codex_provider import stream_codex + + chunks: list[str] = [] + + async def _collect(): + async for c in stream_codex( + messages=[{"role": "user", "content": "hi"}], + model="gpt-5.5", + parallel_calls=1, + ): + chunks.append(c) + + asyncio.run(_collect()) + assert events_seen["turn_called"], "thread.turn() never called" + assert not events_seen["run_streaming_called"] + body = "".join(chunks) + # Each text chunk wraps in its own SSE delta, so check both pieces. + assert '"content": "hello "' in body + assert '"content": "from turn.stream"' in body + async def _consume_first(gen): """Drive an async generator until it raises or yields its first diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 0c557f1b01..7a5554342c 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -7,7 +7,10 @@ import { toast } from "@/lib/toast"; import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core"; import type { ChatModelAdapter } from "@assistant-ui/react"; import { + CODEX_DEFAULT_PARALLEL_CALLS, + clampCodexParallelCalls, getExternalProviderApiKey, + isCodexProviderType, isCustomProviderType, isPromptCacheTtl, loadExternalProviders, @@ -862,12 +865,18 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { throw new Error("Connection not found."); } // Local providers (llama.cpp / vLLM / Ollama) allow an empty key — only block hosted providers. + // Codex dispatches via the local CLI / SDK, no HTTP API key. const externalProviderIsCustom = externalProvider ? isCustomProviderType(externalProvider.providerType) : false; - if (isExternalRequest && !externalApiKey && !externalProviderIsCustom) { + const externalProviderIsCodex = externalProvider + ? isCodexProviderType(externalProvider.providerType) + : false; + const externalProviderNeedsApiKey = + isExternalRequest && !externalProviderIsCustom && !externalProviderIsCodex; + if (externalProviderNeedsApiKey && !externalApiKey) { toast.error("Missing API key for selected connection.", { - description: "Open Settings → Connections and set the API key again.", + description: "Open Settings > Connections and set the API key again.", }); throw new Error("Missing connection API key."); } @@ -1489,6 +1498,19 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } : { enable_thinking: reasoningEnabled } : {}), + // Codex provider only: ask the backend to fan the turn out + // across N parallel Codex tasks and synthesise a unified + // answer. The picker UI uses the provider config's + // `codexParallelCalls` field; default of 1 keeps the + // single-call path. Backend clamps to [1, 20]. + ...(externalProviderIsCodex + ? { + parallel_calls: clampCodexParallelCalls( + externalProvider.codexParallelCalls ?? + CODEX_DEFAULT_PARALLEL_CALLS, + ), + } + : {}), }; } diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index 3ffb3a1441..8a95b3a6a6 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -49,6 +49,7 @@ import { } from "./api/providers-api"; import type { ExternalProviderConfig } from "./external-providers"; import { + CODEX_PROVIDER_TYPE, CUSTOM_BACKEND_PROVIDER_TYPE, CUSTOM_PROVIDER_PRESETS, allowsManualModelIdsWithCatalog, @@ -57,6 +58,7 @@ import { customProviderModelIdsPlaceholder, customPresetSkipsApiKeyField, getExternalProviderApiKey, + isCodexProviderType, isCustomProviderType, LEGACY_CUSTOM_PROVIDER_TYPE, removeExternalProviderApiKey, @@ -66,6 +68,7 @@ import { supportsRemoteModelCatalog, toExternalBackendProviderType, } from "./external-providers"; +import { fetchCodexStatus } from "./api/codex-api"; import { useExternalProvidersStore } from "./stores/external-providers-store"; /** Matches navbar / thread layout easing (see index.css --ease-out-quart) */ @@ -239,9 +242,15 @@ export function ChatProvidersSettings({ (s) => s.setConnectionsEnabled, ); const isCustomProvider = isCustomProviderType(providerType); + const isCodexProvider = isCodexProviderType(providerType); // Local presets (Ollama, llama.cpp) never use API keys — hide the field. - // vLLM may optionally use a bearer token on secured deployments. - const showApiKeyField = !customPresetSkipsApiKeyField(providerType); + // vLLM may optionally use a bearer token on secured deployments. Codex + // dispatches via the local CLI / SDK, no HTTP API key either. + const showApiKeyField = + !customPresetSkipsApiKeyField(providerType) && !isCodexProvider; + // Codex behaves like a "custom" provider for the gate logic below: the + // backend skips the api_key requirement entirely for `provider_type=codex`. + const providerSkipsApiKey = isCustomProvider || isCodexProvider; const showReasoningToggle = supportsProviderReasoningToggle(providerType); const registryByType = useMemo( @@ -279,7 +288,7 @@ export function ChatProvidersSettings({ const missingModelCatalogBaseUrl = supportsRemoteModelCatalog(providerType) && baseUrlDraft.trim().length === 0; const missingModelCatalogApiKey = - !isCustomProvider && !isCuratedModelList && apiKey.trim().length === 0; + !providerSkipsApiKey && !isCuratedModelList && apiKey.trim().length === 0; const loadModelsDisabled = modelsLoading || mutatingProvider || @@ -354,12 +363,41 @@ export function ChatProvidersSettings({ } let syncSucceeded = false; try { - const [registryRows, configRows] = await Promise.all([ - listProviderRegistry(), - listProviderConfigs(), - ]); + // Probe Codex availability in parallel with the registry / configs. + // Codex stays `hidden:true` in the backend registry so it is filtered + // out of `/api/providers/registry`; we synthesise a row here when + // the host has both the CLI and the SDK installed. + const [registryRowsRaw, configRows, codexStatusRaw] = await Promise.all( + [ + listProviderRegistry(), + listProviderConfigs(), + fetchCodexStatus().catch(() => null), + ], + ); if (!isMounted) return; syncSucceeded = true; + const registryRows: ProviderRegistryEntry[] = + codexStatusRaw && codexStatusRaw.installed && + !registryRowsRaw.some( + (entry) => entry.provider_type === CODEX_PROVIDER_TYPE, + ) + ? [ + ...registryRowsRaw, + { + provider_type: CODEX_PROVIDER_TYPE, + display_name: "OpenAI Codex (local CLI)", + base_url: "", + default_models: codexStatusRaw.supported_models ?? [], + supports_streaming: true, + supports_vision: false, + supports_tool_calling: true, + model_list_mode: "curated", + notes: codexStatusRaw.logged_in + ? "Dispatches chat turns through the local Codex CLI." + : "Sign in with `codex login` before chatting.", + } as ProviderRegistryEntry, + ] + : registryRowsRaw; setRegistry(registryRows); setProviderType((current) => { if ( @@ -553,7 +591,7 @@ export function ChatProvidersSettings({ ); return; } - if (!isCustomProvider && !apiKey.trim()) { + if (!providerSkipsApiKey && !apiKey.trim()) { toast.error("Add an API key first."); return; } @@ -625,7 +663,7 @@ export function ChatProvidersSettings({ const displayName = isCustomProvider ? customProviderName.trim() || customProviderDisplayName(providerType) : (selectedRegistryEntry?.display_name ?? providerType); - if (!isCustomProvider && !apiKey.trim()) { + if (!providerSkipsApiKey && !apiKey.trim()) { toast.error("API key is required."); return; } @@ -734,7 +772,9 @@ export function ChatProvidersSettings({ } const isEditingCustomProvider = isCustomProviderType(existing.providerType); - if (!isEditingCustomProvider && !apiKey.trim()) { + const editingProviderSkipsApiKey = + isEditingCustomProvider || isCodexProviderType(existing.providerType); + if (!editingProviderSkipsApiKey && !apiKey.trim()) { toast.error("API key is required."); return; } diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index 8b25f60bf2..e95159e14a 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -33,6 +33,13 @@ export interface ExternalProviderConfig { * OpenAI's hard default is 20. Only meaningful for OpenAI cloud. */ openaiContainerTtlMinutes?: number; + /** + * Codex provider only: number of parallel Codex turns to fan a chat + * request out into. Clamped to [1, 20] by `clampCodexParallelCalls`. + * Omitted or 1 takes the single-call path; values > 1 emit per-tab + * `codex_tab_*` SSE events plus a final `codex_gather` synthesis. + */ + codexParallelCalls?: number; createdAt: number; updatedAt: number; }