Repository navigation
Skip Claude models the user cannot call when picking defaults - #1161
nachill-databricks wants to merge 5 commits into
Conversation
The model-services listing does not say which services the caller may invoke, so `ug claude` could default to a model that 403s on every request. Probe the newest candidate per Claude family with a 1-token request, fall back to the next-newest on a 403 permission denial, and drop families with no callable model. Probe results are cached per process; any other probe outcome keeps the current behavior. Co-authored-by: Isaac <no-reply@databricks.com>
|
/ug-review |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
|
|
||
| # Per-process memo of `_claude_model_is_callable`, keyed by (workspace, model id). Only definite | ||
| # answers are stored; a probe that could not tell is never cached. | ||
| _CLAUDE_CALLABLE_CACHE: dict[tuple[str, str], bool] = {} |
There was a problem hiding this comment.
the key has no token or identity in it, so a profile/PAT switch on the same workspace in one process reuses stale answers. can we key on the profile/identity, or clear this when auth changes?
| ) | ||
| if reason is None: | ||
| callable_ = True | ||
| elif _looks_like_permission_failure(reason) and not _looks_like_scope_failure(reason): |
There was a problem hiding this comment.
_looks_like_permission_failure matches any HTTP 403, so IP-ACL, budget or policy 403s can drop a family if another one passes. can we match the EXECUTE-denial wording?
| if candidates: | ||
| claude_models[family] = candidates[0] | ||
| claude_candidates[family] = candidates | ||
| claude_models = _drop_uncallable_claude_models(workspace, token, claude_candidates) |
There was a problem hiding this comment.
this fires real billed inference calls on every configure/live refresh, even when only codex/gemini is wanted (want_claude False in cli.py). can we probe only when claude is requested?
Polly AI ReviewBlocking issues
Security vulnerabilitiesNo confirmed security vulnerabilities found. The probe cache does not grant access; the gateway still enforces authorization. Non-blocking notes
ApproachProbing candidates newest-first is a reasonable approach when listing cannot expose invocation permissions. Keep callable-model selection separate from shared inventory discovery, and distinguish definitive exhaustion from inconclusive probe failures. SummaryDiff size (computed): 3 files; +142 / -2 lines (144 changed text lines); 0 binary files.
Shares use additions + deletions, including generated text and lockfiles. Binary files have no line count. Tests include fixtures/helpers under test paths and colocated test/spec files; categories are path-based, not coverage measurements. ChangesThe implementation adds cached invocation probes and filters Claude family defaults. The remaining changes add focused tests and an offline fixture guard. There are no documentation, dependency, or generated-file changes. The implementation and support setup are proportional to the outcome; no concrete redundant scaffolding warrants removal. Restricting probing to Claude selection would reduce inference side effects without changing the intended behavior. TestsTest-by-test assessmentThe following accounts for every added test and changed support element.
Existing discovery, HTTP transport, cache, and CLI fallback tests provide useful surrounding coverage. The global probe stub means existing CLI tests do not establish the new selection-to-launch behavior. The missing boundaries are exhausted-result handling through configuration/status and the explicit OAuth-scope fail-open branch. Mocked-transport component tests are sufficient; a new live E2E is not required. The PR reports 410 passing focused tests and passing lint. Those results were not independently reproduced. Network-free isolated checks reproduced denied-model restoration and downstream fallback behavior; repository tests were not executed. The frozen diff checksum was verified. Surrounding context came from the trusted default checkout because the supplied base/head objects were unavailable locally. ScopeAll changed files support permission-aware Claude default selection or its test isolation. No clearly unrelated changes were found. The broader inference side effects at Automated review by Polly · workflow run |
…y for Claude Co-authored-by: Isaac <no-reply@databricks.com>
|
/ug-review |
| return models[0] if models and tool in ("gemini", "opencode", "kilo", "copilot", "pi") else None | ||
|
|
||
|
|
||
| _CLAUDE_MODEL_CONSUMERS = ("claude", "opencode", "kilo", "copilot", "pi") |
There was a problem hiding this comment.
can we reuse _tools_want_claude for want_claude in configure_shared_state? right now the consumer list lives in two places, so adding a Claude consumer in only one of them would leave configure and status disagreeing about whether to probe.
| def pick(models: list[str]) -> str | None: | ||
| return next((m for m in models if _claude_model_is_callable(workspace, token, m)), None) |
There was a problem hiding this comment.
each family checks its candidates one at a time with a 10s timeout each. a slow gateway could stall configure for N×10s. can we cap the candidates per family, or use a shorter total deadline?
| chosen = {family: model for family, model in picks.items() if model} | ||
| return chosen or newest |
There was a problem hiding this comment.
inconclusive probes already return True, so chosen is empty only when every candidate got a definite EXECUTE denial. can we return {} (or warn) here rather than default to models we know the user can't call?
| def discover_model_services( | ||
| workspace: str, token: str | ||
| workspace: str, token: str, *, probe_claude: bool = True | ||
| ) -> tuple[dict[str, str], list[str], list[str], list[str], str | None]: |
There was a problem hiding this comment.
with probe_claude=True as the default, any caller that doesn't pass it (e.g. ug status with no tool filter) sends billed requests. can we make probing opt-in (default False)?
| assert db_mod._claude_model_is_callable(WS, "token-a", "system.ai.claude-opus-5-5") | ||
| assert db_mod._claude_model_is_callable(WS, "token-a", "system.ai.claude-opus-5-5") | ||
| assert len(calls) == 1 |
There was a problem hiding this comment.
is this necessary? the cache-hit check already exists in test_accessible_models_unaffected_and_cached; can we keep only the token-b re-probe part here?
Polly AI ReviewBlocking issues
Security vulnerabilitiesNone identified in the frozen diff. The probe cache separates identities using a token fingerprint, and requests use the existing workspace origin. Non-blocking notes
ApproachProbing at the HTTP boundary, reusing the existing request helper and OAuth-scope classifier, and caching per identity are consistent with existing patterns. Distinguish definitive exclusion from unavailable discovery so fail-open behavior does not undo a confirmed denial. SummaryDiff size (computed): 6 files; +260 / -16 lines (276 changed text lines); 0 binary files.
Shares use additions + deletions, including generated text and lockfiles. Binary files have no line count. Tests include fixtures/helpers under test paths and colocated test/spec files; categories are path-based, not coverage measurements. ChangesThe change adds per-family Claude access probing, selection fallback, identity-scoped process caching, and CLI gating. The remaining changes adapt test setup and exercise selection and failure behavior. No documentation, dependency, or generated-file changes are included. The size is proportionate to the intended outcome. One concrete simplification is to reuse TestsTest-by-test assessment
No tests were removed. Existing coverage addresses family bucketing, discovery failures, HTTP timeout conversion, OAuth-scope classification, and live-status authentication and state preservation. It does not establish the new probe payload, negative/inconclusive cache behavior, or all-denied caller handling. No PR code, tests, or lint were executed during this review. The reported passing tests and lint are author-provided evidence, not independently verified. After the fix, run ScopeAll changes support access-aware Claude default selection or the necessary test adaptations. No clearly unrelated work was found. The all-denied fallback is a correctness gap within the stated outcome, not a separate scope concern. Automated review by Polly · workflow run |
…fallback Co-authored-by: Isaac <no-reply@databricks.com>
|
/ug-review |
| discover_model_services( | ||
| workspace, | ||
| token, | ||
| probe_claude=_tools_want_claude(tools), |
There was a problem hiding this comment.
is this necessary for status? every ug status with claude configured now sends up to 8 billed 1-token calls and can wait up to 15s. the cache only lasts one process. can we skip the probe here (probe_claude=False)?
Suggestion:
| probe_claude=_tools_want_claude(tools), | |
| probe_claude=False, |
| ): | ||
| return model | ||
| if len(models) > _CLAUDE_PROBE_MAX_CANDIDATES: | ||
| return models[_CLAUDE_PROBE_MAX_CANDIDATES] # beyond the probe cap: unprobed, keep |
There was a problem hiding this comment.
if the first two models are definitely denied, the unprobed third one becomes the default, but with exactly two denied models the family is dropped. is that inconsistency intended? the third model is likely denied too
|
|
||
| with ThreadPoolExecutor(max_workers=len(candidates)) as pool: | ||
| picks = dict(zip(candidates, pool.map(pick, candidates.values()), strict=True)) | ||
| return {family: model for family, model in picks.items() if model} |
There was a problem hiding this comment.
what happens if every family is denied? {} with reason None makes configure_shared_state fall back to discover_claude_models, which lists the same denied models unprobed. can we return a reason so callers skip that fallback?
| assert not cli_mod._tools_want_claude({"codex"}) | ||
| assert not cli_mod._tools_want_claude({"gemini", "codex"}) | ||
|
|
||
| def test_live_refresh_skips_probe_for_codex_only(self): |
There was a problem hiding this comment.
can we also assert probe_claude in configure_shared_state? that's the main billed path, and every configure stub here swallows the kwarg with **_kw, so a codex-only configure that probes would still pass.
| claude, *_ = db_mod.discover_model_services(WS, "token", probe_claude=True) | ||
| assert claude == {} | ||
|
|
||
| def test_probing_is_opt_in(self, monkeypatch): |
There was a problem hiding this comment.
is this needed? test_probe_skipped_when_claude_not_wanted already checks that no probes are sent and the unprobed pick is kept
…skip-uncallable-models Co-authored-by: Isaac <no-reply@databricks.com>
Polly AI ReviewBlocking issues
Security vulnerabilitiesNone demonstrated. Fail-open selection does not bypass the gateway’s authorization checks. Non-blocking notes
ApproachThe bounded, fail-open probe fits the existing discovery helpers. The main correction is carrying an intentional-denial result through the caller boundary. The separate SummaryDiff size (computed): 6 files; +301 / -25 lines (326 changed text lines); 0 binary files.
Shares use additions + deletions, including generated text and lockfiles. Binary files have no line count. Tests include fixtures/helpers under test paths and colocated test/spec files; categories are path-based, not coverage measurements. ChangesThe implementation adds permission-aware Claude candidate selection, process-local caching, bounded probing, and caller opt-in. The remaining changes provide unit coverage and adapt existing discovery stubs. No dependencies, generated files, or documentation change. The size is proportional to the outcome. Concrete reductions are the duplicate cache-clear interface/call and overlapping opt-in and live-refresh test setup; no broad restructuring is needed. TestsTest-by-test assessmentThe frozen snapshot’s checksum matched, and both trusted policy files were read in full. Tests were not executed; the PR’s reported test, lint, and live results remain unverified. The supplied base/head commits were unavailable locally, so surrounding caller behavior was checked against the provided trusted default-branch checkout.
No tests are removed. Modified existing tests change stub signatures rather than weaken assertions. Existing coverage protects discovery bucketing, pagination, request errors, live-status non-persistence, and search-registration preservation. It does not cover the demonstrated all-denied caller transition. Probe request shape and non-caching of inconclusive results also lack direct assertions. These gaps belong at unit/component seams; duplicate live/E2E coverage is unnecessary. The real gateway’s denial wording remains unverified. Synthetic denial tests establish local classification behavior, not the provider response contract. To verify the blocking fix offline, advertise the same denied candidates through both discovery sources, exercise configure and live refresh, and assert those candidates cannot reappear. Run the four changed test files in a credential-free, network-disabled checkout. ScopeAll changes support permission-aware defaults or the necessary adaptation of existing tests. No clearly unrelated outcome was found. Probing during live status is a related behavior expansion whose necessity needs clarification, not evidence of unrelated work. Automated review by Polly · workflow run |
Co-authored-by: Isaac <no-reply@databricks.com>
1 participant
Problem
The model-services listing does not say which services the caller may invoke, so
ug claudecould default to a model that returns 403 on every request.Change
Probe the newest candidate per Claude family with a 1-token request. On a 403 permission denial, fall back to the next-newest model; families with no callable model are dropped. Results are cached per process. The probe is fail-open: any other outcome (success, OAuth-scope 403, 5xx, network error) keeps current behavior.
Before: the default Claude model could be one the user is denied. After: denied models are skipped.
Testing
tests/test_databricks.py.uv run pytest -q tests/test_databricks.py: 410 passed.uvx --from rust-just just lint: passes (ruff check, ruff format --check, ty check).tests/test_agents_init.pyandtests/test_e2e_user_agent.pyfail onmaintoo, independent of this change.main(no regression).Not yet verified
This pull request and its description were written by Isaac.