feat(agent): discover Databricks Unity Catalog models - #6918
Conversation
1cf6c66 to
0364cf4
Compare
wpfleger96
left a comment
There was a problem hiding this comment.
🤖 hey @kalvinnchau, thanks for this — I had the team go through it (source review against the model-capabilities manifest work, plus a live E2E run) and I'm consolidating everything here.
The good news first: the manifest stays the single owner of routing, labels, and effort everywhere it's valid. scripts/model-capabilities.json and both interpreters are untouched, there are no new hand-maintained name/route tables, DATABRICKS_MODEL_FILTER stays pure visibility with filtered-empty catalogs authoritative in both the ACP and Desktop paths, and the databricks_v2_dispatch_is_pure_manifest_projection mutation guard is retained. We also verified the whole path live at this head: the PR binary's catalog contained all 58 FQNs returned by /api/2.1/unity-catalog/model-services (no additions, no omissions), and the desktop picker rendered humanized labels for every one the manifest knows — system.ai.claude-opus-5 → "Claude Opus 5", data_workflow_tools.goose.goose-kimi-k3 → "Kimi K3", etc. CI is green across the required gates; the red "Mark Previous Review Stale" job is a 403 from the housekeeping workflow's token, unrelated to this change.
One thing we think should be fixed before merge — the UC-FQN capability rule only exists on the Rust side, so the desktop effort picker and the wire path can now disagree about the same model id. Details in the inline comment on databricks_v2_capability_model, plus one trivial dead assignment also inline.
Two observations that look intentional but are worth an explicit yes:
- Embedding endpoints (
databricks-bge-large-enetc.) are back in the unfiltered picker now that the keyword exclusion is deleted. That matches the task's acceptance criteria, but workspaces withoutDATABRICKS_MODEL_FILTERwill see them again, and selecting one still 400s at send time (the removed code's own comment said as much). Flagging so it's a deliberate choice rather than a surprise. - In the live run, 29 of the 58 UC model-services render as raw FQNs because they have no manifest record or alias (the deepseek/grok/glm families and a few others). That's the designed behavior — no guessed labels — so humanizing them is a manifest edit in a follow-up, not a change to this PR.
Extend Databricks v2 discovery with Unity Catalog model services and a shared visibility filter. Route model-service FQNs through MLflow Chat Completions while preserving Desktop and ACP fallback parity. Signed-off-by: Kalvin Chau <kalvin@block.xyz> Signed-off-by: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz>
Enclose each catalog attempt in one deadline covering response headers, bounded body consumption, and JSON parsing. Retry transient statuses, body failures, malformed JSON, transport failures, and attempt timeouts without logging provider bodies or bearer-derived data. Add deterministic stalled-body, transient-recovery, exhaustion, malformed-response, and redaction coverage. Signed-off-by: Kalvin Chau <kalvin@block.xyz> Signed-off-by: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz>
Co-authored-by: Kalvin Chau <kalvin@block.xyz> Signed-off-by: Kalvin Chau <kalvin@block.xyz> Signed-off-by: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz>
0364cf4 to
93dd353
Compare
Co-authored-by: Kalvin Chau <kalvin@block.xyz> Signed-off-by: Kalvin Chau <kalvin@block.xyz> Signed-off-by: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz>
wpfleger96
left a comment
There was a problem hiding this comment.
🤖 Approving at 33ab5db69 — all three review lanes are clean at this exact head:
- Source (manifest alignment): UC-FQN capability parity fixed in
resolve()in both interpreters, routing back to a pure manifest projection, and the UC chat filter screens onsupported_api_typesmetadata with backward-compatible fallback — no name heuristics. - Independent source re-review: confirms the filter criterion matches the actual UC wire path (
mlflow/v1/chat/completions) and that the parser test fails if the production filter is removed. - Live E2E: 60 FQNs fetched from the UC endpoint, 55/55 chat-eligible services in the agent catalog with all five non-chat services excluded, humanized labels rendering (
Claude Opus 5,Kimi K3), neutral FQN effort set in the picker, and a real prompt through a UC FQN completing on the MLflow route.
CI green at the head; the lone red Mark Previous Review Stale job is the known housekeeping-token 403, not product CI. Nice work on the quick turnarounds.
Summary
Add Unity Catalog model-service discovery to
buzz-agentDatabricks v2, preserve full Unity Catalog FQNs through ACP model discovery, route model services through MLflow chat, and add the optionalDATABRICKS_MODEL_FILTERvisibility filter.The resilience follow-up bounds each catalog attempt across headers, bounded body reads, and JSON parsing; retries transient 499/5xx, transport/body failures, malformed JSON, and attempt timeouts; and emits only safe catalog/status diagnostics.
Validation
The live validation and test gates below were run on the pre-rewrite tip
1cf6c66e7c5bf9a00c357bda24df48e3beeb310b. The rewritten tip0364cf44f3eb5cfbd5771b0a0dd7868ce0ad2c19has identical trees for both commits, so the attribution rewrite changes no tested code.cargo check -p buzz-agentpassedcargo clippy -p buzz-agent --all-targets -- -D warningspassedcargo fmt --checkpassedgit diff --checkpassedThe full
buzz-agentpackage has one separately reproduced pre-existing failure:regressions::cancel_kills_inflight_tool_via_mcp_notification, failing on both baselineeb8e97cfeand the follow-up worktree. No Databricks path reaches that test. The live validation proves discovery and ACP exposure, not inference against every listed model.