Skip to content

Add UCODE_USE_UC_SECURABLES opt-in for UC model-services and MCP-services discovery#151

Merged
AarushiShah-db merged 3 commits into
mainfrom
model-services-discovery
Jun 12, 2026
Merged

Add UCODE_USE_UC_SECURABLES opt-in for UC model-services and MCP-services discovery#151
AarushiShah-db merged 3 commits into
mainfrom
model-services-discovery

Conversation

@AarushiShah-db

@AarushiShah-db AarushiShah-db commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator

Summary

When UCODE_USE_UC_SECURABLES is set, ucode discovers two kinds of UC
securables that are otherwise invisible:

  1. Models via the Unity Catalog model-services API, addressed as
    system.ai.<model-name> instead of the per-family AI Gateway listings.
    Base URLs are unchanged — only the model name differs.

  2. Curated MCP servers via the Unity Catalog mcp-services API
    (system.ai.github, system.ai.slack, system.ai.atlassian, …),
    surfaced in ucode configure mcp alongside the existing UC connection-
    based MCP entries. Selecting one registers an HTTP MCP server pointing at
    /<workspace>/ai-gateway/mcp-services/<full-name>.

The flag is sticky — persisted into state at configure time so launches
honor it without re-exporting the env var (an explicit env var still wins).
With it unset, ucode behavior is unchanged: legacy AI Gateway model lists
and the existing MCP picker.

Discovery uses a small page size with per-page retries because the metastore listing is slow and frequently 504s, and filters to system.ai.* so model services from other schemas don't leak into the family buckets.

Test plan

  • uv run ruff check .
  • uv run pytest tests/test_databricks.py
  • UCODE_USE_UC_SECURABLES=1 uv run ucode configure on e2-dogfood
    pulls Claude/Codex models as system.ai.*.
  • UCODE_USE_UC_SECURABLES=1 uv run ucode configure mcp surfaces
    system.ai.github, etc.; selection registers correctly and
    tools/list works after one-time managed-OAuth authorize.
  • With the flag unset, model listings and MCP picker are unchanged.
  • uv run pytest (full suite)

MCP services surfaced:
Uploading Screenshot 2026-06-11 at 4.43.34 PM.png…
Screenshot 2026-06-11 at 4 50 15 PM

System.ai models surfaced:

In Claude:
Screenshot 2026-06-11 at 4 53 15 PM
Screenshot 2026-06-11 at 4 51 54 PM

In Pi (all 3 model families supported shown):
Screenshot 2026-06-11 at 4 53 37 PM

Co-authored-by: Isaac

When UCODE_USE_MODEL_SERVICES is set, discover models via the Unity Catalog
model-services API and address them as system.ai.<model-name> instead of the
per-family AI Gateway listings. Base URLs are unchanged — only the model name
differs.

The flag is sticky: it's persisted into state at configure time so launches
honor it without re-exporting the env var (an explicit env var still wins).
Discovery uses a small page size with per-page retries because the metastore
listing is slow and frequently 504s, and filters to system.ai.* so model
services from other schemas don't leak into the family buckets.

Co-authored-by: Isaac
@asujithan asujithan changed the title Add UCODE_USE_MODEL_SERVICES opt-in for UC model-services discovery Add UCODE_USE_UC_SECURABLES opt-in for UC model-services and MCP-services discovery Jun 11, 2026
@asujithan asujithan requested a review from rohita5l June 11, 2026 23:56

@rohita5l rohita5l left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@asujithan can you add some e2e tests

Comment thread src/ucode/databricks.py Outdated

@rohita5l rohita5l left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@AarushiShah-db can you also PTAL

@AarushiShah-db AarushiShah-db merged commit 0d4dc1b into main Jun 12, 2026
2 checks passed
@AarushiShah-db AarushiShah-db deleted the model-services-discovery branch June 12, 2026 19:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants