Skip to content

route launches by budget tier and show workspaces spend - #285

Open
asujithan wants to merge 1 commit into
mainfrom
asujithan/managed-budget
Open

route launches by budget tier and show workspaces spend#285
asujithan wants to merge 1 commit into
mainfrom
asujithan/managed-budget

Conversation

@asujithan

Copy link
Copy Markdown
Collaborator

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.

  1. 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.

  2. 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.

  3. 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.

  4. 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.

  5. 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.

  6. 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

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.

1 participant