feat(pi): Add searchable model picker for new sessions with Databricks Unity AI Gateway OAuth (#4961)

* feat(pi): add searchable start model picker

Signed-off-by: Anthony Ivan <anthony.ivan@databricks.com>

* fix(pi): harden model picker compatibility

Signed-off-by: Anthony Ivan <anthony.ivan@databricks.com>

* refactor(pi): simplify model picker filtering

Signed-off-by: Anthony Ivan <anthony.ivan@databricks.com>

---------

Signed-off-by: Anthony Ivan <anthony.ivan@databricks.com>
Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>
Co-authored-by: Anthony Ivan <anthony.ivan@databricks.com>

omnigent/host/connect.py

@@ -2322,6 +2322,24 @@ class HostProcess:
                 models=models,
             )
 
+        if harness == "pi-native":
+            try:
+                from omnigent.pi_native_credentials import pi_native_model_options
+
+                pi_models = await asyncio.to_thread(pi_native_model_options)
+            except Exception:
+                _logger.exception("Failed to resolve pre-launch Pi model options")
+                return HostModelOptionsResultFrame(
+                    request_id=frame.request_id,
+                    status="failed",
+                    error="failed to resolve Pi model options",
+                )
+            return HostModelOptionsResultFrame(
+                request_id=frame.request_id,
+                status="ok",
+                models=pi_models,
+            )
+
         if harness != "claude-native":
             return HostModelOptionsResultFrame(
                 request_id=frame.request_id,

omnigent/pi_native_credentials.py

@@ -95,6 +95,7 @@ _SURFACE_PROVIDER_IDS: dict[DatabricksPiSurface, str] = {
     DatabricksPiSurface.MLFLOW: _PI_MLFLOW_PROVIDER_ID,
 }
 
+_PI_MANAGED_PROVIDER_IDS = frozenset({_PI_PROVIDER_ID, *_SURFACE_PROVIDER_IDS.values()})
 # Databricks AI Gateway Anthropic Messages surface. Pi speaks this protocol
 # natively (``api: anthropic-messages``); the gateway authenticates with a
 # workspace bearer token, so we set ``authHeader`` (Authorization: Bearer).
@@ -120,6 +121,16 @@ _is_databricks_ai_gateway_url = is_databricks_ai_gateway_url
 _PiModelEntry: TypeAlias = PiModelEntry
 
 
+def _split_pi_native_model_selection(selection: str | None) -> tuple[str, str] | None:
+    """Split an Omnigent-managed ``provider/model`` picker value."""
+    if not selection:
+        return None
+    provider_id, separator, model_id = selection.partition("/")
+    if separator and provider_id in _PI_MANAGED_PROVIDER_IDS and model_id:
+        return provider_id, model_id
+    return None
+
+
 class _PiProviderCompat(TypedDict):
     supportsDeveloperRole: bool
     supportsStore: bool
@@ -356,6 +367,25 @@ class PiProviderConfig:
         )
 
 
+def pi_native_model_options() -> list[dict[str, object]]:
+    """Return pre-launch Pi choices configured through ``omni setup``."""
+    provider = resolve_pi_native_provider()
+    if provider is None:
+        return []
+
+    options: dict[str, dict[str, object]] = {}
+    for provider_id, payload in provider.to_models_config()["providers"].items():
+        for model in payload["models"]:
+            model_id = model["id"]
+            qualified = f"{provider_id}/{model_id}"
+            options[qualified] = {
+                "id": qualified,
+                "model": qualified,
+                "displayName": model.get("name") or model_id,
+            }
+    return [options[model_id] for model_id in sorted(options)]
+
+
 # DATABRICKS-PATCH(pi-live-model-discovery)
 def _default_claude_model_from(entries: list[_PiModelEntry]) -> str | None:
     """Pick pi's launch model from the workspace's live Claude entries.
@@ -979,6 +1009,9 @@ def resolve_pi_native_provider(
     :returns: The resolved provider config, or ``None`` to fall back to Pi's
         own credentials.
     """
+    selection = _split_pi_native_model_selection(model)
+    if selection is not None:
+        _, model = selection
     try:
         config = config_loader()
         # Pi is multi-family; ``omnigent setup`` marks defaults per family, not
@@ -1085,12 +1118,16 @@ def write_pi_models_config(
 
 
 def pi_native_provider_launch(
-    agent_dir: Path, provider: PiProviderConfig
+    agent_dir: Path,
+    provider: PiProviderConfig,
+    *,
+    selection: str | None = None,
 ) -> tuple[dict[str, str], list[str]]:
     """Write the managed config and return the launch env + CLI args for Pi.
 
     :param agent_dir: The managed Pi config dir for this session.
     :param provider: The resolved provider config.
+    :param selection: Optional picker value used to select a generated provider.
     :returns: ``(env, args)`` — the env vars to merge into the terminal spec
         (relocating Pi's config dir) and the ``--provider``/``--model`` args to
         append to the Pi command.
@@ -1098,6 +1135,30 @@ def pi_native_provider_launch(
     # Render once and reuse: rendering logs how an uncataloged model was routed,
     # and this function both writes the config and reads it back for --provider.
     rendered = provider.to_models_config()
+    # Resolve which provider the selected model lives in. Non-Claude models
+    # (GLM, GPT, Llama…) are in secondary providers; Claude models are in the
+    # primary provider. Read the rendered config so family fallbacks agree.
+    selected_model = provider.model
+    model_provider_id = provider.provider_id
+    selection_parts = _split_pi_native_model_selection(selection)
+    if selection_parts is not None:
+        candidate_provider, candidate_model = selection_parts
+        configured = rendered["providers"].get(candidate_provider)
+        if not configured or not any(
+            model.get("id") == candidate_model for model in configured.get("models", [])
+        ):
+            raise ValueError(
+                f"Pi model selection {selection!r} is not available in managed configuration"
+            )
+        model_provider_id = candidate_provider
+        selected_model = candidate_model
+    else:
+        for extra_id, extra_cfg in rendered["providers"].items():
+            if extra_id == provider.provider_id:
+                continue
+            if any(m.get("id") == provider.model for m in extra_cfg.get("models", [])):
+                model_provider_id = extra_id
+                break
     write_pi_models_config(agent_dir, provider, rendered)
     # Copy the user's global Pi settings but suppress defaultThinkingLevel.
     # In TUI mode Pi applies the setting from ~/.pi/agent/settings.json; for
@@ -1111,25 +1172,13 @@ def pi_native_provider_launch(
 
     prepare_managed_pi_agent_dir(agent_dir, overlay={"defaultThinkingLevel": None})
     env = {PI_CODING_AGENT_DIR_ENV_VAR: str(agent_dir)}
-    # Resolve which provider the selected model lives in. Non-Claude models
-    # (GLM, GPT, Llama…) are in secondary providers (omnigent-openai); Claude
-    # models are in the primary provider (omnigent). Read the *rendered* config
-    # rather than additional_providers so a model routed by the family fallback
-    # gets the same --provider that models.json registered it under.
-    model_provider_id = provider.provider_id
-    for extra_id, extra_cfg in rendered["providers"].items():
-        if extra_id == provider.provider_id:
-            continue
-        if any(m.get("id") == provider.model for m in extra_cfg.get("models", [])):
-            model_provider_id = extra_id
-            break
     # When the model id contains a "/" Pi's arg parser splits on the first
     # slash and treats the left part as a provider name, overriding
     # --provider. Pass the fully-qualified "provider/model" reference so Pi's
     # findExactModelReferenceMatch matches the canonical form exactly and
     # routes to our custom provider, not a builtin with the same model id.
     model_arg = (
-        f"{model_provider_id}/{provider.model}" if "/" in provider.model else provider.model
+        f"{model_provider_id}/{selected_model}" if "/" in selected_model else selected_model
     )
     args = ["--provider", model_provider_id, "--model", model_arg]
     # For non-Claude models on openai-completions/responses, disable thinking.

omnigent/resources/pi_native/omnigent_pi_native_extension.js

@@ -357,7 +357,11 @@ function piResultFromMcpResponse(json) {
   // masquerading as a successful tool result. callOmnigentTool detects and
   // resolves the ASK round-trip BEFORE calling this; treat a stray one as a
   // fail-closed error so an unresolved approval never reports success.
-  if (result && typeof result === "object" && result.resultType === "input_required") {
+  if (
+    result &&
+    typeof result === "object" &&
+    result.resultType === "input_required"
+  ) {
     return {
       content: [
         {
@@ -371,7 +375,11 @@ function piResultFromMcpResponse(json) {
   if (result && Array.isArray(result.content)) {
     const parts = [];
     for (const block of result.content) {
-      if (block && typeof block === "object" && typeof block.text === "string") {
+      if (
+        block &&
+        typeof block === "object" &&
+        typeof block.text === "string"
+      ) {
         parts.push(block.text);
       }
     }
@@ -400,7 +408,11 @@ function piResultFromMcpResponse(json) {
  */
 function mcpInputRequired(json) {
   const result = json && typeof json === "object" ? json.result : undefined;
-  if (!result || typeof result !== "object" || result.resultType !== "input_required") {
+  if (
+    !result ||
+    typeof result !== "object" ||
+    result.resultType !== "input_required"
+  ) {
     return null;
   }
   const inputRequests =
@@ -894,7 +906,27 @@ async function applyModelChange(pi, config, ctx, modelId) {
   }
   let model;
   try {
-    model = listModels().find((m) => m && m.id === id);
+    const models = listModels();
+    const separator = id.indexOf("/");
+    if (separator > 0) {
+      const provider = id.slice(0, separator);
+      const bareId = id.slice(separator + 1);
+      model = models.find(
+        (candidate) =>
+          candidate && candidate.id === bareId && candidate.provider === provider,
+      );
+      if (!model) model = models.find((candidate) => candidate && candidate.id === id);
+      if (!model && registry && typeof registry.find === "function") {
+        model = registry.find(provider, bareId);
+      }
+      if (!model) {
+        model = models.find(
+          (candidate) => candidate && candidate.id === bareId && !candidate.provider,
+        );
+      }
+    } else {
+      model = models.find((candidate) => candidate && candidate.id === id);
+    }
   } catch (_err) {
     model = undefined;
   }
@@ -940,6 +972,13 @@ async function postModelChangeError(config, message) {
   });
 }
 
+function modelReference(model) {
+  const modelId = model && typeof model.id === "string" ? model.id : "";
+  if (!modelId) return "";
+  const provider = model && typeof model.provider === "string" ? model.provider : "";
+  return provider ? `${provider}/${modelId}` : modelId;
+}
+
 /**
  * Report Pi's live model catalog to Omnigent for the Web UI model picker.
  *
@@ -977,11 +1016,13 @@ async function postModelOptions(config, ctx) {
   const options = [];
   const seen = new Set();
   for (const model of models) {
-    const id = model && typeof model.id === "string" ? model.id : "";
+    const modelId = model && typeof model.id === "string" ? model.id : "";
+    const id = modelReference(model);
     if (!id || seen.has(id)) continue;
     seen.add(id);
-    const name = model && typeof model.name === "string" && model.name ? model.name : id;
-    options.push({ id, displayName: name });
+    const name =
+      model && typeof model.name === "string" && model.name ? model.name : modelId;
+    options.push({ id, model: id, displayName: name });
   }
   if (options.length === 0) return;
   await postEvent(config, {
@@ -990,7 +1031,13 @@ async function postModelOptions(config, ctx) {
   });
 }
 
-function startInboxPoller(pi, config, handleInterrupt, handleCompact, handleModelChange) {
+function startInboxPoller(
+  pi,
+  config,
+  handleInterrupt,
+  handleCompact,
+  handleModelChange,
+) {
   if (!config || !config.inboxDir || pi.__omnigentInboxPoller) return;
   // Bound the dedup set (FIFO eviction) — delivered files are unlinked, so a
   // long-lived TUI mustn't grow it unboundedly.
@@ -1408,7 +1455,8 @@ module.exports = function (pi) {
   // carries no identity field at all.
   function usageMessageKey(message, usage) {
     if (message && typeof message === "object") {
-      if (typeof message.id === "string" && message.id) return `id:${message.id}`;
+      if (typeof message.id === "string" && message.id)
+        return `id:${message.id}`;
       if (typeof message.responseId === "string" && message.responseId)
         return `rid:${message.responseId}`;
       if (typeof message.timestamp === "number")
@@ -1721,12 +1769,11 @@ module.exports = function (pi) {
     // ``/login`` session (no Omnigent ``model_override``, no ``llm_model``)
     // shows no active model until the user switches. Mirrors the
     // ``model_select`` handler, but for the startup value ``ctx.model``.
-    const startupModelId =
-      ctx && ctx.model && typeof ctx.model.id === "string" ? ctx.model.id : "";
-    if (startupModelId) {
+    const startupModel = modelReference(ctx ? ctx.model : undefined);
+    if (startupModel) {
       await postEvent(config, {
         type: "external_model_change",
-        data: { model: startupModelId },
+        data: { model: startupModel },
       });
     }
     await postEvent(config, {
@@ -1749,14 +1796,15 @@ module.exports = function (pi) {
     // web-side override. The server dedups against ``model_override``, so a
     // web-initiated switch (which already persisted the value before queuing
     // the inbox ``model_change``) round-trips here as a no-op.
-    const source = event && typeof event.source === "string" ? event.source : "";
+    const source =
+      event && typeof event.source === "string" ? event.source : "";
     if (source === "restore") return;
     const model = event && event.model ? event.model : undefined;
-    const modelId = model && typeof model.id === "string" ? model.id : "";
-    if (!modelId) return;
+    const selectedModel = modelReference(model);
+    if (!selectedModel) return;
     await postEvent(config, {
       type: "external_model_change",
-      data: { model: modelId },
+      data: { model: selectedModel },
     });
   });
 
@@ -1814,7 +1862,8 @@ module.exports = function (pi) {
     if (changed) await postSessionUsage();
     // Reuse the agent_start response_id so the web client matches the idle
     // edge and clears the "streaming" status, unblocking queued follow-ups.
-    const endResponseId = turnStatusResponseId ?? `pi-${Date.now()}-${++sequence}`;
+    const endResponseId =
+      turnStatusResponseId ?? `pi-${Date.now()}-${++sequence}`;
     turnStatusResponseId = null;
     await postEvent(config, {
       type: "external_session_status",
@@ -1965,9 +2014,13 @@ module.exports = function (pi) {
     // unsupported API types) as visible error items in the web UI so users
     // aren't left staring at an empty turn.
     const stopReason =
-      message && typeof message.stopReason === "string" ? message.stopReason : "";
+      message && typeof message.stopReason === "string"
+        ? message.stopReason
+        : "";
     const errorMessage =
-      message && typeof message.errorMessage === "string" ? message.errorMessage : "";
+      message && typeof message.errorMessage === "string"
+        ? message.errorMessage
+        : "";
     if (stopReason === "error" && errorMessage) {
       await postEvent(config, {
         type: "external_conversation_item",

omnigent/runner/native/orchestration.py

@@ -2158,17 +2158,16 @@ async def _auto_create_pi_terminal(
             resolve_pi_native_provider,
         )
 
-        # Thread the agent spec's pinned model (``executor.model``) into the
-        # resolved provider so the generated ``models.json`` — and the
-        # appended ``--model`` arg (see ``pi_native_provider_launch``) — select
-        # it, reaching parity with claude-native / cursor-native. ``None``
-        # (no model declared) keeps the provider's default model.
-        # model_override (set by /model or sys_session_create's model arg)
-        # takes precedence over the spec's pinned executor.model.
+        # Provider-qualified picker values select one of the models rendered
+        # from the provider configured through ``omni setup``.
         spec_model = launch_config.model_override or _pi_native_model_from_spec(agent_spec)
         provider = resolve_pi_native_provider(model=spec_model)
         if provider is not None:
-            cred_env, cred_args = pi_native_provider_launch(bridge_dir / "pi-agent", provider)
+            cred_env, cred_args = pi_native_provider_launch(
+                bridge_dir / "pi-agent",
+                provider,
+                selection=spec_model,
+            )
             pi_env.update(cred_env)
             pi_args.extend(cred_args)
             # An unroutable model leaves Pi unable to select it, which looks

tests/host/test_connect.py

@@ -123,6 +123,42 @@ async def test_handle_model_options_uses_host_claude_configuration(
     )
 
 
+async def test_handle_model_options_uses_host_pi_configuration(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Pi's launch picker uses only models configured through Omnigent."""
+    from omnigent import pi_native_credentials
+
+    monkeypatch.setattr(
+        pi_native_credentials,
+        "pi_native_model_options",
+        lambda: [
+            {
+                "id": "omnigent-openai/system.ai.gpt-5-6-sol",
+                "model": "omnigent-openai/system.ai.gpt-5-6-sol",
+                "displayName": "omnigent-openai/GPT 5.6 Sol",
+            }
+        ],
+    )
+    host = _make_host_process()
+
+    result = await host._handle_model_options(
+        HostModelOptionsFrame(request_id="req_pi_models", harness="pi-native"),
+    )
+
+    assert result == HostModelOptionsResultFrame(
+        request_id="req_pi_models",
+        status="ok",
+        models=[
+            {
+                "id": "omnigent-openai/system.ai.gpt-5-6-sol",
+                "model": "omnigent-openai/system.ai.gpt-5-6-sol",
+                "displayName": "omnigent-openai/GPT 5.6 Sol",
+            }
+        ],
+    )
+
+
 async def test_handle_model_options_uses_codex_provider_catalog(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:

tests/runner/test_app_sessions_native_terminals_autocreate.py

@@ -257,7 +257,7 @@ async def test_auto_create_pi_terminal_surfaces_credential_warning(
     monkeypatch.setattr(
         pi_native_credentials,
         "pi_native_provider_launch",
-        lambda _agent_dir, _provider: ({}, []),
+        lambda _agent_dir, _provider, **_kwargs: ({}, []),
     )
 
     async def _fake_launch_config(**_kwargs: Any) -> _PiNativeLaunchConfig:

tests/test_pi_native_credentials.py

@@ -154,6 +154,29 @@ def test_key_provider_resolves_to_inline_family() -> None:
     assert provider.model == "claude-sonnet-4-6"
 
 
+def test_managed_picker_prefix_is_not_part_of_provider_model() -> None:
+    """A managed picker value resolves its provider-local model id."""
+    config = {
+        "providers": {
+            "anthropic": {
+                "kind": "key",
+                "default": True,
+                "anthropic": {
+                    "base_url": "https://api.anthropic.com",
+                    "api_key": "sk-test-literal",
+                },
+            }
+        }
+    }
+
+    provider = creds.resolve_pi_native_provider(
+        model="omnigent/claude-opus-4-7", config_loader=lambda: config
+    )
+
+    assert provider is not None
+    assert provider.model == "claude-opus-4-7"
+
+
 def test_subscription_default_returns_none() -> None:
     """A subscription (CLI-login) default isn't reusable by Pi → None."""
     config = {"providers": {"claude": {"kind": "subscription", "default": True, "cli": "claude"}}}
@@ -283,6 +306,99 @@ def test_pi_native_provider_launch_namespaced_model_uses_qualified_arg(
     assert args == ["--provider", "omnigent", "--model", "omnigent/moonshotai/kimi-k2.5"]
 
 
+def test_provider_launch_accepts_provider_qualified_selection(tmp_path: Path) -> None:
+    """A start-picker selection chooses its generated Pi provider and model."""
+    provider = creds.PiProviderConfig(
+        provider_id="omnigent",
+        base_url="https://api.anthropic.com",
+        api="anthropic-messages",
+        model="claude-sonnet-4-6",
+        api_key="sk-secret",
+        auth_header=False,
+        additional_providers={
+            "omnigent-openai": {
+                "baseUrl": "https://api.openai.com/v1",
+                "api": "openai-responses",
+                "apiKey": "sk-openai",
+                "models": [{"id": "gpt-5.6-sol"}],
+            }
+        },
+    )
+
+    _, args = creds.pi_native_provider_launch(
+        tmp_path / "pi-agent",
+        provider,
+        selection="omnigent-openai/gpt-5.6-sol",
+    )
+
+    assert args == [
+        "--provider",
+        "omnigent-openai",
+        "--model",
+        "gpt-5.6-sol",
+        "--thinking",
+        "off",
+    ]
+
+
+def test_provider_launch_rejects_unavailable_qualified_selection(tmp_path: Path) -> None:
+    """A stale picker value must not silently launch the provider default."""
+    provider = creds.PiProviderConfig(
+        provider_id="omnigent",
+        base_url="https://api.anthropic.com",
+        api="anthropic-messages",
+        model="claude-sonnet-4-6",
+        api_key="sk-secret",
+        auth_header=False,
+    )
+    agent_dir = tmp_path / "pi-agent"
+
+    with pytest.raises(ValueError, match="not available"):
+        creds.pi_native_provider_launch(
+            agent_dir,
+            provider,
+            selection="omnigent-openai/gpt-missing",
+        )
+
+    assert not agent_dir.exists()
+
+
+def test_pi_native_model_options_lists_only_managed_models(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Pre-launch choices come only from the provider built by ``omni setup``."""
+    provider = creds.PiProviderConfig(
+        provider_id="omnigent",
+        base_url="https://api.anthropic.com",
+        api="anthropic-messages",
+        model="claude-sonnet-4-6",
+        api_key="sk-secret",
+        auth_header=False,
+        additional_providers={
+            "omnigent-openai": {
+                "baseUrl": "https://api.openai.com/v1",
+                "api": "openai-responses",
+                "apiKey": "sk-openai",
+                "models": [{"id": "gpt-5.6-sol", "name": "GPT 5.6 Sol"}],
+            }
+        },
+    )
+    monkeypatch.setattr(creds, "resolve_pi_native_provider", lambda: provider)
+
+    assert creds.pi_native_model_options() == [
+        {
+            "id": "omnigent-openai/gpt-5.6-sol",
+            "model": "omnigent-openai/gpt-5.6-sol",
+            "displayName": "GPT 5.6 Sol",
+        },
+        {
+            "id": "omnigent/claude-sonnet-4-6",
+            "model": "omnigent/claude-sonnet-4-6",
+            "displayName": "claude-sonnet-4-6",
+        },
+    ]
+
+
 def test_openai_chat_wire_api_resolves_to_completions(monkeypatch: pytest.MonkeyPatch) -> None:
     """An OpenAI family with wire_api: chat → openai-completions API.
 

tests/test_pi_native_extension.py

@@ -2668,9 +2668,9 @@ const setModelCalls = [];
 // The catalog Pi's modelRegistry exposes; setModel returns false for a model
 // with no configured API key (mirrors Pi's real contract).
 const catalog = [
-  { id: "databricks-claude-sonnet-4-6", name: "Sonnet", hasKey: true },
-  { id: "databricks-claude-opus-4-1", name: "Opus", hasKey: true },
-  { id: "no-key-model", name: "NoKey", hasKey: false },
+  { provider: "omnigent", id: "databricks-claude-sonnet-4-6", name: "Sonnet", hasKey: true },
+  { provider: "omnigent", id: "databricks-claude-opus-4-1", name: "Opus", hasKey: true },
+  { provider: "omnigent", id: "no-key-model", name: "NoKey", hasKey: false },
 ];
 const pi = {
   registerCommand() {},
@@ -2692,10 +2692,11 @@ const ctx = {
   // external_model_change). ``getAvailable`` returns only auth-configured
   // models (what the picker should show); ``getAll`` is Pi's full built-in
   // catalog (the fallback for older Pi).
-  model: { id: "databricks-claude-sonnet-4-6", name: "Sonnet" },
+  model: { provider: "omnigent", id: "databricks-claude-sonnet-4-6", name: "Sonnet" },
   modelRegistry: {
     getAll: () => catalog,
     getAvailable: () => catalog.filter((m) => m.hasKey),
+    find: (provider, id) => catalog.find((m) => m.provider === provider && m.id === id),
   },
 };
 
@@ -2747,7 +2748,8 @@ def test_inbox_model_change_applies_via_set_model(tmp_path: Path) -> None:
         + r"""
 (async () => {
   await handlers.session_start({}, ctx); // starts the inbox poller
-  await deliverModelChange("databricks-claude-opus-4-1");
+  delete ctx.modelRegistry.find;
+  await deliverModelChange("omnigent/databricks-claude-opus-4-1");
 
   assert.equal(setModelCalls.length, 1, JSON.stringify(setModelCalls));
   assert.equal(setModelCalls[0].id, "databricks-claude-opus-4-1");
@@ -2814,23 +2816,23 @@ def test_model_select_mirrors_to_external_model_change(tmp_path: Path) -> None:
   // resolves from the start; ignore that when checking the user switch.
   const startupChanges = posted.filter((e) => e.type === "external_model_change");
   assert.equal(startupChanges.length, 1, JSON.stringify(posted));
-  assert.equal(startupChanges[0].data.model, "databricks-claude-sonnet-4-6");
+  assert.equal(startupChanges[0].data.model, "omnigent/databricks-claude-sonnet-4-6");
 
   // A genuine user switch mirrors back.
   await handlers.model_select(
-    { source: "set", model: { id: "databricks-claude-opus-4-1" } },
+    { source: "set", model: { provider: "omnigent", id: "databricks-claude-opus-4-1" } },
     ctx,
   );
   // A startup restore must be ignored (could clobber a pending web override).
   await handlers.model_select(
-    { source: "restore", model: { id: "databricks-claude-sonnet-4-6" } },
+    { source: "restore", model: { provider: "omnigent", id: "databricks-claude-sonnet-4-6" } },
     ctx,
   );
 
   const changes = posted.filter((e) => e.type === "external_model_change");
   // Two total: the startup mirror + the one user switch (restore ignored).
   assert.equal(changes.length, 2, JSON.stringify(posted));
-  assert.equal(changes[1].data.model, "databricks-claude-opus-4-1");
+  assert.equal(changes[1].data.model, "omnigent/databricks-claude-opus-4-1");
   finish();
 })().catch((error) => {
   finish();
@@ -2870,7 +2872,10 @@ def test_session_start_posts_model_options_from_registry(tmp_path: Path) -> None
   // getAvailable() filters out ``no-key-model`` (no configured auth).
   assert.deepEqual(
     models.map((m) => m.id),
-    ["databricks-claude-sonnet-4-6", "databricks-claude-opus-4-1"],
+    [
+      "omnigent/databricks-claude-sonnet-4-6",
+      "omnigent/databricks-claude-opus-4-1",
+    ],
     JSON.stringify(models),
   );
   // Display name falls back to the model's ``name``.
@@ -2879,7 +2884,55 @@ def test_session_start_posts_model_options_from_registry(tmp_path: Path) -> None
   // The launch model is mirrored so the pill/active-row resolve immediately.
   const changes = posted.filter((e) => e.type === "external_model_change");
   assert.equal(changes.length, 1, JSON.stringify(posted));
-  assert.equal(changes[0].data.model, "databricks-claude-sonnet-4-6");
+  assert.equal(changes[0].data.model, "omnigent/databricks-claude-sonnet-4-6");
+  finish();
+})().catch((error) => {
+  finish();
+  console.error(error && error.stack ? error.stack : error);
+  process.exit(1);
+});
+"""
+    )
+    _run_extension_script(node, _extension_path(), script)
+
+
+def test_models_without_provider_keep_bare_id_behavior(tmp_path: Path) -> None:
+    """Older Pi model objects without ``provider`` still populate and mirror."""
+    node = shutil.which("node")
+    if node is None:
+        pytest.skip("node is required for the pi-native extension e2e test")
+
+    script = (
+        _MODEL_SWITCH_HARNESS
+        + r"""
+(async () => {
+  const legacyCatalog = catalog.map(({ provider: _provider, ...model }) => model);
+  const legacyCtx = {
+    ...ctx,
+    model: { id: "databricks-claude-sonnet-4-6", name: "Sonnet" },
+    modelRegistry: {
+      getAll: () => legacyCatalog,
+      getAvailable: () => legacyCatalog.filter((model) => model.hasKey),
+    },
+  };
+
+  await handlers.session_start({}, legacyCtx);
+  const opts = posted.filter((event) => event.type === "external_model_options");
+  assert.deepEqual(
+    opts[0].data.models.map((model) => model.id),
+    ["databricks-claude-sonnet-4-6", "databricks-claude-opus-4-1"],
+  );
+  assert.equal(opts[0].data.models[0].displayName, "Sonnet");
+
+  await handlers.model_select(
+    { source: "set", model: { id: "databricks-claude-opus-4-1" } },
+    legacyCtx,
+  );
+  const changes = posted.filter((event) => event.type === "external_model_change");
+  assert.deepEqual(
+    changes.map((event) => event.data.model),
+    ["databricks-claude-sonnet-4-6", "databricks-claude-opus-4-1"],
+  );
   finish();
 })().catch((error) => {
   finish();

web/src/components/HarnessConfigControls.tsx

@@ -127,10 +127,12 @@ export function ConfigRow({
   label,
   description,
   children,
+  controlClassName,
 }: {
   label: string;
   description?: string;
   children: ReactNode;
+  controlClassName?: string;
 }) {
   return (
     // Stacked on mobile (label above a full-width control) so the label never
@@ -141,7 +143,7 @@ export function ConfigRow({
         <div className="text-ui font-medium">{label}</div>
         {description && <div className="text-sm text-muted-foreground">{description}</div>}
       </div>
-      <div className="w-full sm:w-52 sm:shrink-0">{children}</div>
+      <div className={cn("w-full sm:w-52 sm:shrink-0", controlClassName)}>{children}</div>
     </div>
   );
 }

web/src/lib/nativeCodingAgents.ts

@@ -17,7 +17,7 @@ export type NativeCodingAgentIconKind =
   | "kimi"
   | "hermes";
 export type NativeCodingAgentCapability =
-  "permissionMode" | "approvalMode" | "cursorMode" | "skipPermissions";
+  "permissionMode" | "approvalMode" | "cursorMode" | "skipPermissions" | "modelPicker";
 
 export interface NativeCodingAgentSpec {
   key: NativeCodingAgentIconKind;
@@ -53,7 +53,7 @@ export const NATIVE_CODING_AGENTS = [
     displayName: "Claude Code",
     iconKind: "claude",
     sortRank: 10,
-    capabilities: ["permissionMode"],
+    capabilities: ["permissionMode", "modelPicker"],
     fullySupported: true,
   },
   {
@@ -106,6 +106,7 @@ export const NATIVE_CODING_AGENTS = [
     displayName: "Pi",
     iconKind: "pi",
     sortRank: 40,
+    capabilities: ["modelPicker"],
   },
   {
     key: "kiro",

web/src/shell/NewChatDialog.flow.test.tsx

@@ -9,7 +9,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 
 import { authenticatedFetch } from "@/lib/identity";
 import type { Host } from "@/hooks/useHosts";
-import { useHosts } from "@/hooks/useHosts";
+import { useHostModelOptions, useHosts } from "@/hooks/useHosts";
 import type { AvailableAgent } from "@/hooks/useAvailableAgents";
 import { useAvailableAgents } from "@/hooks/useAvailableAgents";
 import { NewChatLandingScreen, resetLandingDraft, sanitizeInitialPrompt } from "./NewChatDialog";
@@ -265,6 +265,14 @@ beforeEach(() => {
   // left behind by an unmounting test doesn't seed the next one.
   resetLandingDraft();
   localStorage.clear();
+  vi.mocked(useHostModelOptions).mockReturnValue({
+    data: [
+      { id: "opus", displayName: "Opus" },
+      { id: "sonnet", displayName: "Sonnet" },
+      { id: "haiku", displayName: "Haiku" },
+    ],
+    isLoading: false,
+  } as unknown as ReturnType<typeof useHostModelOptions>);
   // Seed host_1's recent so the working directory pre-fills deterministically
   // (the create body must carry SEEDED_WORKSPACE through).
   localStorage.setItem(RECENT_KEY, JSON.stringify({ host_1: [SEEDED_WORKSPACE] }));
@@ -1047,6 +1055,62 @@ describe("NewChatLandingScreen create flow", () => {
     expect(body.reasoning_effort).toBe("high");
   });
 
+  it("rides an omni-setup model along to create for pi-native", async () => {
+    setAgents([
+      agent({
+        id: "ag_pi",
+        name: "pi-native-ui",
+        display_name: "Pi",
+        harness: "pi-native",
+      }),
+    ]);
+    vi.mocked(useHostModelOptions).mockReturnValue({
+      data: [
+        {
+          id: "omnigent-openai/system.ai.gpt-5-6-sol",
+          model: "omnigent-openai/system.ai.gpt-5-6-sol",
+          displayName: "GPT 5.6 Sol",
+        },
+        {
+          id: "omnigent/databricks-claude-sonnet-4-6",
+          model: "omnigent/databricks-claude-sonnet-4-6",
+          displayName: "Claude Sonnet 4.6",
+        },
+      ],
+      isLoading: false,
+    } as unknown as ReturnType<typeof useHostModelOptions>);
+    vi.mocked(authenticatedFetch).mockResolvedValueOnce({
+      ok: true,
+      json: async () => ({ id: "conv_pi" }),
+    } as unknown as Response);
+
+    renderLanding();
+    await waitForWorkspaceSeed();
+    openAgentConfig("ag_pi");
+    fireEvent.click(screen.getByTestId("new-chat-landing-config-model"));
+    const fullNameRow = document.querySelector(
+      '[data-model-id="omnigent-openai/system.ai.gpt-5-6-sol"]',
+    );
+    expect(fullNameRow).not.toBeNull();
+    expect(fullNameRow).toHaveAttribute("title", "GPT 5.6 Sol");
+    fireEvent.change(screen.getByTestId("new-chat-landing-config-model-search"), {
+      target: { value: "gpt sol" },
+    });
+    expect(screen.getByText("GPT 5.6 Sol")).toBeInTheDocument();
+    expect(screen.queryByText("Claude Sonnet 4.6")).toBeNull();
+    fireEvent.click(screen.getByText("GPT 5.6 Sol"));
+    saveConfig();
+    typeMessage("go");
+    fireEvent.click(screen.getByTestId("new-chat-landing-submit"));
+
+    await waitFor(() => expect(authenticatedFetch).toHaveBeenCalledTimes(1));
+    const [, init] = vi.mocked(authenticatedFetch).mock.calls[0] as [string, RequestInit];
+    const body = JSON.parse(init.body as string);
+    expect(body.model_override).toBe("omnigent-openai/system.ai.gpt-5-6-sol");
+    expect(body.reasoning_effort).toBeUndefined();
+    expect(body.labels?.["omnigent.wrapper"]).toBe("pi-native-ui");
+  });
+
   it("seeds the model + effort from the last pick for claude-native on a new session", async () => {
     // A returning user's last model/effort pick for this harness is on record;
     // the new session must auto-fill it and post it WITHOUT re-opening the

web/src/shell/NewChatDialog.tsx

@@ -46,6 +46,13 @@ import {
 } from "@/components/ui/select";
 import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
+import {
+  Command,
+  CommandEmpty,
+  CommandInput,
+  CommandItem,
+  CommandList,
+} from "@/components/ui/command";
 import {
   CLAUDE_NATIVE_EFFORTS,
   ConfigRow,
@@ -1370,9 +1377,93 @@ export function AgentHarnessPicker({
   );
 }
 
+function SearchableModelPicker({
+  value,
+  options,
+  loading,
+  onValueChange,
+}: {
+  value: string;
+  options: readonly { id: string; displayName: string }[];
+  loading: boolean;
+  onValueChange: (value: string) => void;
+}) {
+  const [open, setOpen] = useState(false);
+  const selectedLabel =
+    value === MODEL_SELECT_DEFAULT
+      ? "Default"
+      : (options.find((option) => option.id === value)?.displayName ?? value);
+  const select = (nextValue: string) => {
+    onValueChange(nextValue);
+    setOpen(false);
+  };
+
+  return (
+    <Popover open={open} onOpenChange={setOpen}>
+      <PopoverTrigger asChild>
+        <Button
+          type="button"
+          variant="outline"
+          role="combobox"
+          aria-expanded={open}
+          aria-label="Model"
+          className="h-8 w-full justify-between gap-2 px-2.5 font-normal"
+          data-testid="new-chat-landing-config-model"
+        >
+          <span className="min-w-0 truncate">{selectedLabel}</span>
+          <ChevronDownIcon className="size-4 shrink-0 text-muted-foreground" />
+        </Button>
+      </PopoverTrigger>
+      <PopoverContent
+        align="start"
+        className="max-h-[var(--radix-popover-content-available-height)] w-[var(--radix-popover-trigger-width)] overflow-hidden p-0"
+      >
+        <Command className="h-auto min-h-0">
+          <CommandInput
+            placeholder="Search models…"
+            data-testid="new-chat-landing-config-model-search"
+          />
+          <CommandList
+            className="max-h-72 min-h-0 overflow-y-auto overscroll-contain"
+            onWheel={(event) => event.stopPropagation()}
+          >
+            <CommandItem
+              value={MODEL_SELECT_DEFAULT}
+              data-checked={value === MODEL_SELECT_DEFAULT}
+              onSelect={() => select(MODEL_SELECT_DEFAULT)}
+            >
+              Default
+            </CommandItem>
+            {options.map((option) => (
+              <CommandItem
+                key={option.id}
+                value={option.id}
+                keywords={[option.displayName]}
+                title={option.displayName}
+                data-model-id={option.id}
+                data-checked={value === option.id}
+                onSelect={() => select(option.id)}
+              >
+                <span className="min-w-0 truncate">{option.displayName}</span>
+              </CommandItem>
+            ))}
+            {!loading && <CommandEmpty>No models found</CommandEmpty>}
+            {loading && (
+              <div className="px-2 py-3 text-center text-xs text-muted-foreground">
+                Loading models…
+              </div>
+            )}
+          </CommandList>
+        </Command>
+      </PopoverContent>
+    </Popover>
+  );
+}
+
 /**
  * Harness-configuration modal opened from the composer's gear icon. Shows the
  * selected agent's run-config knobs — Claude: model / effort / permissions;
+ * Pi: model;
  * Codex/OpenCode: approval mode (+ Codex's dangerous full-bypass opt-in);
  * Cursor: exec mode; bundle agents: brain-harness override. On the fully-auto
  * harness the router owns harness and model, so every harness-specific knob
@@ -1402,6 +1493,8 @@ function HarnessConfigModal({
   claudeModelsLoading,
   codexModelOptions,
   codexModelsLoading,
+  piModelOptions,
+  piModelsLoading,
   pickedEffort,
   pickedHarness,
   costControlMode,
@@ -1432,6 +1525,8 @@ function HarnessConfigModal({
   claudeModelsLoading: boolean;
   codexModelOptions: readonly Pick<NativeModelOption, "id" | "displayName" | "isDefault">[];
   codexModelsLoading: boolean;
+  piModelOptions: readonly { id: string; displayName: string }[];
+  piModelsLoading: boolean;
   pickedEffort: string;
   pickedHarness: string | null;
   costControlMode: CostControlMode;
@@ -1453,9 +1548,8 @@ function HarnessConfigModal({
   const hasApproval = nativeAgentHasCapability(agent, "approvalMode");
   const hasCursor = nativeAgentHasCapability(agent, "cursorMode");
   const hasAgySkip = nativeAgentHasCapability(agent, "skipPermissions");
+  const hasModelPicker = nativeAgentHasCapability(agent, "modelPicker");
   const isCodex = entryHarness === "codex-native";
-  const modelOptions = isCodex ? codexModelOptions : claudeModelOptions;
-  const modelsLoading = isCodex ? codexModelsLoading : claudeModelsLoading;
   const brainDefault =
     agent.harness != null && agent.harness in brainHarnessLabels ? agent.harness : null;
 
@@ -1553,6 +1647,9 @@ function HarnessConfigModal({
           effort: draftEffort,
           mode: draftPermission,
         });
+    } else if (hasModelPicker) {
+      setPickedModel(draftModel);
+      if (entryHarness) writeHarnessOption(entryHarness, { model: draftModel });
     } else if (hasApproval) {
       if (isCodex) setPickedModel(draftModel);
       setApprovalMode(draftApproval);
@@ -1607,7 +1704,10 @@ function HarnessConfigModal({
 
   return (
     <Dialog open={open} onOpenChange={onOpenChange}>
-      <DialogContent className="sm:max-w-md" data-testid="new-chat-landing-config-modal">
+      <DialogContent
+        className={cn("sm:max-w-md", entryHarness === "pi-native" && "sm:max-w-xl")}
+        data-testid="new-chat-landing-config-modal"
+      >
         <DialogHeader>
           <DialogTitle>Configure {configTitleName}</DialogTitle>
           <DialogDescription className="sr-only">
@@ -1616,6 +1716,17 @@ function HarnessConfigModal({
         </DialogHeader>
 
         <div className="flex flex-col gap-5 py-1">
+          {!autoRouting && hasModelPicker && !hasPermission && (
+            <ConfigRow label="Model" description="Underlying LLM" controlClassName="sm:w-80">
+              <SearchableModelPicker
+                value={modelValue}
+                options={piModelOptions}
+                loading={piModelsLoading}
+                onValueChange={onModelChange}
+              />
+            </ConfigRow>
+          )}
+
           {!autoRouting && hasPermission && (
             <>
               <ConfigRow label="Model" description="Underlying LLM">
@@ -1698,10 +1809,10 @@ function HarnessConfigModal({
                   defaultLabel={defaultModelLabel(codexModelOptions, displayModelId)}
                   contentClassName="[&_[data-slot=select-item]]:pl-2.5"
                 >
-                  {modelsLoading && (
+                  {codexModelsLoading && (
                     <div className="px-2.5 py-1 text-sm text-muted-foreground">Loading models…</div>
                   )}
-                  {!modelsLoading && modelOptions.length === 0 && (
+                  {!codexModelsLoading && codexModelOptions.length === 0 && (
                     <div className="px-2.5 py-1 text-sm text-muted-foreground">
                       Models unavailable
                     </div>
@@ -2107,6 +2218,11 @@ export function NewChatLandingScreen() {
     "codex-native",
     !sandboxSelected,
   );
+  const { data: hostPiModelOptions, isLoading: hostPiModelsLoading } = useHostModelOptions(
+    selectedHostId,
+    "pi-native",
+    !sandboxSelected,
+  );
   const claudeModelOptions = useMemo(
     () =>
       sandboxSelected
@@ -2124,6 +2240,16 @@ export function NewChatLandingScreen() {
     () => (sandboxSelected ? [] : (hostCodexModelOptions ?? [])),
     [hostCodexModelOptions, sandboxSelected],
   );
+  const piModelOptions = useMemo(
+    () =>
+      sandboxSelected
+        ? []
+        : (hostPiModelOptions ?? []).map((option) => ({
+            id: option.id,
+            displayName: option.displayName ?? option.id,
+          })),
+    [hostPiModelOptions, sandboxSelected],
+  );
   // Desktop-shell host status for THIS machine (null outside Electron), so the
   // picker can tag the current machine and offer to auto-connect it.
   const [desktopHost, setDesktopHost] = useState<HostIdentity | null>(null);
@@ -2659,15 +2785,16 @@ export function NewChatLandingScreen() {
         : agentList.find((a) => a.id === effectiveAgentId),
     [agentList, effectiveAgentId, pendingAgent],
   );
+  const selectedNativeHarness = nativeCodingAgentForAvailableAgent(selectedAgent)?.harness ?? null;
   const supportsPermissionMode = nativeAgentHasCapability(selectedAgent, "permissionMode");
   const supportsApprovalMode = nativeAgentHasCapability(selectedAgent, "approvalMode");
   const supportsCursorMode = nativeAgentHasCapability(selectedAgent, "cursorMode");
   const supportsAgySkipPermissions = nativeAgentHasCapability(selectedAgent, "skipPermissions");
+  const supportsModelPicker = nativeAgentHasCapability(selectedAgent, "modelPicker");
   const hideUnconfiguredHarnesses = useMemo(() => readHideUnconfiguredHarnesses(), []);
   // The selected native harness, used to persist/seed its option knobs (mode /
   // model / effort), which are harness-specific. null for non-native agents,
   // which have no knobs to remember.
-  const selectedNativeHarness = nativeCodingAgentForAvailableAgent(selectedAgent)?.harness ?? null;
   const selectedHost = allHosts.find((h) => h.host_id === selectedHostId);
   // Warn-only readiness signal for the agent picker: only meaningful when
   // a connected host is selected (a sandbox provisions its own tooling).
@@ -2705,6 +2832,7 @@ export function NewChatLandingScreen() {
     supportsApprovalMode ||
     supportsCursorMode ||
     supportsAgySkipPermissions ||
+    supportsModelPicker ||
     smartRoutingEligible ||
     (selectedAgent?.harness != null && selectedAgent.harness in brainHarnessLabelsAll);
   // Label/value pairs summarizing the selected agent's current run-config, for
@@ -2731,6 +2859,11 @@ export function NewChatLandingScreen() {
       // previously selected native harness.
       return [{ label: "Permissions", value: AUTO_PERMISSION_MODE.label }];
     }
+    if (supportsModelPicker && !supportsPermissionMode) {
+      const modelValue =
+        piModelOptions.find((model) => model.id === pickedModel)?.displayName ?? "Default";
+      return [{ label: "Model", value: modelValue }];
+    }
     if (supportsPermissionMode) {
       const modelValue = routingOn
         ? SMART_ROUTING_LABEL
@@ -2804,12 +2937,14 @@ export function NewChatLandingScreen() {
     supportsApprovalMode,
     supportsCursorMode,
     supportsAgySkipPermissions,
+    supportsModelPicker,
     selectedAgent,
     brainHarnessLabelsAll,
     routingOn,
     pickedModel,
     claudeModelOptions,
     codexModelOptions,
+    piModelOptions,
     pickedEffort,
     permissionMode,
     approvalMode,
@@ -2859,6 +2994,13 @@ export function NewChatLandingScreen() {
     // this holds on every run of this effect — including the re-run when the
     // model catalog resolves, which lands after the routing seed below.
     const storedRoutingOn = stored.routing === "on";
+    if (selectedNativeHarness === "pi-native") {
+      setPickedModel(
+        stored.model != null && piModelOptions.some((model) => model.id === stored.model)
+          ? stored.model
+          : "",
+      );
+    }
     if (supportsPermissionMode) {
       setPermissionMode(
         resolve(CLAUDE_NATIVE_PERMISSION_MODES, CLAUDE_NATIVE_DEFAULT_PERMISSION_MODE),
@@ -2903,7 +3045,7 @@ export function NewChatLandingScreen() {
     // Reseed on harness changes and when the selected host's catalog resolves;
     // capability flags are derived from the same harness and stay omitted.
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [selectedNativeHarness, claudeModelOptions, codexModelOptions]);
+  }, [selectedNativeHarness, claudeModelOptions, codexModelOptions, piModelOptions]);
   // Smart Routing is remembered per harness alongside the mode/model
   // knobs, in its own effect because eligibility depends on the server flag
   // (which resolves after mount — this must reseed when it lands). A stored
@@ -3618,6 +3760,7 @@ export function NewChatLandingScreen() {
       const agentSupportsApprovalMode = nativeAgentHasCapability(agent, "approvalMode");
       const agentSupportsCursorMode = nativeAgentHasCapability(agent, "cursorMode");
       const agentSupportsAgySkip = nativeAgentHasCapability(agent, "skipPermissions");
+      const agentSupportsModelPicker = nativeAgentHasCapability(agent, "modelPicker");
       // Smart Routing — server-side. The fully-auto harness always routes
       // (harness + model), so send "on" to keep the persisted state consistent
       // with the lit routing icon. Otherwise only send it when routing is
@@ -3785,13 +3928,13 @@ export function NewChatLandingScreen() {
                       ? (AGY_NATIVE_SKIP_MODES.find((m) => m.value === agySkipMode)?.args ?? [])
                       : undefined,
             // Model + reasoning effort, persisted on the session row before
-            // the runner launches. Claude and Codex read model_override at
+            // the runner launches. Claude, Codex, and Pi read model_override at
             // terminal launch; an unselected ("") knob is omitted so the
             // harness keeps its own configured/default model.
             model_override:
               !smartRoutingHarnessSelected &&
               !routingOwnsModel &&
-              (agentSupportsPermissionMode || nativeAgent?.harness === "codex-native") &&
+              (agentSupportsModelPicker || nativeAgent?.harness === "codex-native") &&
               pickedModel
                 ? pickedModel
                 : undefined,
@@ -4363,6 +4506,10 @@ export function NewChatLandingScreen() {
                     codexModelsLoading={
                       !sandboxSelected && selectedHostId !== null && hostCodexModelsLoading
                     }
+                    piModelOptions={piModelOptions}
+                    piModelsLoading={
+                      !sandboxSelected && selectedHostId !== null && hostPiModelsLoading
+                    }
                     pickedEffort={pickedEffort}
                     pickedHarness={pickedHarness}
                     costControlMode={costControlMode}