route launches by budget tier and show workspaces spend - #285
Open
asujithan wants to merge 1 commit into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Changes
Wires the AI Gateway’s :recommendModel API into the launch path, so an admin’s budget policy steers which agent and model a developer lands on as spend accumulates, and shows them the budget they’re spending against. All behavior stays behind ENABLE_MANAGED_AGENT_CONFIG and inside a managed config; enforcement remains server-side, so a failed read never blocks a launch.
recommendModel client and normalization. fetch_model_recommendation posts an empty request, and get_model_recommendation normalizes the response into {agent, model, current_spend, effective_threshold}. Money fields arrive as decimal strings and are parsed once here so callers work in numbers. Each field is optional server-side and normalized independently—an unknown agent enum is dropped rather than failing the read.
Bare ucode launches the tier’s agent. recommended_agent prefers the tier’s agent and falls back to the config’s default_agent, so a policy can move the organization to a cheaper harness without every developer reconfiguring. ucode still launches what was requested and prints a note explaining why the session isn’t using the tier’s model.
A tier’s model applies only to the agent it was recommended for. managed_launch_model checks the recommendation’s agent before overriding the config’s default_model. Without this, a tier moving the organization to OpenCode + Kimi would launch ucode claude pinned to a Kimi model, which its Anthropic-dialect endpoint can’t serve. A model returned with no agent still applies to any tool, matching the proto’s independently optional fields.
Budget panel at launch (managed_budget.py), showing spend, threshold, percent used, and a color-coded fill bar. The amber point is derived from the admin’s own lowest activating tier rather than a fixed 80%, so the color agrees with where their policy actually steps down; tiers at 0 are skipped, since they activate from the first dollar and would leave the panel.
The recommendation is read once per launch. Bare ucode needs it to pick the agent, so it’s fetched there and threaded into _launch_tool via a recommendation= parameter, mirroring the existing managed= pattern. Refetching would double the control-plane round trip and any warning it printed.
A failed budget read never blocks the launch. Both a non-None reason and a RuntimeError from token resolution fold into one warning; the config’s own default_model still applies. --skip-preflight skips the read entirely.
Notes:
RecommendModel’s recommended_agent / recommended_model pair (#2376259) is not live in staging yet: https://devportal.prod.databricks-corp.com/releases/project/ai-gateway-api?commit=71e9f894f77f9f4262fb67b315673b645edef6a2. The API currently returns spend figures only. This branch degrades correctly against that build: the panel renders from the spend figures, while agent and model fall back to the config’s defaults. The tier-routing paths are covered by unit tests rather than end-to-end tests.
Test Plan
uv run pytest → 1349 passed, 36 skipped
ruff check + ruff format clean
Verified end-to-end against eng-ml-inference.staging: normalization, fallback to config defaults, and the rendered panel ($10.09 / $5,000.00)
Added unit tests for cross-agent gating, warn-fraction derivation, and no-block-on-failure paths
Screen.Recording.2026-08-06.at.5.12.02.PM.mov