Skip to content

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

Open
anthonyivn2 wants to merge 3 commits into
omnigent-ai:mainfrom
anthonyivn2:feat/pi-start-model-picker
Open

feat(pi): Add searchable model picker for new sessions with Databricks Unity AI Gateway OAuth#3481
anthonyivn2 wants to merge 3 commits into
omnigent-ai:mainfrom
anthonyivn2:feat/pi-start-model-picker

Conversation

@anthonyivn2

@anthonyivn2 anthonyivn2 commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Related issue

N/A — no issue is currently linked.

Summary

This PR extends the existing host-resolved model selection path to Pi. After configuring a provider with omni setup, users can search and select one of its models before starting Pi. Omnigent sends the exact provider and model through launch, so Pi starts and reports model changes using the same unambiguous identifier.

  • Resolve pi-native model options on the selected host from the provider generated by omni setup; local Pi configuration such as ~/.pi is intentionally excluded.
  • Add Pi to the existing harness configuration modal with a wider, searchable, scrollable picker while preserving Claude, Codex, and Smart Routing behavior from current main.
  • Carry provider-qualified provider/model selections through model_override, Pi terminal launch, namespaced model IDs, startup reporting, and live model-change events.
  • Cover host catalog resolution, provider-qualified launch selection, extension behavior, and the new-session create flow.
  • Configuration UI does not have an option to configure reasoning level at the start or during runtime, this would be resolved in another PR
flowchart LR
    A[omni setup] --> B[Managed Pi provider catalog]
    B --> C[Host model_options for pi-native]
    C --> D[Searchable new-session picker]
    D --> E[provider/model override]
    E --> F[Pi terminal launch]
    E --> G[Startup and live model events]

    H[Local ~/.pi configuration] -. intentionally excluded .-> C
Loading

Test Plan

  • Backend, host, launch, and extension coverage — 258 passed:
    • uv run pytest -q tests/host/test_connect.py tests/test_pi_native_credentials.py tests/runner/test_app_sessions_native_terminals_autocreate.py tests/test_pi_native_extension.py
  • TypeScript project build:
    • web/node_modules/.bin/tsc -b web
  • Focused new-session dialog flow coverage — 57 passed:
    • cd web && ./node_modules/.bin/vitest run src/shell/NewChatDialog.flow.test.tsx
  • Full repository pre-commit suite:
    • uv run --frozen pre-commit run --all-files
  • Manual verification steps:
    1. Configure a Pi provider using omni setup.
    2. Start a new Pi session and open its configuration dialog.
    3. Confirm only Omnigent-configured models appear.
    4. Search with multiple keywords, such as gpt sol, and scroll the results.
    5. Hover a truncated model and confirm its full provider-qualified name appears.
    6. Select a model and confirm the Pi session launches with that model.

Demo

Gear option on the home page

Screenshot 2026-08-06 at 19 15 31

Searchable model picker

Screenshot 2026-08-06 at 19 15 19 Screenshot 2026-08-06 at 19 15 10

Running with the selected model

Screenshot 2026-08-06 at 19 16 29 Screenshot 2026-08-06 at 19 20 23

Type of change

  • Bug fix
  • Feature
  • UI / frontend change
  • Refactor / chore
  • Docs
  • Test / CI
  • Breaking change

Test coverage

  • Unit tests added / updated
  • Integration tests added / updated
  • E2E tests added / updated
  • Manual verification completed
  • Existing tests cover this change
  • Not applicable

Coverage notes

Automated coverage exercises host catalog resolution, provider-qualified launch selection, extension model events, and the new-session create request. The attached screenshots demonstrate the intended UI behavior; no additional manual UI pass was performed during the rebase.

Changelog

Pi sessions now provide a searchable model picker for providers configured through omni setup.

@github-actions github-actions Bot added the size/XL Pull request size: XL label Jul 29, 2026
@anthonyivn2
anthonyivn2 marked this pull request as ready for review July 29, 2026 05:30
@anthonyivn2 anthonyivn2 changed the title [pi] Add searchable model picker for new sessions feat(pi): Add searchable model picker for new sessions Jul 29, 2026
@github-actions
github-actions Bot requested a review from serena-ruan July 29, 2026 05:35
Signed-off-by: Anthony Ivan <anthony.ivan@databricks.com>
@anthonyivn2
anthonyivn2 force-pushed the feat/pi-start-model-picker branch from 2f27bdc to 3583086 Compare August 6, 2026 05:29
@anthonyivn2 anthonyivn2 changed the title feat(pi): Add searchable model picker for new sessions feat(pi): Add searchable model picker for new sessions with Databricks Unity AI Gateway OAuth Aug 6, 2026
@TomeHirata

TomeHirata commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

What's omnigent-openai in the UI?

@TomeHirata

Copy link
Copy Markdown
Contributor

/review

@omnigent-ci

omnigent-ci Bot commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Polly AI Review

Review: Pi searchable model picker

The change is well-scoped and its Python-side slash-parsing/launch logic is correct and consistent with the pre-existing namespaced-model path. However, both an independent review and a second-vendor review independently flagged the same two problems in the JS extension, centered on the new "qualified provider/model id" contract not being made robust to Pi builds/objects that predate or omit that shape.

Blocking issues

B1 — applyModelChange fallback can never match a qualified id on Pi builds without registry.find (omnigent/resources/pi_native/omnigent_pi_native_extension.js ~L896-906).
postModelOptions now emits qualified ids (provider/modelId), so the web picker and inbox model_change deliver qualified ids into applyModelChange. The resolve logic is:

const separator = id.indexOf("/");
if (separator > 0 && registry && typeof registry.find === "function") {
  model = registry.find(id.slice(0, separator), id.slice(separator + 1));
}
if (!model) {
  model = listModels().find((m) => m && m.id === id); // m.id is BARE
}

The fallback compares the qualified id (omnigent/databricks-claude-opus-4-1) against m.id, which is the bare model id — so it can never match. When registry.find exists this is masked, but this file explicitly accommodates "older Pi" builds exposing only getAll/getAvailable. On any such build, every web/inbox-initiated model switch now fails with "model … is not available." The fallback should split the id and match on the bare model id (m.id === id.slice(separator + 1)), ideally also checking m.provider. The updated test only exercises the find path, so this gap is uncovered.

B2 — Hard provider requirement replaces the previous bare-id behavior on three paths → silent regression if Pi's model objects don't always carry .provider.

  • postModelOptions (~L995): if (!modelId || !provider) continue; — if models lack provider, all are skipped, options is empty, and the picker is never posted (disappears).
  • startup mirror (~L1753): now gated on startupProvider && startupModelId — no startup external_model_change, so the composer pill/active row show nothing.
  • model_select (~L1783): if (!provider || !modelId) return; — in-TUI /model switches stop mirroring back to the web UI.

Previously these used the bare id. This is a genuine regression iff Pi's modelRegistry entries / ctx.model / model_select event.model don't always populate provider — a contract that lives in the external Pi CLI and can't be confirmed from this repo. The tests were updated to add provider to the mock catalog, which would mask a real-world regression rather than catch it. Recommend either confirming the Pi contract guarantees provider on all these objects, or falling back to the bare id when provider is absent (preserving existing Databricks/Anthropic//login behavior) instead of dropping the model entirely.

B3 — Provider-qualified launch selection silently falls back to the provider's default model when unmatched (omnigent/runner/native/orchestration.py ~L2127 + omnigent/pi_native_credentials.py ~L948).
When spec_model contains /, resolve_pi_native_provider() is called without the selected model, and pi_native_provider_launch only overrides model_provider_id/selected_model if the selection is found in the freshly regenerated provider config; otherwise it silently keeps the provider default. So a selection like omnigent/databricks-claude-opus-4-1 can launch the default (e.g. Sonnet) instead of Opus with no error if the launch-time model list differs from what populated the picker. The second review argues the picker and launch stay in lockstep because both iterate the same to_models_config() on a provider resolved the same way, which makes this benign in the common path — but the failure mode is a silent wrong-model launch. Prefer failing loud (reject/diagnose an unmatched qualified selection) rather than silently substituting the default.

Security vulnerabilities

None found. _logger.exception("Failed to resolve pre-launch Pi model options") logs server-side only; pi_native_model_options/to_models_config don't raise with api_key embedded, and the client-facing error is

Both cross-vendor reviews are in and converge strongly. Synthesizing.

Review: Pi searchable model picker

The change is well-scoped and its Python-side slash-parsing/launch logic is correct and consistent with the pre-existing namespaced-model path. However, both an independent review and a second-vendor review independently flagged the same two problems in the JS extension, centered on the new "qualified provider/model id" contract not being made robust to Pi builds/objects that predate or omit that shape.

Blocking issues

B1 — applyModelChange fallback can never match a qualified id on Pi builds without registry.find (omnigent/resources/pi_native/omnigent_pi_native_extension.js ~L896-906).
postModelOptions now emits qualified ids (provider/modelId), so the web picker and inbox model_change deliver qualified ids into applyModelChange. The resolve logic is:

const separator = id.indexOf("/");
if (separator > 0 && registry && typeof registry.find === "function") {
  model = registry.find(id.slice(0, separator), id.slice(separator + 1));
}
if (!model) {
  model = listModels().find((m) => m && m.id === id); // m.id is BARE
}

The fallback compares the qualified id (omnigent/databricks-claude-opus-4-1) against m.id, which is the bare model id — so it can never match. When registry.find exists this is masked, but this file explicitly accommodates "older Pi" builds exposing only getAll/getAvailable. On any such build, every web/inbox-initiated model switch now fails with "model … is not available." The fallback should split the id and match on the bare model id (m.id === id.slice(separator + 1)), ideally also checking m.provider. The updated test only exercises the find path, so this gap is uncovered.

B2 — Hard provider requirement replaces the previous bare-id behavior on three paths → silent regression if Pi's model objects don't always carry .provider.

  • postModelOptions (~L995): if (!modelId || !provider) continue; — if models lack provider, all are skipped, options is empty, and the picker is never posted (disappears).
  • startup mirror (~L1753): now gated on startupProvider && startupModelId — no startup external_model_change, so the composer pill/active row show nothing.
  • model_select (~L1783): if (!provider || !modelId) return; — in-TUI /model switches stop mirroring back to the web UI.

Previously these used the bare id. This is a genuine regression iff Pi's modelRegistry entries / ctx.model / model_select event.model don't always populate provider — a contract that lives in the external Pi CLI and can't be confirmed from this repo. The tests were updated to add provider to the mock catalog, which would mask a real-world regression rather than catch it. Recommend either confirming the Pi contract guarantees provider on all these objects, or falling back to the bare id when provider is absent (preserving existing Databricks/Anthropic//login behavior) instead of dropping the model entirely.

B3 — Provider-qualified launch selection silently falls back to the provider's default model when unmatched (omnigent/runner/native/orchestration.py ~L2127 + omnigent/pi_native_credentials.py ~L948).
When spec_model contains /, resolve_pi_native_provider() is called without the selected model, and pi_native_provider_launch only overrides model_provider_id/selected_model if the selection is found in the freshly regenerated provider config; otherwise it silently keeps the provider default. So a selection like omnigent/databricks-claude-opus-4-1 can launch the default (e.g. Sonnet) instead of Opus with no error if the launch-time model list differs from what populated the picker. The second review argues the picker and launch stay in lockstep because both iterate the same to_models_config() on a provider resolved the same way, which makes this benign in the common path — but the failure mode is a silent wrong-model launch. Prefer failing loud (reject/diagnose an unmatched qualified selection) rather than silently substituting the default.

Security vulnerabilities

None found. _logger.exception("Failed to resolve pre-launch Pi model options") logs server-side only; pi_native_model_options/to_models_config don't raise with api_key embedded, and the client-facing error is a static string. Credentials still flow only through the generated models.json, never CLI args.

Non-blocking notes

  • The _logger.exception traceback is low-risk, but consider _logger.warning without the full traceback to avoid any chance of config values surfacing in a chained exception repr.
  • pi_native_model_options() dedup-by-qualified-id, deterministic sort, and displayName fallback all check out.
  • Slash qualification logic verified correct: omnigent-openai/gpt-5.6-sol--provider omnigent-openai --model gpt-5.6-sol (bare), while namespaced moonshotai/kimi-k2.5 retains its slash and stays qualified. Consistent with existing behavior. registry is confirmed in scope in applyModelChange.
  • Models not renderable without threading the model into resolve_pi_native_provider also won't appear in the picker — acceptable since picker and launch stay in lockstep, but worth being aware of.

Summary

Solid, well-tested feature with correct Python launch-arg logic and no security concerns; screenshots adequately demonstrate the UI. The blockers all stem from one theme: the new qualified provider/model id contract wasn't made resilient to Pi builds/objects that lack registry.find or omit provider (B1, B2), and the launch path substitutes the default model silently on an unmatched selection (B3). B1 is a concrete, provable defect; B2 and B3 are real regression/correctness risks gated on the external Pi contract and launch-time consistency. Address the fallback matching, add a bare-id fallback (or confirm the provider guarantee), and make an unmatched selection fail loud before merge.


Automated review by Polly · workflow run

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

anthonyivn2 commented Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

@TomeHirata I have made changes related to the review, and removed omnigent* from the model name shown on the model picker list. Screenshots in the PR body is updated

Signed-off-by: Anthony Ivan <anthony.ivan@databricks.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size/XL Pull request size: XL waiting-for-review

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants