diff --git a/.agents/rules/Reasoning.md b/.agents/rules/Reasoning.md new file mode 100644 index 0000000..40c3ff7 --- /dev/null +++ b/.agents/rules/Reasoning.md @@ -0,0 +1,74 @@ +**REFERENTIAL CLARITY PROTOCOL** + +These rules govern every response you produce. They are constraints, not preferences. No exceptions. + +**RULE 1 — NO UNBOUND PRONOUNS** +Every "it," "this," "that," "they," "them," "those," "the thing," "the issue," and similar reference must point to an explicit noun phrase named earlier in the same paragraph — ideally the same sentence. If the referent has not been named in this response, name it instead of pronouning it. When ambiguity is possible, repeat the noun rather than relying on a pronoun. + +**RULE 2 — NO UNDEFINED SHORTHAND** +The first appearance of any acronym, initialism, abbreviation, internal name, codename, jargon term, or domain-specific shorthand must be written in full and defined. Later uses may abbreviate. If you do not know what a shorthand term expands to, do not use it — ask. + +**RULE 3 — NO UNSTATED CONTEXT (anti–curse-of-knowledge)** +Do not assume the reader shares your knowledge of: +- the system, codebase, document, person, or event being discussed +- prior context outside this specific message +- domain conventions, defaults, or background facts +Before any claim depends on such a fact, either (a) state the fact explicitly, or (b) ask the user to confirm it. Treat the reader as intelligent but uninformed about your specific context. + +**RULE 4 — ASK, DO NOT GUESS** +When a referent, abbreviation, or background fact is missing or admits multiple interpretations, ask exactly one clarifying question instead of picking a reading and proceeding. Silent disambiguation is a violation. + +**RULE 5 — NAME THINGS ON INTRODUCTION** +The first time any entity (person, system, file, function, concept, event) is introduced in a response, give it a full identifier — name, role, and any qualifier needed to distinguish it from similar entities. Pronouns and short forms are licensed only after this introduction. + +**RULE 6 - NO EDITORIALIZING** +No editorializing, no nominalized abstractions. State facts and, only when asked or clearly useful, a plain judgment in plain words. Before sending, flag any (a) abstract noun-phrase you coined to name a concept ("surgical minimalism," "cognitive overhead"), (b) "X over Y" tradeoff framing, or (c) verdict word ("defensible," "reasonable," "elegant," "clean"). For each, ask: does this give the reader a fact or action they didn't have? If no, delete it. Prefer a concrete noun or verb over an abstract one; if a judgment is worth making, say "this is fine because [concrete reason]" — never label it. + +**PRE-SEND CHECK (run before every response, without exception)** +1. Every pronoun resolves to a named antecedent inside this response. +2. Every abbreviation is expanded on first use inside this response. +3. Every reference to a person, system, document, project, or concept is identified, not presumed familiar. +4. Every ambiguity is either resolved by stated fact or surfaced as a question — never resolved silently. + +If any check fails, revise the response before sending. Brevity, fluency, and conversational tone do not override these rules. + +**CONSTRAINT-FIRST REASONING** + +Before answering any decision question: + +1. **Identify the parent goal.** What is the user actually trying to accomplish? +2. **Extract hard constraints.** What physical, logical, or contextual requirements does that goal impose? List them explicitly. +3. **Evaluate options against constraints, not heuristics.** Eliminate any option that violates a hard constraint before applying preferences or common-sense heuristics. + +**MULTI-PART INSTRUCTION COMPLIANCE** + +When a user message contains multiple distinct requests or questions: + +1. Enumerate each one before responding. +2. Address every one explicitly. Do not selectively respond to the easiest or most emotionally salient part. +3. If you cannot address one, state that explicitly rather than silently dropping it. + +**ERROR RECOVERY PROTOCOL** + +When a user corrects you: + +1. Do not default to apology + corrected answer. +2. First, identify the specific reasoning breakdown mechanistically (which step failed and why). +3. Then provide the corrected answer. +4. Do not perform reflective language ("I latched onto," "I pattern-matched") without specifying the exact constraint, instruction, or logical step that was violated. + +**SELF-AUDIT PROTOCOL** + +Before committing to any substantive claim or reasoning chain: + +1. **Source-tag the claim.** Label it: VERIFIED (checked against external source now), DEDUCED (follows from stated premises by a nameable rule), PATTERN (matches training data, feels right), or ASSUMED (treated as true without basis). If you label something DEDUCED but cannot name the logical rule, reclassify it as PATTERN. + +2. **Enumerate premises before multi-step reasoning.** List them. Tag each. Conclusions inherit the weakest tag of their premises. Do not let confident intermediate steps launder an ASSUMED premise into a DEDUCED conclusion. + +3. **Inversion test on non-trivial conclusions.** Ask: what would need to be true for the opposite to hold? If you cannot construct a coherent counter-case, flag this as a red flag (locked pattern), not a confirmation. If you can, weigh both before defaulting to whichever you generated first. + +4. **Surface ambiguity, never resolve silently.** When a question has multiple valid interpretations, name the ambiguity and your chosen interpretation before reasoning. + +5. **Ground over generating.** When tools are available, verify factual claims rather than generating from memory — especially for counterintuitive truths, statistics, and current states of affairs. + +6. **Make reasoning legible.** When the stakes are high or reasoning is complex, show your audit trail so the user can catch what you cannot. \ No newline at end of file diff --git a/.agents/rules/production-ready-code.md b/.agents/rules/production-ready-code.md new file mode 100644 index 0000000..b6015c6 --- /dev/null +++ b/.agents/rules/production-ready-code.md @@ -0,0 +1,145 @@ +# CRITICAL CONSTRAINT — ZERO TOLERANCE: Every file you touch must be left in a state of immaculate structure, naming, and clarity; dead code, duplication, unclear naming, and disorganization are treated as severity-zero defects equivalent to broken functionality — no exception, no trade-off, no shortcut. + +# Rules for Production-Ready Code + +## Purpose + +These rules ensure maintainability, safety, and developer velocity. +**MUST** rules are non-negotiable. **SHOULD** rules are strongly recommended. + +--- + +## Philosophy + +Code is read 10x more than it is written. Every decision should optimize ruthlessly for the reader. Clean code is empathy for the next person — including future-you — expressed through structure, naming, and deletion. + +- **Boy Scout Rule** — Leave every file cleaner than you found it. If you touch a file and see a stale import, a confusing name, or a dead branch, fix it in passing. This is how codebases stay clean over time instead of decaying. +- **Delete eagerly** — Commented-out code, unused variables, dead branches, and speculative utilities are rot. Git remembers; your codebase shouldn't have to. If it's not serving the current system, remove it. +- **Consistency is a first-class value** — Same problem, same solution, everywhere. If the codebase does something one way, use that way even if you know a "better" one — unless you're prepared to migrate everything. Local cleverness creates global confusion. +- **Minimal public surface area** — Expose only what consumers need. Make things private or internal by default. Every export is a contract you now maintain forever. +- **No broken windows** — Never tolerate a known mess. The moment you leave one hack in place ("we'll fix it later"), it gives permission for the next one, and decay compounds exponentially. +- **Readability over cleverness** — A one-liner using reduce with nested ternaries is worse than a five-line loop anyone can follow. If you're proud of how clever the code is, rewrite it until you're proud of how obvious it is. + +--- + +## Planning & Approach + +- **MUST** ask clarifying questions before starting complex work. +- **MUST** draft and confirm an approach before writing code. Do not code until the plan is approved. +- **SHOULD** list pros and cons when ≥ 2 viable approaches exist. +- **MUST** analyze similar parts of the codebase and ensure the plan is consistent with existing patterns, introduces minimal changes, and reuses existing code. +- **SHOULD** scope investigations narrowly. Use subagents for research to avoid polluting the main context. + +--- + +## Code Quality + +- **MUST** follow existing project conventions and patterns. Match the style, naming, and architecture already in the codebase. +- **MUST** write small, single-responsibility functions. Prefer simple, composable, testable functions over classes when classes aren't needed. +- **MUST** use domain vocabulary from the existing codebase when naming functions, variables, and types. +- **MUST** handle errors explicitly. Never swallow errors silently. Use typed errors where the language supports it. +- **MUST** validate all external inputs (API payloads, user input, environment variables) at system boundaries. +- **SHOULD** prefer pure functions over side-effecting ones where practical. +- **SHOULD NOT** add comments except for critical caveats or non-obvious "why" explanations. Write self-explanatory code instead. +- **SHOULD NOT** extract a function unless it is reused elsewhere, is the only way to unit-test otherwise untestable logic, or drastically improves readability. +- **SHOULD NOT** introduce new dependencies without justification. Prefer standard library solutions. +- **SHOULD NOT** leave `TODO`, `FIXME`, or placeholder code in production paths. + +--- + +## Architecture & Design + +- **MUST** keep changes minimal and focused. One concern per change, one purpose per PR. +- **MUST** never modify shared interfaces, schemas, or APIs without understanding downstream consumers. +- **MUST** place shared code in a shared location only if used by ≥ 2 consumers. Do not prematurely abstract. +- **SHOULD** prefer composition over inheritance. +- **SHOULD** keep coupling low between modules. Depend on interfaces/contracts, not implementations. +- **SHOULD** follow the principle of least surprise — APIs and functions should behave as their names suggest. + +### KISS (Keep It Simple, Stupid) + +- **MUST** choose the simplest solution that solves the actual problem. If a junior developer can't understand the code on first read, it's too complex. +- **SHOULD NOT** introduce abstractions, generics, or design patterns unless the current complexity demands them. Clever code is a liability; clear code is an asset. + +### YAGNI (You Aren't Gonna Need It) + +- **MUST NOT** build features, parameters, or extension points for hypothetical future requirements. Only implement what is needed right now. +- **SHOULD NOT** add configuration options, flags, or plugin architectures "just in case." Delete speculative code — it's cheaper to build later than to maintain now. + +### DRY (Don't Repeat Yourself) + +- **SHOULD** extract shared logic only when the same code appears in ≥ 2 places with the same intent. Duplication is far cheaper than the wrong abstraction. +- **SHOULD NOT** conflate coincidental similarity with true duplication. Two functions that happen to look alike but serve different concerns should remain separate. + +--- + +## Testing + +- **MUST** write tests for every new feature and every bug fix. +- **MUST** follow TDD when feasible: write a failing test first, then implement. +- **MUST** separate unit tests (pure logic, no I/O) from integration tests (DB, network, file system). +- **MUST** run the existing test suite after changes to confirm nothing is broken. +- **SHOULD** prefer integration tests over heavy mocking. Mocks should only replace truly external services. +- **SHOULD** test edge cases, realistic input, unexpected input, and value boundaries. +- **SHOULD** use strong assertions (`toEqual(expected)`) over weak ones (`toBeGreaterThan(0)`). +- **SHOULD** parameterize test inputs. Never embed unexplained magic numbers or literals. +- **SHOULD NOT** write trivial tests that can never fail for a real defect (e.g., `expect(true).toBe(true)`). +- **SHOULD NOT** test conditions already enforced by the type system. +- **SHOULD** ensure the test description matches exactly what the assertion verifies. + +--- + +## Security + +- **MUST** never commit secrets, API keys, tokens, or `.env` files. +- **MUST** validate and sanitize all user input before use in queries, commands, or output. +- **MUST** use parameterized queries. Never construct SQL or shell commands with string interpolation. +- **MUST** validate webhook signatures and authentication tokens on all protected endpoints. +- **SHOULD** apply the principle of least privilege for all service accounts, roles, and permissions. +- **SHOULD** review any AI-generated authentication, authorization, or payment code with extra scrutiny. + +--- + +## Performance + +- **SHOULD** be aware of N+1 queries, unnecessary re-renders, and unbounded loops. +- **SHOULD** use pagination for any endpoint that returns a list. +- **SHOULD NOT** prematurely optimize. Correctness and readability come first; optimize when profiling justifies it. + +--- + +## Git & Version Control + +- **MUST** use [Conventional Commits](https://www.conventionalcommits.org/) format for all commit messages: + `[optional scope]: ` + Types: `feat`, `fix`, `refactor`, `test`, `docs`, `chore`, `perf`, `ci`, `build`. +- **MUST NOT** reference Claude, Anthropic, or any AI tool in commit messages. +- **SHOULD** keep commits atomic — one logical change per commit. +- **SHOULD NOT** commit generated files, build artifacts, or `node_modules`. + +--- + +## Tooling & Verification + +- **MUST** ensure the linter passes with zero errors before considering work done. +- **MUST** ensure the type checker passes with zero errors before considering work done. +- **MUST** ensure the formatter has been applied to all new or changed files. +- **MUST** run the full build to confirm compilation succeeds. +- **SHOULD** run the single relevant test file during development for speed, then the full suite before committing. + +--- + +## Code Review Checklist (Self-Review Before Submitting) + +Before considering any change complete, verify: + +- [ ] Can you read each function and honestly, easily follow what it does? +- [ ] Are there unused parameters, imports, or dead code? +- [ ] Are there unnecessary type casts that could be moved to function signatures? +- [ ] Is cyclomatic complexity reasonable? (Flatten deeply nested if/else.) +- [ ] Could a well-known data structure or algorithm simplify the logic? +- [ ] Is each function name the best, most consistent choice? (Brainstorm 3 alternatives.) +- [ ] Are all new functions easily testable without mocking core infrastructure? +- [ ] Are there hidden dependencies that should be explicit parameters instead? + + diff --git a/.agents/rules/testing.md b/.agents/rules/testing.md new file mode 100644 index 0000000..b6cc355 --- /dev/null +++ b/.agents/rules/testing.md @@ -0,0 +1,72 @@ +1. ALWAYS perform comprehensive live testing of every single feature, fix, or refactor you implement. + +2. PROVE every new gate, seal, contract, or guard test can catch a real defect: after it + passes, show it FAILING against a deliberately wrong-but-*present* implementation — a + near-miss such as a flipped condition, an off-by-one, or the wrong-but-same-shaped value + — not merely an absent or stubbed one. A test that stays green against a plausible wrong + implementation is a proxy assertion; rewrite it until the near-miss goes red. (Red against + a stub proves only that the test fails when the code is missing, never that it fails when + the code is present and wrong — that gap is where weak assertions hide.) + +3. NEVER assert a stand-in for the property under test. On structured artifacts (YAML, JSON, + config, generated files) a substring check (`toContain`, `indexOf`), a bare count, or a + value coupled to a moving target is a proxy that passes for the wrong reason — e.g. + `toContain('V-1')` is also satisfied by `V-10`. Anchor to the exact property with a + structural or exact-equality assertion that matches it and nothing else. + +4. For every "must fail loud" criterion — a seal gate, a deny rule, a validation boundary — + ALWAYS write an explicit BREACH test that drives the failure path and asserts the exact + error, not only the happy-path success. + +5. NEVER hand-maintain the file set a source-scan security gate inspects (the no-eval / + no-persistence / no-network grep tests). Derive it from a glob over the package's runtime + source (e.g. `src/**/*.ts` excluding `*.test.*` and type-only re-export barrels), by rule + not by manual omission — a literal file array silently under-covers a new or moved module + and over-couples the gate to unrelated stubs owned by later work. Within the scan, strip + comments and string literals before matching and anchor each forbidden token to a + call/word boundary (e.g. `\beval\s*\(`) so the count is a capability check, not a + substring a comment can satisfy. + +6. ALWAYS exercise a source-scan security gate (no-eval / no-persistence / no-network) + through its OWN exported gate artifacts — the single forbidden-pattern array and the + glob-derived file set — never through a private re-declared copy or an inline literal. + (a) Its near-miss/breach tests MUST iterate the actual exported pattern array and, for + every branch of each token alternation, seed one real matching call and assert that + exact branch matches exactly once, so a typo in any single branch (e.g. `\bsendBeacon\b` + → `\bsendBecon\b`) turns a test RED. (b) The glob's exclusion of barrels MUST anchor to + exact top-level source-relative paths (`== 'index.ts'` / `== 'internal.ts'`), never a + basename match at any depth, and that anchoring MUST have its own unit assertion proving + a nested same-named module (`builders/index.ts`) is NOT excluded. A gate whose breach + test verifies a private regex copy, or whose file set silently under-covers via an + over-broad exclusion, passes for the wrong reason and cannot catch a real capability + regression (extends rules 2, 3, 5). + +7. ALWAYS cover the WIRING POINT, not only the extracted unit. When a story moves + behavior into a factory or helper and unit-tests it in isolation, the call site + that wires it in — a provider fallback (`x ?? createDefaultTelemetrySink()`), a + dispose/cleanup effect, or a new context field — is itself a deliverable and MUST + carry a test at that integration point that goes RED when the wiring is reverted or + deleted: mount the real consumer WITHOUT the injected dependency, assert the + resolved dependency's observable behavior, then tear down and assert release. + Unit-testing the extracted unit alone does NOT satisfy this. Corollary: when a + story swaps a default or fallback, grep existing suites for tests that describe the + OLD behavior and rewrite them — an assertion like `emit(event) === undefined` that + BOTH the removed and the new implementation satisfy is a proxy (rule 3) that now + passes for the wrong reason; fix it and add every touched test file to the story + File List (extends rules 2, 3). + +8. VERIFY rendered CSS behavior through the CSSOM, never through an inline-style string. + When a test asserts a keyframe animation, an `@media (prefers-reduced-motion: reduce)` + exemption, or any cascade-applied value, it MUST read the parsed rule from + `document.styleSheets` — assert a `CSSKeyframesRule` by name, or the `CSSMediaRule` + selector→declaration — never assert only an element's inline `style.animation*` / + `style.*`. jsdom's `getComputedStyle` applies neither stylesheet cascade nor `@media`, + so an inline-style assertion passes identically whether or not the backing `@keyframes` + / `@media` rule shipped, and stays green against a component that omits the rule + entirely (a proxy, rules 2-3). Before landing the test, name the wrong-but-present + implementation it must turn RED — a missing keyframe, a blanket rule that also zeros the + exempt animation, a hardcoded duration; if you cannot name one, rewrite it. And NEVER + assert a negative (zero `` rows, no export button) against a component that + structurally cannot produce the positive — the assertion is vacuously true; place it at + the integration point (the dispatch, the fallback-table renderer) where the affordance + could actually appear (extends rules 2, 3). \ No newline at end of file diff --git a/.claude/rules/Reasoning.md b/.claude/rules/Reasoning.md new file mode 100644 index 0000000..40c3ff7 --- /dev/null +++ b/.claude/rules/Reasoning.md @@ -0,0 +1,74 @@ +**REFERENTIAL CLARITY PROTOCOL** + +These rules govern every response you produce. They are constraints, not preferences. No exceptions. + +**RULE 1 — NO UNBOUND PRONOUNS** +Every "it," "this," "that," "they," "them," "those," "the thing," "the issue," and similar reference must point to an explicit noun phrase named earlier in the same paragraph — ideally the same sentence. If the referent has not been named in this response, name it instead of pronouning it. When ambiguity is possible, repeat the noun rather than relying on a pronoun. + +**RULE 2 — NO UNDEFINED SHORTHAND** +The first appearance of any acronym, initialism, abbreviation, internal name, codename, jargon term, or domain-specific shorthand must be written in full and defined. Later uses may abbreviate. If you do not know what a shorthand term expands to, do not use it — ask. + +**RULE 3 — NO UNSTATED CONTEXT (anti–curse-of-knowledge)** +Do not assume the reader shares your knowledge of: +- the system, codebase, document, person, or event being discussed +- prior context outside this specific message +- domain conventions, defaults, or background facts +Before any claim depends on such a fact, either (a) state the fact explicitly, or (b) ask the user to confirm it. Treat the reader as intelligent but uninformed about your specific context. + +**RULE 4 — ASK, DO NOT GUESS** +When a referent, abbreviation, or background fact is missing or admits multiple interpretations, ask exactly one clarifying question instead of picking a reading and proceeding. Silent disambiguation is a violation. + +**RULE 5 — NAME THINGS ON INTRODUCTION** +The first time any entity (person, system, file, function, concept, event) is introduced in a response, give it a full identifier — name, role, and any qualifier needed to distinguish it from similar entities. Pronouns and short forms are licensed only after this introduction. + +**RULE 6 - NO EDITORIALIZING** +No editorializing, no nominalized abstractions. State facts and, only when asked or clearly useful, a plain judgment in plain words. Before sending, flag any (a) abstract noun-phrase you coined to name a concept ("surgical minimalism," "cognitive overhead"), (b) "X over Y" tradeoff framing, or (c) verdict word ("defensible," "reasonable," "elegant," "clean"). For each, ask: does this give the reader a fact or action they didn't have? If no, delete it. Prefer a concrete noun or verb over an abstract one; if a judgment is worth making, say "this is fine because [concrete reason]" — never label it. + +**PRE-SEND CHECK (run before every response, without exception)** +1. Every pronoun resolves to a named antecedent inside this response. +2. Every abbreviation is expanded on first use inside this response. +3. Every reference to a person, system, document, project, or concept is identified, not presumed familiar. +4. Every ambiguity is either resolved by stated fact or surfaced as a question — never resolved silently. + +If any check fails, revise the response before sending. Brevity, fluency, and conversational tone do not override these rules. + +**CONSTRAINT-FIRST REASONING** + +Before answering any decision question: + +1. **Identify the parent goal.** What is the user actually trying to accomplish? +2. **Extract hard constraints.** What physical, logical, or contextual requirements does that goal impose? List them explicitly. +3. **Evaluate options against constraints, not heuristics.** Eliminate any option that violates a hard constraint before applying preferences or common-sense heuristics. + +**MULTI-PART INSTRUCTION COMPLIANCE** + +When a user message contains multiple distinct requests or questions: + +1. Enumerate each one before responding. +2. Address every one explicitly. Do not selectively respond to the easiest or most emotionally salient part. +3. If you cannot address one, state that explicitly rather than silently dropping it. + +**ERROR RECOVERY PROTOCOL** + +When a user corrects you: + +1. Do not default to apology + corrected answer. +2. First, identify the specific reasoning breakdown mechanistically (which step failed and why). +3. Then provide the corrected answer. +4. Do not perform reflective language ("I latched onto," "I pattern-matched") without specifying the exact constraint, instruction, or logical step that was violated. + +**SELF-AUDIT PROTOCOL** + +Before committing to any substantive claim or reasoning chain: + +1. **Source-tag the claim.** Label it: VERIFIED (checked against external source now), DEDUCED (follows from stated premises by a nameable rule), PATTERN (matches training data, feels right), or ASSUMED (treated as true without basis). If you label something DEDUCED but cannot name the logical rule, reclassify it as PATTERN. + +2. **Enumerate premises before multi-step reasoning.** List them. Tag each. Conclusions inherit the weakest tag of their premises. Do not let confident intermediate steps launder an ASSUMED premise into a DEDUCED conclusion. + +3. **Inversion test on non-trivial conclusions.** Ask: what would need to be true for the opposite to hold? If you cannot construct a coherent counter-case, flag this as a red flag (locked pattern), not a confirmation. If you can, weigh both before defaulting to whichever you generated first. + +4. **Surface ambiguity, never resolve silently.** When a question has multiple valid interpretations, name the ambiguity and your chosen interpretation before reasoning. + +5. **Ground over generating.** When tools are available, verify factual claims rather than generating from memory — especially for counterintuitive truths, statistics, and current states of affairs. + +6. **Make reasoning legible.** When the stakes are high or reasoning is complex, show your audit trail so the user can catch what you cannot. \ No newline at end of file diff --git a/.claude/rules/production-ready-code.md b/.claude/rules/production-ready-code.md new file mode 100644 index 0000000..b6015c6 --- /dev/null +++ b/.claude/rules/production-ready-code.md @@ -0,0 +1,145 @@ +# CRITICAL CONSTRAINT — ZERO TOLERANCE: Every file you touch must be left in a state of immaculate structure, naming, and clarity; dead code, duplication, unclear naming, and disorganization are treated as severity-zero defects equivalent to broken functionality — no exception, no trade-off, no shortcut. + +# Rules for Production-Ready Code + +## Purpose + +These rules ensure maintainability, safety, and developer velocity. +**MUST** rules are non-negotiable. **SHOULD** rules are strongly recommended. + +--- + +## Philosophy + +Code is read 10x more than it is written. Every decision should optimize ruthlessly for the reader. Clean code is empathy for the next person — including future-you — expressed through structure, naming, and deletion. + +- **Boy Scout Rule** — Leave every file cleaner than you found it. If you touch a file and see a stale import, a confusing name, or a dead branch, fix it in passing. This is how codebases stay clean over time instead of decaying. +- **Delete eagerly** — Commented-out code, unused variables, dead branches, and speculative utilities are rot. Git remembers; your codebase shouldn't have to. If it's not serving the current system, remove it. +- **Consistency is a first-class value** — Same problem, same solution, everywhere. If the codebase does something one way, use that way even if you know a "better" one — unless you're prepared to migrate everything. Local cleverness creates global confusion. +- **Minimal public surface area** — Expose only what consumers need. Make things private or internal by default. Every export is a contract you now maintain forever. +- **No broken windows** — Never tolerate a known mess. The moment you leave one hack in place ("we'll fix it later"), it gives permission for the next one, and decay compounds exponentially. +- **Readability over cleverness** — A one-liner using reduce with nested ternaries is worse than a five-line loop anyone can follow. If you're proud of how clever the code is, rewrite it until you're proud of how obvious it is. + +--- + +## Planning & Approach + +- **MUST** ask clarifying questions before starting complex work. +- **MUST** draft and confirm an approach before writing code. Do not code until the plan is approved. +- **SHOULD** list pros and cons when ≥ 2 viable approaches exist. +- **MUST** analyze similar parts of the codebase and ensure the plan is consistent with existing patterns, introduces minimal changes, and reuses existing code. +- **SHOULD** scope investigations narrowly. Use subagents for research to avoid polluting the main context. + +--- + +## Code Quality + +- **MUST** follow existing project conventions and patterns. Match the style, naming, and architecture already in the codebase. +- **MUST** write small, single-responsibility functions. Prefer simple, composable, testable functions over classes when classes aren't needed. +- **MUST** use domain vocabulary from the existing codebase when naming functions, variables, and types. +- **MUST** handle errors explicitly. Never swallow errors silently. Use typed errors where the language supports it. +- **MUST** validate all external inputs (API payloads, user input, environment variables) at system boundaries. +- **SHOULD** prefer pure functions over side-effecting ones where practical. +- **SHOULD NOT** add comments except for critical caveats or non-obvious "why" explanations. Write self-explanatory code instead. +- **SHOULD NOT** extract a function unless it is reused elsewhere, is the only way to unit-test otherwise untestable logic, or drastically improves readability. +- **SHOULD NOT** introduce new dependencies without justification. Prefer standard library solutions. +- **SHOULD NOT** leave `TODO`, `FIXME`, or placeholder code in production paths. + +--- + +## Architecture & Design + +- **MUST** keep changes minimal and focused. One concern per change, one purpose per PR. +- **MUST** never modify shared interfaces, schemas, or APIs without understanding downstream consumers. +- **MUST** place shared code in a shared location only if used by ≥ 2 consumers. Do not prematurely abstract. +- **SHOULD** prefer composition over inheritance. +- **SHOULD** keep coupling low between modules. Depend on interfaces/contracts, not implementations. +- **SHOULD** follow the principle of least surprise — APIs and functions should behave as their names suggest. + +### KISS (Keep It Simple, Stupid) + +- **MUST** choose the simplest solution that solves the actual problem. If a junior developer can't understand the code on first read, it's too complex. +- **SHOULD NOT** introduce abstractions, generics, or design patterns unless the current complexity demands them. Clever code is a liability; clear code is an asset. + +### YAGNI (You Aren't Gonna Need It) + +- **MUST NOT** build features, parameters, or extension points for hypothetical future requirements. Only implement what is needed right now. +- **SHOULD NOT** add configuration options, flags, or plugin architectures "just in case." Delete speculative code — it's cheaper to build later than to maintain now. + +### DRY (Don't Repeat Yourself) + +- **SHOULD** extract shared logic only when the same code appears in ≥ 2 places with the same intent. Duplication is far cheaper than the wrong abstraction. +- **SHOULD NOT** conflate coincidental similarity with true duplication. Two functions that happen to look alike but serve different concerns should remain separate. + +--- + +## Testing + +- **MUST** write tests for every new feature and every bug fix. +- **MUST** follow TDD when feasible: write a failing test first, then implement. +- **MUST** separate unit tests (pure logic, no I/O) from integration tests (DB, network, file system). +- **MUST** run the existing test suite after changes to confirm nothing is broken. +- **SHOULD** prefer integration tests over heavy mocking. Mocks should only replace truly external services. +- **SHOULD** test edge cases, realistic input, unexpected input, and value boundaries. +- **SHOULD** use strong assertions (`toEqual(expected)`) over weak ones (`toBeGreaterThan(0)`). +- **SHOULD** parameterize test inputs. Never embed unexplained magic numbers or literals. +- **SHOULD NOT** write trivial tests that can never fail for a real defect (e.g., `expect(true).toBe(true)`). +- **SHOULD NOT** test conditions already enforced by the type system. +- **SHOULD** ensure the test description matches exactly what the assertion verifies. + +--- + +## Security + +- **MUST** never commit secrets, API keys, tokens, or `.env` files. +- **MUST** validate and sanitize all user input before use in queries, commands, or output. +- **MUST** use parameterized queries. Never construct SQL or shell commands with string interpolation. +- **MUST** validate webhook signatures and authentication tokens on all protected endpoints. +- **SHOULD** apply the principle of least privilege for all service accounts, roles, and permissions. +- **SHOULD** review any AI-generated authentication, authorization, or payment code with extra scrutiny. + +--- + +## Performance + +- **SHOULD** be aware of N+1 queries, unnecessary re-renders, and unbounded loops. +- **SHOULD** use pagination for any endpoint that returns a list. +- **SHOULD NOT** prematurely optimize. Correctness and readability come first; optimize when profiling justifies it. + +--- + +## Git & Version Control + +- **MUST** use [Conventional Commits](https://www.conventionalcommits.org/) format for all commit messages: + `[optional scope]: ` + Types: `feat`, `fix`, `refactor`, `test`, `docs`, `chore`, `perf`, `ci`, `build`. +- **MUST NOT** reference Claude, Anthropic, or any AI tool in commit messages. +- **SHOULD** keep commits atomic — one logical change per commit. +- **SHOULD NOT** commit generated files, build artifacts, or `node_modules`. + +--- + +## Tooling & Verification + +- **MUST** ensure the linter passes with zero errors before considering work done. +- **MUST** ensure the type checker passes with zero errors before considering work done. +- **MUST** ensure the formatter has been applied to all new or changed files. +- **MUST** run the full build to confirm compilation succeeds. +- **SHOULD** run the single relevant test file during development for speed, then the full suite before committing. + +--- + +## Code Review Checklist (Self-Review Before Submitting) + +Before considering any change complete, verify: + +- [ ] Can you read each function and honestly, easily follow what it does? +- [ ] Are there unused parameters, imports, or dead code? +- [ ] Are there unnecessary type casts that could be moved to function signatures? +- [ ] Is cyclomatic complexity reasonable? (Flatten deeply nested if/else.) +- [ ] Could a well-known data structure or algorithm simplify the logic? +- [ ] Is each function name the best, most consistent choice? (Brainstorm 3 alternatives.) +- [ ] Are all new functions easily testable without mocking core infrastructure? +- [ ] Are there hidden dependencies that should be explicit parameters instead? + + diff --git a/.claude/rules/testing.md b/.claude/rules/testing.md new file mode 100644 index 0000000..b6cc355 --- /dev/null +++ b/.claude/rules/testing.md @@ -0,0 +1,72 @@ +1. ALWAYS perform comprehensive live testing of every single feature, fix, or refactor you implement. + +2. PROVE every new gate, seal, contract, or guard test can catch a real defect: after it + passes, show it FAILING against a deliberately wrong-but-*present* implementation — a + near-miss such as a flipped condition, an off-by-one, or the wrong-but-same-shaped value + — not merely an absent or stubbed one. A test that stays green against a plausible wrong + implementation is a proxy assertion; rewrite it until the near-miss goes red. (Red against + a stub proves only that the test fails when the code is missing, never that it fails when + the code is present and wrong — that gap is where weak assertions hide.) + +3. NEVER assert a stand-in for the property under test. On structured artifacts (YAML, JSON, + config, generated files) a substring check (`toContain`, `indexOf`), a bare count, or a + value coupled to a moving target is a proxy that passes for the wrong reason — e.g. + `toContain('V-1')` is also satisfied by `V-10`. Anchor to the exact property with a + structural or exact-equality assertion that matches it and nothing else. + +4. For every "must fail loud" criterion — a seal gate, a deny rule, a validation boundary — + ALWAYS write an explicit BREACH test that drives the failure path and asserts the exact + error, not only the happy-path success. + +5. NEVER hand-maintain the file set a source-scan security gate inspects (the no-eval / + no-persistence / no-network grep tests). Derive it from a glob over the package's runtime + source (e.g. `src/**/*.ts` excluding `*.test.*` and type-only re-export barrels), by rule + not by manual omission — a literal file array silently under-covers a new or moved module + and over-couples the gate to unrelated stubs owned by later work. Within the scan, strip + comments and string literals before matching and anchor each forbidden token to a + call/word boundary (e.g. `\beval\s*\(`) so the count is a capability check, not a + substring a comment can satisfy. + +6. ALWAYS exercise a source-scan security gate (no-eval / no-persistence / no-network) + through its OWN exported gate artifacts — the single forbidden-pattern array and the + glob-derived file set — never through a private re-declared copy or an inline literal. + (a) Its near-miss/breach tests MUST iterate the actual exported pattern array and, for + every branch of each token alternation, seed one real matching call and assert that + exact branch matches exactly once, so a typo in any single branch (e.g. `\bsendBeacon\b` + → `\bsendBecon\b`) turns a test RED. (b) The glob's exclusion of barrels MUST anchor to + exact top-level source-relative paths (`== 'index.ts'` / `== 'internal.ts'`), never a + basename match at any depth, and that anchoring MUST have its own unit assertion proving + a nested same-named module (`builders/index.ts`) is NOT excluded. A gate whose breach + test verifies a private regex copy, or whose file set silently under-covers via an + over-broad exclusion, passes for the wrong reason and cannot catch a real capability + regression (extends rules 2, 3, 5). + +7. ALWAYS cover the WIRING POINT, not only the extracted unit. When a story moves + behavior into a factory or helper and unit-tests it in isolation, the call site + that wires it in — a provider fallback (`x ?? createDefaultTelemetrySink()`), a + dispose/cleanup effect, or a new context field — is itself a deliverable and MUST + carry a test at that integration point that goes RED when the wiring is reverted or + deleted: mount the real consumer WITHOUT the injected dependency, assert the + resolved dependency's observable behavior, then tear down and assert release. + Unit-testing the extracted unit alone does NOT satisfy this. Corollary: when a + story swaps a default or fallback, grep existing suites for tests that describe the + OLD behavior and rewrite them — an assertion like `emit(event) === undefined` that + BOTH the removed and the new implementation satisfy is a proxy (rule 3) that now + passes for the wrong reason; fix it and add every touched test file to the story + File List (extends rules 2, 3). + +8. VERIFY rendered CSS behavior through the CSSOM, never through an inline-style string. + When a test asserts a keyframe animation, an `@media (prefers-reduced-motion: reduce)` + exemption, or any cascade-applied value, it MUST read the parsed rule from + `document.styleSheets` — assert a `CSSKeyframesRule` by name, or the `CSSMediaRule` + selector→declaration — never assert only an element's inline `style.animation*` / + `style.*`. jsdom's `getComputedStyle` applies neither stylesheet cascade nor `@media`, + so an inline-style assertion passes identically whether or not the backing `@keyframes` + / `@media` rule shipped, and stays green against a component that omits the rule + entirely (a proxy, rules 2-3). Before landing the test, name the wrong-but-present + implementation it must turn RED — a missing keyframe, a blanket rule that also zeros the + exempt animation, a hardcoded duration; if you cannot name one, rewrite it. And NEVER + assert a negative (zero `` rows, no export button) against a component that + structurally cannot produce the positive — the assertion is vacuously true; place it at + the integration point (the dispatch, the fallback-table renderer) where the affordance + could actually appear (extends rules 2, 3). \ No newline at end of file diff --git a/README.md b/README.md index 0420190..0e26e01 100644 --- a/README.md +++ b/README.md @@ -75,6 +75,23 @@ For CI or headless environments where the profile holds a personal access token ucode configure --profiles DEFAULT --agents claude,codex --use-pat --skip-validate --skip-upgrade ``` +### Choosing a model + +`ucode` discovers the models your workspace serves and picks a sensible default per agent. To see what is available and override the choice: + +```bash +ucode models list # every agent +ucode models list --agent pi # one agent; marks (auto) and (pinned) + +ucode pi --model system.ai.claude-fable-5 # this session only +ucode models pin pi system.ai.claude-fable-5 # every launch, until unpinned +ucode models unpin pi +``` + +Pins are saved per workspace. Claude Code is the exception: it selects by family, so use its in-session `/model` picker, or `ucode claude -- --model ` for one session. + +Because `--model` is a `ucode` flag, reaching an agent's own `--model` needs the `--` separator: `ucode opencode -- --model `. + ### MCP servers (optional) ```bash @@ -99,6 +116,9 @@ Discovered external MCP connections are listed directly. MCP auth uses a Databri | Command | Description | |---------|-------------| | `ucode status` | Show current workspace, base URLs, managed config files, and selected models | +| `ucode models list` | List the models each agent can launch with | +| `ucode models pin ` | Always launch `` with `` | +| `ucode models unpin ` | Return `` to automatic model selection | | `ucode usage` | Show AI Gateway usage summary | | `ucode revert` | Clear saved state and restore backed-up config files | | `ucode configure --dry-run` | Preview config files without writing them | diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index f44b91b..9499140 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -24,7 +24,7 @@ map_bedrock_claude_models, resolve_provider_service, ) -from ucode.state import get_provider_service, load_state, save_state +from ucode.state import get_model_override, get_provider_service, load_state, save_state from ucode.telemetry import agent_version from ucode.ui import ( console, @@ -245,12 +245,94 @@ def default_model_for_tool(tool: str, state: dict) -> str | None: return _MODULES[tool].default_model(state) +#: Agents whose model ucode pins. Claude Code selects by family env var +#: (`ANTHROPIC_DEFAULT_{OPUS,SONNET,HAIKU}_MODEL`) and ships its own `/model` +#: picker, so there is no single model to pin for it. +PINNABLE_TOOLS = ("codex", "gemini", "opencode", "copilot", "pi") + +_MAX_LISTED_MODELS = 10 + + +def available_models_for_tool(tool: str, state: dict) -> list[str]: + """Every Databricks model id ``tool`` can be launched with, from state.""" + claude_ids = state.get("claude_model_ids") or list((state.get("claude_models") or {}).values()) + codex_models = state.get("codex_models") or [] + gemini_models = state.get("gemini_models") or [] + if tool == "opencode": + buckets = state.get("opencode_models") or {} + return [m for bucket in buckets.values() for m in bucket] + sources = { + "claude": claude_ids, + "codex": codex_models, + "gemini": gemini_models, + "oss": state.get("oss_models") or [], + "external": state.get("external_models") or [], + } + return [m for source in _TOOL_DISCOVERY_SOURCES.get(tool, ()) for m in sources[source]] + + +def _model_list_hint(models: list[str]) -> str: + shown = ", ".join(models[:_MAX_LISTED_MODELS]) + extra = len(models) - _MAX_LISTED_MODELS + return f"{shown}, ... ({len(models)} total)" if extra > 0 else shown + + +def ensure_model_available(tool: str, state: dict, model: str, *, pinned: bool = False) -> None: + """Raise RuntimeError when ``model`` is not a model ``tool`` can launch. + + Pi and OpenCode also accept provider-qualified selectors (e.g. + ``databricks-claude/system.ai.claude-opus-4-8``); the bare id is validated. + """ + bare = model.split("/", 1)[1] if "/" in model else model + models = available_models_for_tool(tool, state) + if bare in models or model in models: + return + display = TOOL_SPECS[tool]["display"] + if not models: + raise RuntimeError( + f"No models discovered for {display} on this workspace. " + f"Run `ucode configure --agent {tool}` first." + ) + if pinned: + raise RuntimeError( + f"Pinned {display} model '{model}' is no longer available on this workspace " + f"(model lists refresh on every launch). Run `ucode models unpin {tool}` to return " + f"to automatic selection, or `ucode models pin {tool} ` with one of: " + f"{_model_list_hint(models)}" + ) + raise RuntimeError( + f"Model '{model}' is not available for {display} on this workspace. " + f"Available: {_model_list_hint(models)}. " + f"Run `ucode models list --agent {tool}` to see this list again." + ) + + def resolve_launch_model( tool: str, state: dict, explicit_model: str | None, + *, + strict_pin: bool = True, ) -> tuple[dict, str | None]: - model = explicit_model or default_model_for_tool(tool, state) + """Resolve the model to launch: explicit flag > persisted pin > agent default. + + A pin naming a model the workspace no longer serves is an error at launch — + never silently run a model other than the one asked for. During configure + (``strict_pin=False``) it degrades to a warning instead, since re-running + configure is how a user recovers from a rotated-out model. + """ + pinned_model = None if explicit_model else get_model_override(state, tool) + model = explicit_model or pinned_model + if model: + try: + ensure_model_available(tool, state, model, pinned=bool(pinned_model)) + except RuntimeError as exc: + if explicit_model or strict_pin: + raise + print_warning(f"{exc} Falling back to the default model.") + model = None + if not model: + model = default_model_for_tool(tool, state) if not model: raise RuntimeError( f"No models available for {tool}. Run `ucode configure` to set up your workspace." @@ -317,8 +399,8 @@ def configure_tool( return result -def launch(tool: str, state: dict, tool_args: list[str]) -> None: - _MODULES[tool].launch(state, tool_args) +def launch(tool: str, state: dict, tool_args: list[str], model: str | None = None) -> None: + _MODULES[tool].launch(state, tool_args, model) def check_gateway_endpoint(state: dict, tool: str) -> bool: @@ -328,27 +410,36 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: if tool == "opencode": return bool(state.get("opencode_models")) if tool == "codex": - return bool(state.get("codex_models")) + # A non-empty list isn't enough: codex pins the model id verbatim and + # `default_model` drops any id that doesn't parse as `gpt-`. + # Testing the list alone reports codex available and then fails at + # launch with "No models available for codex". + return codex.default_model(state) is not None if tool == "gemini": return bool(state.get("gemini_models")) if tool == "copilot": - return bool(state.get("claude_models")) or bool(state.get("codex_models")) + return ( + bool(state.get("claude_models")) + or bool(state.get("codex_models")) + or bool(state.get("external_models")) + ) if tool == "pi": return ( bool(state.get("claude_models")) or bool(state.get("codex_models")) or bool(state.get("gemini_models")) + or bool(state.get("external_models")) ) return False _TOOL_DISCOVERY_SOURCES: dict[str, tuple[str, ...]] = { "claude": ("claude",), - "opencode": ("claude", "gemini", "oss"), + "opencode": ("claude", "gemini", "oss", "external"), "codex": ("codex",), "gemini": ("gemini",), - "copilot": ("claude", "codex"), - "pi": ("claude", "codex", "gemini"), + "copilot": ("claude", "codex", "external"), + "pi": ("claude", "codex", "gemini", "external"), } @@ -392,7 +483,7 @@ def _configure_one(tool: str, state: dict, provider: str | None) -> dict: return configure_tool(tool, state, None, provider=provider, provider_models=provider_models) if tool == "codex": return configure_tool("codex", state) - state, model = resolve_launch_model(tool, state, None) + state, model = resolve_launch_model(tool, state, None, strict_pin=False) return configure_tool(tool, state, model) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 422e763..2cb144f 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -24,6 +24,7 @@ build_auth_shell_command, build_tool_base_url, get_databricks_token, + is_responses_model_id, ) from ucode.launcher import exec_or_spawn from ucode.state import mark_tool_managed, save_state @@ -51,15 +52,20 @@ def is_update_available() -> tuple[str, str] | None: def _resolve_web_search_model(state: dict) -> str | None: """Pick the model the web_search MCP server should call. Prefers an explicit override in state, otherwise the first endpoint discovered as - Responses-API-capable. Returns None if no GPT endpoint is available — - callers should skip the MCP wiring in that case.""" + Responses-API-capable. Returns None if no such endpoint is available — + callers should skip the MCP wiring in that case. + + The MCP server POSTs to `/ai-gateway/codex/v1/responses`, which only + forwards models that speak the Responses dialect. `codex_models` holds + exactly those, so the guard below is a backstop against a discovery + regression silently repointing the MCP at a chat-completions model.""" override = state.get("web_search_model") if isinstance(override, str) and override.strip(): return override.strip() codex_models = state.get("codex_models") or [] if isinstance(codex_models, list) and codex_models: first = codex_models[0] - if isinstance(first, str) and first.strip(): + if isinstance(first, str) and first.strip() and is_responses_model_id(first.strip()): return first.strip() return None @@ -634,7 +640,11 @@ def _build_claude_argv(binary: str, tool_args: list[str]) -> list[str]: return [binary, "--settings", json.dumps(merged, separators=(",", ":")), *remaining] -def launch(state: dict, tool_args: list[str]) -> None: +def launch(state: dict, tool_args: list[str], model: str | None = None) -> None: + # Claude Code selects by family env var and ships its own `/model` picker, + # so ucode pins no single model for it. `model` is accepted so every agent's + # launch() has the same signature; `ucode claude -- --model X` still works. + del model binary = SPEC["binary"] workspace = state.get("workspace") if workspace: diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 2a5c657..ce1845b 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -391,7 +391,11 @@ def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]]) return max(parsed, key=_gpt_version_key)[0] -def launch(state: dict, tool_args: list[str]) -> None: +def launch(state: dict, tool_args: list[str], model: str | None = None) -> None: + # `model` is already pinned into ~/.codex/ucode.config.toml by the + # pre-launch configure_tool call; accepted here so every agent's launch() + # has the same signature. + del model binary = SPEC["binary"] workspace = state.get("workspace") if workspace: diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 91192e4..49cce59 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -33,9 +33,10 @@ from ucode.databricks import ( TOKEN_REFRESH_INTERVAL_SECONDS, build_copilot_base_url, + build_external_serving_base_url, get_databricks_token, ) -from ucode.state import mark_tool_managed, save_state +from ucode.state import get_model_override, mark_tool_managed, save_state COPILOT_CONFIG_DIR = Path.home() / ".copilot" COPILOT_ENV_PATH = COPILOT_CONFIG_DIR / "ucode.env" @@ -71,7 +72,7 @@ def is_update_available() -> tuple[str, str] | None: def default_model(state: dict) -> str | None: - """Prefer Claude sonnet, then opus/haiku, then codex.""" + """Prefer Claude sonnet, then opus/haiku, then codex, then external.""" claude_models = state.get("claude_models") or {} for family in ("sonnet", "opus", "haiku"): if claude_models.get(family): @@ -79,13 +80,26 @@ def default_model(state: dict) -> str | None: codex_models = state.get("codex_models") or [] if codex_models: return codex_models[0] - return None + external_models = state.get("external_models") or [] + return external_models[0] if external_models else None -def render_env_overlay(workspace: str, model: str, token: str) -> dict[str, str]: +def _resolve_base_url(state: dict, model: str) -> str: + """Pick Copilot's single provider base URL for ``model``. + + External-model serving endpoints are reached via the classic + `/serving-endpoints` path (the unified MLflow gateway 404s on them); every + other model routes through the gateway's OpenAI-compatible chat-completions. + """ + if model in (state.get("external_models") or []): + return build_external_serving_base_url(state["workspace"]) + return build_copilot_base_url(state["workspace"]) + + +def render_env_overlay(state: dict, model: str, token: str) -> dict[str, str]: return { "COPILOT_PROVIDER_TYPE": "openai", - "COPILOT_PROVIDER_BASE_URL": build_copilot_base_url(workspace), + "COPILOT_PROVIDER_BASE_URL": _resolve_base_url(state, model), "COPILOT_MODEL": model, "COPILOT_PROVIDER_BEARER_TOKEN": token, "COPILOT_OFFLINE": "true", @@ -93,9 +107,9 @@ def render_env_overlay(workspace: str, model: str, token: str) -> dict[str, str] } -def build_runtime_env(workspace: str, model: str, token: str) -> dict[str, str]: +def build_runtime_env(state: dict, model: str, token: str) -> dict[str, str]: env = os.environ.copy() - env.update(render_env_overlay(workspace, model, token)) + env.update(render_env_overlay(state, model, token)) return env @@ -146,7 +160,7 @@ def write_tool_config( token = get_databricks_token( state["workspace"], state.get("profile"), force_refresh=force_refresh ) - overlay = render_env_overlay(state["workspace"], model, token) + overlay = render_env_overlay(state, model, token) existing = parse_dotenv(COPILOT_ENV_PATH) for key in LEGACY_ENV_KEYS: existing.pop(key, None) @@ -157,30 +171,32 @@ def write_tool_config( return state, token -def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> tuple[str, str]: - model = default_model(state) +def _refresh_token_once( + state: dict, model: str | None = None, *, force_refresh: bool = False +) -> tuple[str, str]: + model = model or default_model(state) if not model: raise RuntimeError("No Copilot model is available on this workspace.") _, token = write_tool_config(state, model, force_refresh=force_refresh) return model, token -def _refresh_forever(state: dict, stop_event: threading.Event) -> None: +def _refresh_forever(state: dict, model: str | None, stop_event: threading.Event) -> None: while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS): try: - _refresh_token_once(state, force_refresh=True) + _refresh_token_once(state, model, force_refresh=True) except RuntimeError: continue -def launch(state: dict, tool_args: list[str]) -> None: - model, token = _refresh_token_once(state) - env = build_runtime_env(state["workspace"], model, token) +def launch(state: dict, tool_args: list[str], model: str | None = None) -> None: + model, token = _refresh_token_once(state, model) + env = build_runtime_env(state, model, token) stop_event = threading.Event() refresher = threading.Thread( target=_refresh_forever, - args=(state, stop_event), + args=(state, model, stop_event), daemon=True, ) refresher.start() @@ -219,8 +235,8 @@ def validate_env(state: dict) -> dict[str, str]: workspace = state.get("workspace") if not workspace: raise RuntimeError("No workspace configured.") - model = default_model(state) + model = get_model_override(state, "copilot") or default_model(state) if not model: raise RuntimeError("No Copilot model is available on this workspace.") token = get_databricks_token(workspace, state.get("profile")) - return build_runtime_env(workspace, model, token) + return build_runtime_env(state, model, token) diff --git a/src/ucode/agents/gemini.py b/src/ucode/agents/gemini.py index d8499e4..638cda2 100644 --- a/src/ucode/agents/gemini.py +++ b/src/ucode/agents/gemini.py @@ -25,7 +25,7 @@ build_tool_base_url, get_databricks_token, ) -from ucode.state import mark_tool_managed, save_state +from ucode.state import get_model_override, mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version GEMINI_CONFIG_DIR = Path.home() / ".gemini" @@ -188,25 +188,27 @@ def default_model(state: dict) -> str | None: return gemini_models[0] if gemini_models else None -def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: - model = default_model(state) +def _refresh_token_once( + state: dict, model: str | None = None, *, force_refresh: bool = False +) -> str: + model = model or default_model(state) if not model: raise RuntimeError("No Gemini model is configured.") _, token = write_tool_config(state, model, force_refresh=force_refresh) return token -def _refresh_forever(state: dict, stop_event: threading.Event) -> None: +def _refresh_forever(state: dict, model: str | None, stop_event: threading.Event) -> None: while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS): try: - _refresh_token_once(state, force_refresh=True) + _refresh_token_once(state, model, force_refresh=True) except RuntimeError: continue -def launch(state: dict, tool_args: list[str]) -> None: - token = _refresh_token_once(state) - model = default_model(state) +def launch(state: dict, tool_args: list[str], model: str | None = None) -> None: + token = _refresh_token_once(state, model) + model = model or default_model(state) if not model: raise RuntimeError("No Gemini model is configured.") env = build_runtime_env(state["workspace"], model, token) @@ -214,7 +216,7 @@ def launch(state: dict, tool_args: list[str]) -> None: stop_event = threading.Event() refresher = threading.Thread( target=_refresh_forever, - args=(state, stop_event), + args=(state, model, stop_event), daemon=True, ) refresher.start() @@ -245,7 +247,7 @@ def validate_env(state: dict) -> dict[str, str]: workspace = state.get("workspace") if not workspace: raise RuntimeError("No workspace configured.") - model = default_model(state) + model = get_model_override(state, "gemini") or default_model(state) if not model: raise RuntimeError("No Gemini model is configured.") token = get_databricks_token(workspace, state.get("profile")) diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 4b614f3..00fc9d6 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -43,6 +43,7 @@ ["provider", "databricks-anthropic"], ["provider", "databricks-google"], ["provider", "databricks-oss"], + ["provider", "databricks-external"], ] @@ -52,7 +53,9 @@ def is_update_available() -> tuple[str, str] | None: def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) -> str: """Return an OpenCode model selector in provider/model form when possible.""" - if model.startswith(("databricks-anthropic/", "databricks-google/", "databricks-oss/")): + if model.startswith( + ("databricks-anthropic/", "databricks-google/", "databricks-oss/", "databricks-external/") + ): return model anthropic_models = opencode_models.get("anthropic") or [] @@ -67,6 +70,10 @@ def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) - if model in oss_models: return f"databricks-oss/{model}" + external_models = opencode_models.get("external") or [] + if model in external_models: + return f"databricks-external/{model}" + return model @@ -103,6 +110,7 @@ def render_overlay( anthropic_models = opencode_models.get("anthropic") or [] gemini_models = opencode_models.get("gemini") or [] oss_models = opencode_models.get("oss") or [] + external_models = opencode_models.get("external") or [] providers: dict = {} keys: list[list[str]] = [["model"]] @@ -148,6 +156,21 @@ def render_overlay( "models": {m: _oss_model_overlay(m, ua_header) for m in oss_models}, } keys.append(["provider", "databricks-oss"]) + if external_models: + # External-model serving endpoints (e.g. Azure OpenAI) speak OpenAI + # chat-completions on the classic `/serving-endpoints` path — the unified + # gateway 404s on them. The @ai-sdk/openai provider posts to + # `/chat/completions` with the endpoint name as the model. + providers["databricks-external"] = { + "npm": "@ai-sdk/openai", + "options": { + "baseURL": opencode_base_urls["external"], + "apiKey": token, + "headers": auth_headers, + }, + "models": {m: {"headers": ua_header} for m in external_models}, + } + keys.append(["provider", "databricks-external"]) overlay: dict = {"model": _resolve_model_selector(model, opencode_models)} if providers: @@ -184,6 +207,7 @@ def write_tool_config( "databricks-google", "databricks-openai", "databricks-oss", + "databricks-external", ): providers.pop(stale, None) merged = deep_merge_dict(existing, overlay) @@ -237,21 +261,26 @@ def default_model(state: dict) -> str | None: if gemini: return gemini[0] oss = opencode_models.get("oss") or [] - return oss[0] if oss else None + if oss: + return oss[0] + external = opencode_models.get("external") or [] + return external[0] if external else None -def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: - model = default_model(state) +def _refresh_token_once( + state: dict, model: str | None = None, *, force_refresh: bool = False +) -> str: + model = model or default_model(state) if not model: raise RuntimeError("No OpenCode model is configured.") _, token = write_tool_config(state, model, force_refresh=force_refresh) return token -def _refresh_forever(state: dict, stop_event: threading.Event) -> None: +def _refresh_forever(state: dict, model: str | None, stop_event: threading.Event) -> None: while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS): try: - _refresh_token_once(state, force_refresh=True) + _refresh_token_once(state, model, force_refresh=True) except RuntimeError: continue @@ -263,15 +292,15 @@ def build_runtime_env(token: str, state: dict | None = None) -> dict[str, str]: return env -def launch(state: dict, tool_args: list[str]) -> None: +def launch(state: dict, tool_args: list[str], model: str | None = None) -> None: """Launch opencode with background token refresh (same pattern as Gemini).""" - token = _refresh_token_once(state) + token = _refresh_token_once(state, model) env = build_runtime_env(token, state) stop_event = threading.Event() refresher = threading.Thread( target=_refresh_forever, - args=(state, stop_event), + args=(state, model, stop_event), daemon=True, ) refresher.start() diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index e7c1760..eb49743 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -68,6 +68,7 @@ "databricks-claude", "databricks-openai", "databricks-gemini", + "databricks-external", ) PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES] @@ -83,20 +84,23 @@ def is_update_available() -> tuple[str, str] | None: def _resolve_model_selector( model: str, - claude_models: dict[str, str], + claude_ids: list[str], codex_models: list[str], gemini_models: list[str], + external_models: list[str], ) -> str: """Return a Pi model selector in `/` form when possible.""" for name in PROVIDER_NAMES: if model.startswith(f"{name}/"): return model - if model in claude_models.values(): + if model in claude_ids: return f"databricks-claude/{model}" if model in codex_models: return f"databricks-openai/{model}" if model in gemini_models: return f"databricks-gemini/{model}" + if model in external_models: + return f"databricks-external/{model}" return model @@ -104,18 +108,26 @@ def render_overlay( model: str, token: str, pi_base_urls: dict[str, str], - claude_models: dict[str, str], + claude_ids: list[str], codex_models: list[str], gemini_models: list[str], + external_models: list[str], ) -> tuple[dict, list[list[str]]]: - """Return (overlay, managed_key_paths) for ~/.pi/agent/models.json.""" + """Return (overlay, managed_key_paths) for ~/.pi/agent/models.json. + + ``claude_ids`` is every Anthropic-dialect model the workspace serves, not + just the newest per family — Pi has its own model picker, so it should see + the full list. ``external_models`` are external-model serving endpoints + (e.g. Azure OpenAI), reached via the OpenAI chat-completions dialect on the + classic `/serving-endpoints` path rather than the unified gateway. + """ providers: dict = {} keys: list[list[str]] = [["model"]] # Pi expands header values that match an env var name. Our UA contains # `/` and a space so it can never collide — safe to pass as a literal. ua_headers = {"User-Agent": f"ucode/{ucode_version()} pi/{agent_version('pi')}"} - claude_ids = sorted(set(claude_models.values())) + claude_ids = sorted(set(claude_ids)) if claude_ids: providers["databricks-claude"] = { "baseUrl": pi_base_urls["claude"], @@ -150,8 +162,24 @@ def render_overlay( "models": [{"id": m} for m in gemini_models], } keys.append(["providers", "databricks-gemini"]) + if external_models: + # External-model serving endpoints speak OpenAI chat-completions on the + # classic `/serving-endpoints/{name}/invocations` path; `openai-completions` + # appends `/chat/completions` to the base URL and sends the endpoint name + # as the model. + providers["databricks-external"] = { + "baseUrl": pi_base_urls["external"], + "api": "openai-completions", + "apiKey": token, + "authHeader": True, + "headers": ua_headers, + "models": [{"id": m} for m in external_models], + } + keys.append(["providers", "databricks-external"]) overlay: dict = { - "model": _resolve_model_selector(model, claude_models, codex_models, gemini_models), + "model": _resolve_model_selector( + model, claude_ids, codex_models, gemini_models, external_models + ), } if providers: overlay["providers"] = providers @@ -164,20 +192,35 @@ def write_tool_config( token: str | None = None, *, force_refresh: bool = False, + update_settings: bool = True, + pin_model: bool = False, ) -> tuple[dict, str]: + """Write Pi's models.json, and seed settings.json unless told not to. + + ``update_settings=False`` is used by the background token refresh: it must + rewrite models.json to rotate the baked-in bearer, but must not touch the + model the user picked in Pi's own model selector. + + ``pin_model=True`` means the user named this model (`--model` or a pin), so + it wins over whatever they last chose in Pi's selector. + """ backup_existing_file(PI_CONFIG_PATH, PI_BACKUP_PATH) if token is None: token = get_databricks_token( state["workspace"], state.get("profile"), force_refresh=force_refresh ) pi_base_urls = state.get("base_urls", {}).get("pi") or build_pi_base_urls(state["workspace"]) + # State written before claude_model_ids existed only carries the family map; + # fall back to its values so a stale state still renders a usable config. + claude_ids = state.get("claude_model_ids") or list((state.get("claude_models") or {}).values()) overlay, managed_keys = render_overlay( model, token, pi_base_urls, - state.get("claude_models") or {}, + claude_ids, state.get("codex_models") or [], state.get("gemini_models") or [], + state.get("external_models") or [], ) existing = read_json_safe(PI_CONFIG_PATH) providers = existing.get("providers") @@ -186,27 +229,53 @@ def write_tool_config( providers.pop(stale, None) merged = deep_merge_dict(existing, overlay) write_json_file(PI_CONFIG_PATH, merged) - _write_settings(overlay["model"]) + if update_settings: + _write_settings(overlay["model"], overlay.get("providers") or {}, force=pin_model) state = mark_tool_managed(state, "pi", managed_keys) save_state(state) return state, token -def _write_settings(model_selector: str) -> None: - # Pin defaultProvider/defaultModel in settings.json so Pi doesn't fall - # through to an env-key-backed provider (e.g. HF_TOKEN exposing - # huggingface) in `findInitialModel` when no --model is passed. +def _settings_need_repair(existing: dict, providers: dict) -> bool: + """True when ucode should (re)write Pi's defaultProvider/defaultModel. + + Pi persists the user's model-selector choice into settings.json via + `setDefaultModelAndProvider`, so an existing pin is a user preference and + must survive. We only write when there is nothing to preserve, or when what + is there is ucode-authored residue that no longer resolves: + + - nothing pinned yet — seed it, so Pi's `findInitialModel` can't fall + through to an env-key-backed provider (e.g. HF_TOKEN exposing huggingface) + - pinned to a ucode provider we no longer register (e.g. `databricks-openai` + on a workspace with no Responses models, or a `LEGACY_PROVIDER_NAMES` entry) + - pinned to a ucode provider we do register, but to a model it no longer offers + + A pin naming a provider ucode does not manage is the user's own; leave it. + """ + provider = existing.get("defaultProvider") + model_id = existing.get("defaultModel") + if not isinstance(provider, str) or not isinstance(model_id, str) or not model_id: + return True + if provider in providers: + offered = {m["id"] for m in providers[provider].get("models") or []} + return model_id not in offered + return provider in (*PROVIDER_NAMES, *LEGACY_PROVIDER_NAMES) + + +def _write_settings(model_selector: str, providers: dict, *, force: bool = False) -> None: provider, _, model_id = model_selector.partition("/") if not model_id: return - backup_existing_file(PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH) existing = read_json_safe(PI_SETTINGS_PATH) + if not force and not _settings_need_repair(existing, providers): + return + backup_existing_file(PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH) merged = deep_merge_dict(existing, {"defaultProvider": provider, "defaultModel": model_id}) write_json_file(PI_SETTINGS_PATH, merged) def default_model(state: dict) -> str | None: - """Prefer Claude opus → sonnet → haiku; fall back to codex, gemini.""" + """Prefer Claude opus → sonnet → haiku; fall back to codex, gemini, external.""" claude_models = state.get("claude_models") or {} for family in ("opus", "sonnet", "haiku"): if claude_models.get(family): @@ -215,21 +284,40 @@ def default_model(state: dict) -> str | None: if codex_models: return codex_models[0] gemini_models = state.get("gemini_models") or [] - return gemini_models[0] if gemini_models else None + if gemini_models: + return gemini_models[0] + external_models = state.get("external_models") or [] + return external_models[0] if external_models else None -def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: - model = default_model(state) +def _refresh_token_once( + state: dict, + model: str | None = None, + *, + force_refresh: bool = False, + update_settings: bool = True, + pin_model: bool = False, +) -> str: + model = model or default_model(state) if not model: raise RuntimeError("No Pi model is available on this workspace.") - _, token = write_tool_config(state, model, force_refresh=force_refresh) + _, token = write_tool_config( + state, + model, + force_refresh=force_refresh, + update_settings=update_settings, + pin_model=pin_model, + ) return token -def _refresh_forever(state: dict, stop_event: threading.Event) -> None: +def _refresh_forever(state: dict, model: str | None, stop_event: threading.Event) -> None: while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS): try: - _refresh_token_once(state, force_refresh=True) + # Rotate the bearer baked into models.json, but leave settings.json + # alone: mid-session the user may have switched models in Pi's + # selector, and Pi persists that choice there. + _refresh_token_once(state, model, force_refresh=True, update_settings=False) except RuntimeError: continue @@ -241,14 +329,17 @@ def build_runtime_env(token: str) -> dict[str, str]: return env -def launch(state: dict, tool_args: list[str]) -> None: - token = _refresh_token_once(state) +def launch(state: dict, tool_args: list[str], model: str | None = None) -> None: + # A non-None `model` was named by the user (`--model` or a pin), so it wins + # over whatever they last selected inside Pi. None means "use the default, + # but don't disturb their selector choice". + token = _refresh_token_once(state, model, pin_model=model is not None) env = build_runtime_env(token) stop_event = threading.Event() refresher = threading.Thread( target=_refresh_forever, - args=(state, stop_event), + args=(state, model, stop_event), daemon=True, ) refresher.start() diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b2f23be..831a0df 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -9,12 +9,16 @@ from rich.panel import Panel from ucode.agents import ( + PINNABLE_TOOLS, TOOL_SPECS, + available_models_for_tool, check_gateway_endpoint, configure_selected_tools, configure_single_tool, configure_tool, + default_model_for_tool, ensure_bootstrap_dependencies, + ensure_model_available, ensure_provider_state, install_tool_binary, normalize_tool, @@ -35,6 +39,7 @@ build_shared_base_urls, discover_claude_models, discover_codex_models, + discover_external_serving_models, discover_gemini_models, discover_model_services, ensure_ai_gateway_v2, @@ -60,10 +65,12 @@ from ucode.state import ( STATE_PATH, clear_state, + get_model_override, get_provider_service, load_full_state, load_state, save_state, + set_model_override, set_provider_service, ) from ucode.tracing import configure_tracing_command @@ -91,6 +98,7 @@ "codex": ("codex", "copilot", "pi"), "gemini": ("gemini", "opencode", "pi"), "oss": ("opencode",), + "external": ("pi", "copilot", "opencode"), } @@ -105,6 +113,7 @@ def _print_discovery_diagnostics(state: dict) -> None: "codex": "Codex models", "gemini": "Gemini models", "oss": "OSS models", + "external": "External serving models", } for source, reason in reasons.items(): consumers = ", ".join(_DISCOVERY_CONSUMERS.get(source, ())) @@ -274,15 +283,19 @@ def configure_shared_state( want_gemini = fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools want_codex = fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools want_oss = fetch_all or "opencode" in tools + want_external = fetch_all or "pi" in tools or "copilot" in tools or "opencode" in tools claude_reason: str | None = None gemini_reason: str | None = None codex_reason: str | None = None oss_reason: str | None = None + external_reason: str | None = None claude_models = {} + claude_model_ids: list[str] = [] gemini_models = [] codex_models = [] oss_models = [] + external_models: list[str] = [] web_search_model: str | None = None if skip_model_discovery: # Provider mode: the agent routes through a Model Provider Service and @@ -300,30 +313,45 @@ def configure_shared_state( # empty (workspace without UC model-services, or the listing failed), fall # back to the per-family AI Gateway listing for that family only. with spinner("Fetching available models..."): - ms_claude, ms_codex, ms_gemini, ms_oss, ms_reason = discover_model_services( - workspace, token - ) + discovered = discover_model_services(workspace, token) + ms_reason = discovered.reason if want_claude: - claude_models, claude_reason = ms_claude, ms_reason - if not claude_models: - claude_models, claude_reason = discover_claude_models(workspace, token) + claude_models = discovered.claude_models + claude_model_ids = discovered.claude_ids + claude_reason = ms_reason + if not claude_model_ids: + claude_models, claude_model_ids, claude_reason = discover_claude_models( + workspace, token + ) if want_gemini: - gemini_models, gemini_reason = ms_gemini, ms_reason + gemini_models, gemini_reason = discovered.gemini_models, ms_reason if not gemini_models: gemini_models, gemini_reason = discover_gemini_models(workspace, token) if want_codex: - codex_models, codex_reason = ms_codex, ms_reason + codex_models, codex_reason = discovered.codex_models, ms_reason if not codex_models: codex_models, codex_reason = discover_codex_models(workspace, token) if want_oss: - oss_models, oss_reason = ms_oss, ms_reason + oss_models, oss_reason = discovered.oss_models, ms_reason + if want_external: + # External-model serving endpoints live outside UC model-services + # and the unified gateway, so they need their own listing call. + external_models, external_reason = discover_external_serving_models( + workspace, token + ) opencode_models: dict[str, list[str]] = {} if claude_models: + # Deliberately the family map, not claude_model_ids: opencode picks + # `anthropic[0]` as its default, and the family map is ordered + # opus → sonnet → haiku. Widening this to every Claude id would silently + # move opencode's default onto whichever id sorts first. opencode_models["anthropic"] = list(claude_models.values()) if gemini_models: opencode_models["gemini"] = gemini_models if oss_models: opencode_models["oss"] = oss_models + if external_models: + opencode_models["external"] = external_models # Merge into existing workspace state so prior tool configs are preserved. state = load_state() @@ -350,12 +378,15 @@ def configure_shared_state( else: if want_claude: state["claude_models"] = claude_models + state["claude_model_ids"] = claude_model_ids if want_gemini: state["gemini_models"] = gemini_models if want_codex: state["codex_models"] = codex_models if want_oss: state["oss_models"] = oss_models + if want_external: + state["external_models"] = external_models if fetch_all or "opencode" in tools: state["opencode_models"] = opencode_models save_state(state) @@ -370,6 +401,7 @@ def configure_shared_state( "gemini": gemini_reason, "codex": codex_reason, "oss": oss_reason, + "external": external_reason, } return state @@ -632,6 +664,11 @@ def status() -> int: provider_service = get_provider_service(state, tool) if configured and provider_service: print_kv("Model Provider Service", provider_service) + elif configured: + pinned = get_model_override(state, tool) + active = pinned or default_model_for_tool(tool, state) + if active: + print_kv("Model", f"{active} ({'pinned' if pinned else 'auto'})") print_kv("Base URL", base_url) if configured and tool in MCP_CLIENTS: tool_mcp_servers = [ @@ -725,6 +762,101 @@ def revert() -> int: app.add_typer(configure_app, name="configure", help="Configure workspace and tool settings.") mcp_app = typer.Typer(add_completion=False, no_args_is_help=True) app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ucode.") +models_app = typer.Typer(add_completion=False, no_args_is_help=True) +app.add_typer( + models_app, name="models", help="List available Databricks models and pin per-agent defaults." +) + + +def _pinnable_tool(agent: str) -> str: + """Normalize ``agent`` and reject agents ucode does not pin a model for.""" + try: + tool = normalize_tool(agent) + except (KeyError, RuntimeError) as exc: + raise RuntimeError(f"Unknown agent '{agent}'.") from exc + if tool not in PINNABLE_TOOLS: + raise RuntimeError( + f"ucode doesn't pin a model for {TOOL_SPECS[tool]['display']} — use its in-session " + f"/model picker, or `ucode {tool} -- --model ` for a single session." + ) + return tool + + +@models_app.command("list") +def models_list_cmd( + agent: Annotated[ + str | None, typer.Option("--agent", help="Show one agent instead of all.") + ] = None, +) -> None: + """List the model ids each agent can launch with, from the last discovery.""" + state = load_state() + if not state.get("workspace"): + print_err("No workspace configured. Run `ucode configure` first.") + raise typer.Exit(1) + tools = [normalize_tool(agent)] if agent else list(TOOL_SPECS) + print_section("Available models") + for tool in tools: + models = available_models_for_tool(tool, state) + print_heading(TOOL_SPECS[tool]["display"]) + if tool == "claude": + print_note( + "ucode maps opus/sonnet/haiku to this workspace's models automatically. " + "Pick one in-session with /model, or `ucode claude -- --model `." + ) + if not models: + print_note("No models discovered. Run `ucode configure --agent " + tool + "`.") + continue + pinned = get_model_override(state, tool) + default = default_model_for_tool(tool, state) + for model in models: + if model == pinned: + marker = " (pinned)" + elif not pinned and model == default: + marker = " (auto)" + else: + marker = "" + console.print(f" {model}{marker}") + print_note("Model lists refresh on every launch and on `ucode configure`.") + + +@models_app.command("pin") +def models_pin_cmd(agent: str, model: str) -> None: + """Pin the model AGENT launches with, until `ucode models unpin`.""" + try: + tool = _pinnable_tool(agent) + state = load_state() + if not state.get("workspace"): + raise RuntimeError("No workspace configured. Run `ucode configure` first.") + ensure_model_available(tool, state, model) + save_state(set_model_override(state, tool, model)) + except RuntimeError as exc: + print_err(str(exc)) + raise typer.Exit(1) from None + display = TOOL_SPECS[tool]["display"] + if get_provider_service(state, tool): + print_warning( + f"{display} currently routes through a Model Provider Service; the pin applies " + "only when launching without one." + ) + print_success(f"{display} will launch with {model}.") + print_note(f"Clear it with `ucode models unpin {tool}`.") + + +@models_app.command("unpin") +def models_unpin_cmd(agent: str) -> None: + """Return AGENT to automatic model selection.""" + try: + tool = _pinnable_tool(agent) + except RuntimeError as exc: + print_err(str(exc)) + raise typer.Exit(1) from None + state = load_state() + display = TOOL_SPECS[tool]["display"] + if not get_model_override(state, tool): + print_note(f"No model was pinned for {display}.") + return + save_state(set_model_override(state, tool, None)) + print_success(f"{display} is back to automatic model selection.") @mcp_app.command("web-search") @@ -822,7 +954,12 @@ def _auto_configure_tool(tool: str) -> None: raise RuntimeError(f"{spec['display']} validation failed — config reverted.") -def _launch_tool(tool_name: str, ctx: typer.Context, provider: str | None = None) -> None: +def _launch_tool( + tool_name: str, + ctx: typer.Context, + provider: str | None = None, + model: str | None = None, +) -> None: try: tool = normalize_tool(tool_name) existing = load_state() @@ -837,9 +974,22 @@ def _launch_tool(tool_name: str, ctx: typer.Context, provider: str | None = None if needs_auto_configure: _auto_configure_tool(tool) state = ensure_provider_state(tool) + if model and provider: + raise RuntimeError( + "Use either --model or --provider, not both. A Model Provider Service " + "routes by header and pins no Databricks model." + ) # An explicit --provider overrides the persisted choice; otherwise fall - # back to whatever `ucode configure` saved for this tool. - provider = provider or get_provider_service(state, tool) + # back to whatever `ucode configure` saved for this tool. An explicit + # --model likewise beats a persisted provider service for this session. + saved_provider = get_provider_service(state, tool) + if model and saved_provider: + print_note( + f"Ignoring saved Model Provider Service '{saved_provider}' for this session " + "(--model pins a Databricks model). Re-run `ucode configure` to change it." + ) + saved_provider = None + provider = provider or saved_provider # Validate the provider service before launching — it must exist, be a # provider type this tool can route to (e.g. claude can't use an OpenAI # or Foundry service), and, for Bedrock, expose Claude models to pin. @@ -867,8 +1017,12 @@ def _launch_tool(tool_name: str, ctx: typer.Context, provider: str | None = None # provider). Skip model resolution, which would otherwise fail when # the workspace has no matching Databricks models. resolved_model = None + requested_model = None else: - state, resolved_model = resolve_launch_model(tool, state, None) + state, resolved_model = resolve_launch_model(tool, state, model) + # Non-None only when the user named the model, so agents can tell a + # deliberate choice from "whatever the default is" (Pi's settings.json). + requested_model = resolved_model if (model or get_model_override(state, tool)) else None state = configure_tool( tool, state, resolved_model, provider=provider, provider_models=provider_models ) @@ -876,14 +1030,15 @@ def _launch_tool(tool_name: str, ctx: typer.Context, provider: str | None = None if provider: print_kv("Provider", provider) elif resolved_model: - print_kv("Model", resolved_model) + suffix = " (pinned)" if requested_model and not model else "" + print_kv("Model", f"{resolved_model}{suffix}") if tool in ("gemini", "opencode", "copilot", "pi"): print_note( f"{TOOL_SPECS[tool]['display']} token refresh is managed automatically " f"every 30 minutes while the session is running." ) print_success(f"Starting {TOOL_SPECS[tool]['display']}") - launch_agent(tool, state, ctx.args) + launch_agent(tool, state, ctx.args, model=requested_model) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None @@ -892,6 +1047,13 @@ def _launch_tool(tool_name: str, ctx: typer.Context, provider: str | None = None raise typer.Exit(130) from None +_MODEL_OPTION_HELP = ( + "Databricks model id to use for this session (see `ucode models list`). " + "Overrides any pinned model; pass before any `--` separator. Use " + "`ucode models pin ` to make it permanent." +) + + @app.command("codex", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) def codex_cmd( ctx: typer.Context, @@ -904,9 +1066,10 @@ def codex_cmd( "before any `--` separator.", ), ] = None, + model: Annotated[str | None, typer.Option("--model", help=_MODEL_OPTION_HELP)] = None, ) -> None: """Launch Codex via Databricks.""" - _launch_tool("codex", ctx, provider=provider) + _launch_tool("codex", ctx, provider=provider, model=model) @app.command("claude", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) @@ -927,29 +1090,48 @@ def claude_cmd( @app.command("gemini", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) -def gemini_cmd(ctx: typer.Context) -> None: +def gemini_cmd( + ctx: typer.Context, + model: Annotated[str | None, typer.Option("--model", help=_MODEL_OPTION_HELP)] = None, +) -> None: """Launch Gemini CLI via Databricks.""" - _launch_tool("gemini", ctx) + _launch_tool("gemini", ctx, model=model) @app.command( "opencode", context_settings={"allow_extra_args": True, "ignore_unknown_options": True} ) -def opencode_cmd(ctx: typer.Context) -> None: +def opencode_cmd( + ctx: typer.Context, + model: Annotated[str | None, typer.Option("--model", help=_MODEL_OPTION_HELP)] = None, +) -> None: """Launch OpenCode via Databricks.""" - _launch_tool("opencode", ctx) + _launch_tool("opencode", ctx, model=model) @app.command("copilot", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) -def copilot_cmd(ctx: typer.Context) -> None: +def copilot_cmd( + ctx: typer.Context, + model: Annotated[str | None, typer.Option("--model", help=_MODEL_OPTION_HELP)] = None, +) -> None: """Launch GitHub Copilot CLI via Databricks.""" - _launch_tool("copilot", ctx) + _launch_tool("copilot", ctx, model=model) @app.command("pi", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) -def pi_cmd(ctx: typer.Context) -> None: +def pi_cmd( + ctx: typer.Context, + model: Annotated[ + str | None, + typer.Option( + "--model", + help=f"{_MODEL_OPTION_HELP} Accepts provider-qualified ids like " + "databricks-claude/system.ai.claude-opus-4-8.", + ), + ] = None, +) -> None: """Launch Pi coding agent via Databricks.""" - _launch_tool("pi", ctx) + _launch_tool("pi", ctx, model=model) @configure_app.callback(invoke_without_command=True) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9f89c6b..d5d2f65 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -22,6 +22,7 @@ from concurrent.futures import ( TimeoutError as FutureTimeoutError, ) +from dataclasses import dataclass, field from pathlib import Path from typing import Literal, cast, overload from urllib import error as urllib_error @@ -1096,10 +1097,57 @@ def build_auth_shell_command( # Databricks-managed foundation models under `system.ai`. _MODEL_SERVICE_REQUIRED_PREFIX = "system.ai." -# Supported OSS chat families, matched by name substring. Add an entry to -# support a new family. +# API dialects a model-service advertises in its `supported_api_types`. A model +# can speak several — every Claude model also advertises mlflow chat-completions +# — so buckets are assigned by precedence rather than by sweeping each dialect +# independently, or a Claude model would land in both `claude` and `oss`. +_ANTHROPIC_API_TYPE = "anthropic/v1/messages" +_RESPONSES_API_TYPE = "openai/v1/responses" +_GEMINI_API_TYPE = "gemini/v1/generateContent" +_MLFLOW_CHAT_API_TYPE = "mlflow/v1/chat/completions" + +# Name rules, used only for workspaces whose model-services payload predates +# `supported_api_types`. The codex rule requires a version digit: `gpt-oss-120b` +# is an open-weight chat model served on the mlflow route, not a Responses +# endpoint, and a bare `gpt-` substring wrongly claims it. (Mirrors +# `codex._GPT_RE`; duplicated so databricks.py stays free of agent imports.) +_CLAUDE_NAME_FRAGMENT = "claude-" +_GEMINI_NAME_FRAGMENT = "gemini-" +_GPT_VERSIONED_RE = re.compile(r"gpt-\d") + +# Supported OSS chat families, matched by name substring. This is an allowlist of +# families vetted against the agents that serve them (see `_MODEL_TOKEN_LIMITS`), +# not a list of every mlflow-capable model — qwen, llama and gemma all speak the +# same dialect and are deliberately excluded. Add an entry to support a family. _OSS_MODEL_FAMILIES = ("kimi-", "glm-") + +@dataclass(frozen=True) +class DiscoveredModels: + """Workspace model ids bucketed by the API dialect each agent speaks. + + ``claude_models`` maps ``opus``/``sonnet``/``haiku`` to the newest id in that + family, for agents that select by family (Claude Code's + ``ANTHROPIC_DEFAULT_*_MODEL`` env vars). ``claude_ids`` is every + Anthropic-dialect id — including models outside those three families, e.g. + ``claude-fable-5`` — for agents like Pi that register a full model list. + It is grouped by family and newest-first within each, so no entry is + meaningfully "the default"; read ``claude_models`` for that. + + ``codex_models`` and ``gemini_models`` are newest first: consumers that need + a single default take entry ``[0]``. + + ``reason`` is None on success, else it explains why the buckets are empty. + """ + + claude_models: dict[str, str] = field(default_factory=dict) + claude_ids: list[str] = field(default_factory=list) + codex_models: list[str] = field(default_factory=list) + gemini_models: list[str] = field(default_factory=list) + oss_models: list[str] = field(default_factory=list) + reason: str | None = None + + # Per-family token limits (context window + max output tokens). These are a # property of the model + its `/ai-gateway/mlflow/v1` route (the gateway rejects # requests whose output exceeds the cap), not of any one agent — so every agent @@ -1125,11 +1173,13 @@ def model_token_limits(model_id: str) -> dict[str, int] | None: return None -def _model_service_id(service: dict) -> str | None: - """Extract the `system.ai.` id from one model-service entry. +def _model_service_entry(service: dict) -> tuple[str, list[str]] | None: + """Extract `(system.ai., supported_api_types)` from one entry. Returns None for services in any other schema, so user/internal model - services don't leak into the family buckets.""" + services don't leak into the family buckets. The api-type list is empty when + the workspace's API predates `supported_api_types`; callers read that as + "capabilities unknown" and fall back to name rules.""" name = service.get("name") if not isinstance(name, str): return None @@ -1138,7 +1188,11 @@ def _model_service_id(service: dict) -> str | None: name = name[len(_MODEL_SERVICE_NAME_PREFIX) :] if not name.startswith(_MODEL_SERVICE_REQUIRED_PREFIX): return None - return name or None + if not name: + return None + raw = service.get("supported_api_types") + api_types = [t for t in raw if isinstance(t, str)] if isinstance(raw, list) else [] + return name, api_types # The model-services metastore listing REQUIRES a bounded `page_size`: @@ -1174,17 +1228,18 @@ def list_model_services( *, page_size: int = _MODEL_SERVICES_PAGE_SIZE, max_pages: int = 100, -) -> tuple[list[str], str | None]: - """List all `system.ai.*` model ids via the UC model-services API. +) -> tuple[dict[str, list[str]], str | None]: + """List all `system.ai.*` model services and the API dialects each speaks. Pages through ``/api/2.1/unity-catalog/model-services`` (metastore scope) - with a bounded ``page_size`` (the endpoint 499s without one) and returns the - de-duplicated, sorted list of ``system.ai.`` ids. Returns - (ids, reason); reason is None on success, otherwise it describes why the - list is empty (HTTP/network error or no services). + with a bounded ``page_size`` (the endpoint 499s without one) and returns + (services, reason). ``services`` maps each ``system.ai.`` id to + its ``supported_api_types``, sorted by id; the list is empty for a workspace + whose API predates that field. ``reason`` is None on success, otherwise it + describes why the mapping is empty (HTTP/network error or no services). """ hostname = workspace_hostname(workspace) - ids: list[str] = [] + services: dict[str, list[str]] = {} page_token: str | None = None seen_tokens: set[str] = set() last_reason: str | None = None @@ -1202,9 +1257,13 @@ def list_model_services( data = cast(dict, payload) if isinstance(payload, dict) else {} for service in data.get("model_services", []): if isinstance(service, dict): - model_id = _model_service_id(service) - if model_id: - ids.append(model_id) + entry = _model_service_entry(service) + if entry: + model_id, api_types = entry + # A duplicate id across pages keeps whichever entry actually + # advertised capabilities. + if api_types or model_id not in services: + services[model_id] = api_types page_token = data.get("next_page_token") or None if not page_token: last_reason = None @@ -1213,60 +1272,108 @@ def list_model_services( break seen_tokens.add(page_token) - deduped = sorted(set(ids)) - if deduped: - return deduped, None - return [], last_reason or "model-services listing returned no models" + if services: + return dict(sorted(services.items())), None + return {}, last_reason or "model-services listing returned no models" -def discover_model_services( - workspace: str, token: str -) -> tuple[dict[str, str], list[str], list[str], list[str], str | None]: - """Discover models via UC model-services and bucket them by family name. +def is_responses_model_id(model_id: str) -> bool: + """True when ``model_id`` names a GPT model that speaks the Responses dialect. - Returns (claude_models, codex_models, gemini_models, oss_models, reason): + A name-shaped backstop for callers that hold a bare id and no capability + record. `gpt-oss-*` is an open-weight chat model served on the mlflow route, + so a bare `gpt-` substring is not sufficient — a version digit is required. + """ + return bool(_GPT_VERSIONED_RE.search(model_id)) - - ``claude_models`` maps ``opus``/``sonnet``/``haiku`` to the newest - matching ``system.ai.claude-*`` id (mirrors ``discover_claude_models``). - - ``codex_models`` is the list of ``system.ai.*gpt-*`` ids. - - ``gemini_models`` is the list of ``system.ai.*gemini-*`` ids, newest first. - - ``oss_models`` is the list of OSS-model ``system.ai.*`` ids. - ``reason`` is None on success, else explains why nothing was found. Family - bucketing is by name substring because the model-services API does not - expose per-model API dialects. +def _speaks(api_types: list[str], api_type: str, name_matches: bool) -> bool: + """True when a model advertises ``api_type``. + + Falls back to ``name_matches`` for entries whose payload carries no + ``supported_api_types`` (older workspaces), so discovery never degrades + below the name-substring behavior it replaces. """ - ids, reason = list_model_services(workspace, token) - if not ids: - return {}, [], [], [], reason + if api_types: + return api_type in api_types + return name_matches + + +def _claude_family_map(claude_ids: list[str]) -> dict[str, str]: + """Map ``opus``/``sonnet``/``haiku`` to the newest id in that family. - claude_models: dict[str, str] = {} + Ids outside those three families (e.g. ``claude-fable-5``) belong to no + family and are reachable only through the full id list — Claude Code selects + by family env var and has no slot for them. + """ + result: dict[str, str] = {} for family in ("opus", "sonnet", "haiku"): candidates = sorted( - [m for m in ids if f"claude-{family}-" in m], - reverse=True, + [m for m in claude_ids if f"claude-{family}-" in m], + key=model_version_sort_key, ) if candidates: - claude_models[family] = candidates[0] + result[family] = candidates[0] + return result - codex_models = [m for m in ids if "gpt-" in m] - gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key) - oss_models = [m for m in ids if any(family in m for family in _OSS_MODEL_FAMILIES)] +def _is_oss_model(model_id: str, api_types: list[str]) -> bool: + """True for an allowlisted OSS chat family served on the mlflow route.""" + if not any(family in model_id for family in _OSS_MODEL_FAMILIES): + return False + return _speaks(api_types, _MLFLOW_CHAT_API_TYPE, True) + + +def discover_model_services(workspace: str, token: str) -> DiscoveredModels: + """Discover models via UC model-services and bucket them by API dialect. - if not (claude_models or codex_models or gemini_models or oss_models): - sample = ", ".join(ids[:5]) - return ( - {}, - [], - [], - [], - ( + Each service advertises the dialects it speaks in ``supported_api_types``; + buckets are assigned by precedence because a model can speak several (every + Claude model also serves mlflow chat-completions). Workspaces whose API + predates that field fall back to matching the model name. + """ + services, reason = list_model_services(workspace, token) + if not services: + return DiscoveredModels(reason=reason) + + claude_ids: list[str] = [] + codex_models: list[str] = [] + gemini_models: list[str] = [] + oss_models: list[str] = [] + for model_id, api_types in services.items(): + if _speaks(api_types, _ANTHROPIC_API_TYPE, _CLAUDE_NAME_FRAGMENT in model_id): + claude_ids.append(model_id) + elif _speaks(api_types, _RESPONSES_API_TYPE, bool(_GPT_VERSIONED_RE.search(model_id))): + codex_models.append(model_id) + elif _speaks(api_types, _GEMINI_API_TYPE, _GEMINI_NAME_FRAGMENT in model_id): + gemini_models.append(model_id) + elif _is_oss_model(model_id, api_types): + oss_models.append(model_id) + + # `model_version_sort_key` groups by the non-numeric name prefix, then orders + # by version descending. For codex/gemini every id shares a prefix, so [0] is + # the newest — which is what their consumers take. Claude ids span several + # prefixes (opus/sonnet/haiku/fable), so this groups by family rather than + # ranking globally; callers wanting a default read `claude_models`. + claude_ids.sort(key=model_version_sort_key) + codex_models.sort(key=model_version_sort_key) + gemini_models.sort(key=model_version_sort_key) + + if not (claude_ids or codex_models or gemini_models or oss_models): + sample = ", ".join(list(services)[:5]) + return DiscoveredModels( + reason=( "model-services returned model ids but none matched " f"claude/gpt/gemini/oss families (got: {sample})" - ), + ) ) - return claude_models, codex_models, gemini_models, oss_models, None + return DiscoveredModels( + claude_models=_claude_family_map(claude_ids), + claude_ids=claude_ids, + codex_models=codex_models, + gemini_models=gemini_models, + oss_models=oss_models, + ) # --- MCP services (parallel to model services) ----------------------------- @@ -1777,17 +1884,21 @@ def collect_pair(has_fn, pair): return sorted(pairs), None -def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], str | None]: - """Discover Claude families on this workspace's AI Gateway. +def discover_claude_models( + workspace: str, token: str +) -> tuple[dict[str, str], list[str], str | None]: + """Discover Claude models on this workspace's AI Gateway. - Returns (models_by_family, reason). reason is None on success; otherwise it - describes why the dict is empty (HTTP error, network error, or no models + Returns (models_by_family, claude_ids, reason). ``claude_ids`` is every + Claude model the gateway serves, newest first, including ids outside the + opus/sonnet/haiku families. reason is None on success; otherwise it + describes why the result is empty (HTTP error, network error, or no models matching the expected naming convention). """ hostname = workspace_hostname(workspace) payload, reason = _http_get_json(f"https://{hostname}/ai-gateway/anthropic/v1/models", token) if payload is None: - return {}, reason + return {}, [], reason data = cast(dict, payload) if isinstance(payload, dict) else {} raw_ids = [ @@ -1796,28 +1907,25 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], if isinstance(m.get("id"), str) and not m["id"].endswith("-anthropic") ] - result: dict[str, str] = {} - for family, key in [("opus", "opus"), ("sonnet", "sonnet"), ("haiku", "haiku")]: - candidates = sorted( - [m for m in raw_ids if f"databricks-claude-{family}-" in m], - reverse=True, - ) - if candidates: - result[key] = candidates[0] - if result: - return result, None + claude_ids = sorted( + [m for m in raw_ids if _CLAUDE_NAME_FRAGMENT in m], key=model_version_sort_key + ) + result = _claude_family_map(claude_ids) + if claude_ids: + return result, claude_ids, None if not raw_ids: - return {}, "AI Gateway returned no Claude model ids" + return {}, [], "AI Gateway returned no Claude model ids" sample = ", ".join(raw_ids[:5]) - return {}, ( - "AI Gateway returned model ids but none matched " - f"`databricks-claude-{{opus,sonnet,haiku}}-*` (got: {sample})" + return ( + {}, + [], + (f"AI Gateway returned model ids but none matched `databricks-claude-*` (got: {sample})"), ) def fetch_ai_gateway_claude_models(workspace: str, token: str) -> dict[str, str]: - """Backwards-compatible wrapper that discards the diagnostic reason.""" - models, _ = discover_claude_models(workspace, token) + """Backwards-compatible wrapper that discards the ids and diagnostic reason.""" + models, _ids, _reason = discover_claude_models(workspace, token) return models @@ -1926,6 +2034,48 @@ def fetch_codex_models(workspace: str, token: str) -> list[str]: return models +# A classic serving endpoint that hosts an external model (Azure OpenAI, etc.) +# speaks the OpenAI chat-completions dialect and is reachable ONLY via the +# per-endpoint `/serving-endpoints/{name}/invocations` path — the unified AI +# Gateway 404s on it. We surface the chat-capable ones so agents that speak that +# dialect can route to them by endpoint name. +_EXTERNAL_MODEL_ENDPOINT_TYPE = "EXTERNAL_MODEL" +_CHAT_COMPLETIONS_TASK = "llm/v1/chat" + + +def discover_external_serving_models(workspace: str, token: str) -> tuple[list[str], str | None]: + """List external-model serving endpoints that speak OpenAI chat-completions. + + Reads `/api/2.0/serving-endpoints` and keeps endpoints whose ``endpoint_type`` + is ``EXTERNAL_MODEL``, whose ``task`` is ``llm/v1/chat``, and that are ready. + The returned ids are the endpoint names — exactly the ``model`` value the + `/serving-endpoints/chat/completions` route expects. Returns (names, reason); + reason is non-None when the listing failed or matched nothing. + """ + hostname = workspace_hostname(workspace) + payload, reason = _http_get_json(f"https://{hostname}/api/2.0/serving-endpoints", token) + if payload is None: + return [], reason + data = cast(dict, payload) if isinstance(payload, dict) else {} + names: list[str] = [] + for ep in data.get("endpoints") or []: + if not isinstance(ep, dict): + continue + if ep.get("endpoint_type") != _EXTERNAL_MODEL_ENDPOINT_TYPE: + continue + if ep.get("task") != _CHAT_COMPLETIONS_TASK: + continue + ready = (ep.get("state") or {}).get("ready") + if ready is not None and ready != "READY": + continue + name = ep.get("name") + if isinstance(name, str) and name: + names.append(name) + if names: + return sorted(names), None + return [], "no ready external-model chat serving endpoints found" + + def ensure_ai_gateway_v2(workspace: str, token: str) -> None: """Probe AI Gateway v2 and raise if unavailable. @@ -2109,11 +2259,20 @@ def build_tool_base_url(tool: str, workspace: str) -> str: raise RuntimeError(f"Unsupported tool '{tool}'.") +def build_external_serving_base_url(workspace: str) -> str: + # External-model serving endpoints are reached via the classic OpenAI- + # compatible path, NOT the unified AI Gateway (which 404s on them). An agent's + # openai chat-completions provider appends `/chat/completions`, so this stops + # just before that suffix; the `model` field carries the endpoint name. + return f"{workspace}/serving-endpoints" + + def build_opencode_base_urls(workspace: str) -> dict[str, str]: return { "anthropic": build_tool_base_url("claude", workspace) + "/v1", "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", "oss": f"{workspace}/ai-gateway/mlflow/v1", + "external": build_external_serving_base_url(workspace), } @@ -2133,6 +2292,7 @@ def build_pi_base_urls(workspace: str) -> dict[str, str]: "claude": build_tool_base_url("claude", workspace), "openai": build_tool_base_url("codex", workspace), "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", + "external": build_external_serving_base_url(workspace), } diff --git a/src/ucode/state.py b/src/ucode/state.py index 372af7f..0b31e1e 100644 --- a/src/ucode/state.py +++ b/src/ucode/state.py @@ -141,8 +141,10 @@ def build_agent_state(state: dict) -> dict[str, dict]: claude_model = ( claude_models.get("opus") or claude_models.get("sonnet") or claude_models.get("haiku") ) - codex_model = codex_models[0] if codex_models else None - pi_model = claude_model or codex_model or (gemini_models[0] if gemini_models else None) + codex_model = get_model_override(state, "codex") or (codex_models[0] if codex_models else None) + pi_model = get_model_override(state, "pi") or ( + claude_model or codex_model or (gemini_models[0] if gemini_models else None) + ) agents: dict[str, dict] = { "claude": { @@ -227,3 +229,31 @@ def set_provider_service(state: dict, tool: str, full_name: str | None) -> dict: else: state.pop("provider_services", None) return state + + +def get_model_override(state: dict, tool: str) -> str | None: + """Return the model pinned for ``tool`` via `ucode models pin`, if any. + + Launches use this instead of the agent's default-model heuristic, unless an + explicit ``--model`` flag overrides it. Pins live under the workspace entry, + so they never leak across workspaces. + """ + overrides = state.get("model_overrides") + if not isinstance(overrides, dict): + return None + model = overrides.get(tool) + return model if isinstance(model, str) and model else None + + +def set_model_override(state: dict, tool: str, model: str | None) -> dict: + """Persist (or clear, when ``model`` is None) the model pinned for ``tool``.""" + overrides = dict(state.get("model_overrides") or {}) + if model: + overrides[tool] = model + else: + overrides.pop(tool, None) + if overrides: + state["model_overrides"] = overrides + else: + state.pop("model_overrides", None) + return state diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index b2d3a27..8641bcc 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -247,15 +247,22 @@ def test_uses_explicit_override(self): assert claude._resolve_web_search_model({"web_search_model": "explicit"}) == "explicit" def test_falls_back_to_first_codex_model(self): - state = {"codex_models": ["m1", "m2"]} - assert claude._resolve_web_search_model(state) == "m1" + state = {"codex_models": ["system.ai.gpt-5-6", "system.ai.gpt-5-2-codex"]} + assert claude._resolve_web_search_model(state) == "system.ai.gpt-5-6" def test_returns_none_when_no_codex_models(self): assert claude._resolve_web_search_model({}) is None assert claude._resolve_web_search_model({"codex_models": []}) is None + def test_ignores_codex_model_that_is_not_responses_capable(self): + # The MCP POSTs to /ai-gateway/codex/v1/responses. If a discovery + # regression ever leaks a chat-completions model (gpt-oss-*) into + # codex_models, skip the MCP rather than wire it to a route that 400s. + state = {"codex_models": ["system.ai.gpt-oss-120b"]} + assert claude._resolve_web_search_model(state) is None + def test_override_wins_over_codex_models(self): - state = {"web_search_model": "winner", "codex_models": ["loser"]} + state = {"web_search_model": "winner", "codex_models": ["system.ai.gpt-5-6"]} assert claude._resolve_web_search_model(state) == "winner" diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index 91afbdd..7306fb3 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -7,6 +7,8 @@ from ucode.agents import copilot WS = "https://example.databricks.com" +STATE = {"workspace": WS} +STATE_WITH_EXTERNAL = {"workspace": WS, "external_models": ["azure-foundry-gpt-5-6-tera"]} class TestCopilotSpec: @@ -25,40 +27,54 @@ def test_config_path_is_ucode_env_file(self): class TestRenderEnvOverlay: def test_sets_provider_base_url(self): - env = copilot.render_env_overlay(WS, "claude-sonnet-4-6", "tok") + env = copilot.render_env_overlay(STATE, "claude-sonnet-4-6", "tok") assert env["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/ai-gateway/mlflow/v1" def test_sets_provider_type(self): - env = copilot.render_env_overlay(WS, "m", "t") + env = copilot.render_env_overlay(STATE, "m", "t") assert env["COPILOT_PROVIDER_TYPE"] == "openai" def test_sets_model(self): - env = copilot.render_env_overlay(WS, "claude-sonnet-4-6", "tok") + env = copilot.render_env_overlay(STATE, "claude-sonnet-4-6", "tok") assert env["COPILOT_MODEL"] == "claude-sonnet-4-6" def test_sets_bearer_token(self): - env = copilot.render_env_overlay(WS, "m", "tok123") + env = copilot.render_env_overlay(STATE, "m", "tok123") assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "tok123" def test_sets_offline_true(self): - env = copilot.render_env_overlay(WS, "m", "t") + env = copilot.render_env_overlay(STATE, "m", "t") assert env["COPILOT_OFFLINE"] == "true" + def test_external_model_routes_to_serving_endpoints(self): + env = copilot.render_env_overlay(STATE_WITH_EXTERNAL, "azure-foundry-gpt-5-6-tera", "tok") + assert env["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/serving-endpoints" + + def test_non_external_model_stays_on_gateway_when_external_present(self): + # A gateway model must keep the MLflow route even on a workspace that + # also serves external endpoints — the route is chosen per model. + env = copilot.render_env_overlay(STATE_WITH_EXTERNAL, "claude-sonnet-4-6", "tok") + assert env["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/ai-gateway/mlflow/v1" + class TestBuildRuntimeEnv: def test_inherits_path(self): - env = copilot.build_runtime_env(WS, "m", "t") + env = copilot.build_runtime_env(STATE, "m", "t") assert "PATH" in env def test_overrides_copilot_vars(self): - env = copilot.build_runtime_env(WS, "m", "tok") + env = copilot.build_runtime_env(STATE, "m", "tok") assert env["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/ai-gateway/mlflow/v1" assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "tok" def test_sets_oauth_token_for_mcp(self): - env = copilot.build_runtime_env(WS, "m", "tok") + env = copilot.build_runtime_env(STATE, "m", "tok") assert env["OAUTH_TOKEN"] == "tok" + def test_external_model_routes_to_serving_endpoints(self): + env = copilot.build_runtime_env(STATE_WITH_EXTERNAL, "azure-foundry-gpt-5-6-tera", "tok") + assert env["COPILOT_PROVIDER_BASE_URL"] == f"{WS}/serving-endpoints" + class TestMcpServerConfig: def test_builds_http_server_entry_with_oauth_token_env_header(self): diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 0960c64..4121c8f 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -15,6 +15,7 @@ def _base_urls() -> dict[str, str]: "anthropic": f"{WS}/ai-gateway/anthropic/v1", "gemini": f"{WS}/ai-gateway/gemini/v1beta", "oss": f"{WS}/ai-gateway/mlflow/v1", + "external": f"{WS}/serving-endpoints", } @@ -202,6 +203,30 @@ def test_prefixes_oss_model_with_provider_id(self): assert overlay["model"] == "databricks-oss/system.ai.kimi-k2-7-code" +class TestRenderOverlayExternalProvider: + def test_external_provider_uses_openai_on_serving_endpoints(self): + models = {"external": ["azure-foundry-gpt-5-6-tera"]} + overlay, _ = opencode.render_overlay( + "azure-foundry-gpt-5-6-tera", "tok", _base_urls(), models + ) + provider = overlay["provider"]["databricks-external"] + assert provider["npm"] == "@ai-sdk/openai" + assert provider["options"]["baseURL"] == f"{WS}/serving-endpoints" + assert set(provider["models"]) == {"azure-foundry-gpt-5-6-tera"} + + def test_prefixes_external_model_with_provider_id(self): + models = {"external": ["azure-foundry-gpt-5-6-tera"]} + overlay, _ = opencode.render_overlay( + "azure-foundry-gpt-5-6-tera", "tok", _base_urls(), models + ) + assert overlay["model"] == "databricks-external/azure-foundry-gpt-5-6-tera" + + def test_managed_keys_include_external_provider(self): + models = {"external": ["azure-foundry-gpt-5-6-tera"]} + _, keys = opencode.render_overlay("azure-foundry-gpt-5-6-tera", "tok", _base_urls(), models) + assert ["provider", "databricks-external"] in keys + + class TestMcpServerConfig: def test_builds_remote_server_entry_with_oauth_token_env_header(self): entry = opencode.build_mcp_server_entry(f"{WS}/api/2.0/mcp/external/github") diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 0afc5fb..cd6a6d0 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -17,15 +17,17 @@ def _base_urls() -> dict[str, str]: "claude": f"{WS}/ai-gateway/anthropic", "openai": f"{WS}/ai-gateway/codex/v1", "gemini": f"{WS}/ai-gateway/gemini/v1beta", + "external": f"{WS}/serving-endpoints", } def _empty() -> dict: """No-models input bundle for render_overlay.""" return { - "claude_models": {}, + "claude_ids": [], "codex_models": [], "gemini_models": [], + "external_models": [], } @@ -36,9 +38,10 @@ def _overlay(model: str, token: str = "tok", **kwargs): model, token, _base_urls(), - bundle["claude_models"], + bundle["claude_ids"], bundle["codex_models"], bundle["gemini_models"], + bundle["external_models"], ) @@ -64,7 +67,7 @@ def test_no_providers_when_no_models(self): assert "providers" not in overlay def test_claude_provider_uses_anthropic_messages(self): - overlay, _ = _overlay("claude-sonnet", claude_models={"sonnet": "claude-sonnet"}) + overlay, _ = _overlay("claude-sonnet", claude_ids=["claude-sonnet"]) provider = overlay["providers"]["databricks-claude"] assert provider["api"] == "anthropic-messages" assert provider["baseUrl"] == f"{WS}/ai-gateway/anthropic" @@ -84,7 +87,7 @@ def test_gemini_provider_uses_google_generative_ai(self): def test_all_three_providers_when_all_present(self): overlay, _ = _overlay( "claude-sonnet", - claude_models={"sonnet": "claude-sonnet"}, + claude_ids=["claude-sonnet"], codex_models=["gpt-5"], gemini_models=["gemini-2"], ) @@ -95,13 +98,43 @@ def test_all_three_providers_when_all_present(self): } +class TestRenderOverlayExternalProvider: + def test_external_provider_uses_openai_completions_on_serving_endpoints(self): + overlay, _ = _overlay( + "azure-foundry-gpt-5-6-tera", + external_models=["azure-foundry-gpt-5-6-tera"], + ) + provider = overlay["providers"]["databricks-external"] + assert provider["api"] == "openai-completions" + assert provider["baseUrl"] == f"{WS}/serving-endpoints" + assert provider["models"] == [{"id": "azure-foundry-gpt-5-6-tera"}] + + def test_selector_prefixes_external_model(self): + overlay, _ = _overlay( + "azure-foundry-gpt-5-6-tera", + external_models=["azure-foundry-gpt-5-6-tera"], + ) + assert overlay["model"] == "databricks-external/azure-foundry-gpt-5-6-tera" + + def test_managed_keys_include_external_provider(self): + _, keys = _overlay( + "azure-foundry-gpt-5-6-tera", + external_models=["azure-foundry-gpt-5-6-tera"], + ) + assert ["providers", "databricks-external"] in keys + + def test_no_external_provider_without_external_models(self): + overlay, _ = _overlay("gpt-5", codex_models=["gpt-5"]) + assert "databricks-external" not in overlay.get("providers", {}) + + class TestRenderOverlayUserAgent: def test_user_agent_set_on_all_three_providers(self, monkeypatch): monkeypatch.setattr(pi, "ucode_version", lambda: "0.1.0") monkeypatch.setattr(pi, "agent_version", lambda binary: "0.74.0") overlay, _ = _overlay( "claude-sonnet", - claude_models={"sonnet": "claude-sonnet"}, + claude_ids=["claude-sonnet"], codex_models=["gpt-5"], gemini_models=["gemini-2"], ) @@ -115,7 +148,7 @@ def test_claude_disables_eager_tool_input_streaming(self): # Gateway's Anthropic translator rejects per-tool # `eager_input_streaming`; this flag makes pi send the legacy beta # header instead. - overlay, _ = _overlay("claude-sonnet", claude_models={"sonnet": "claude-sonnet"}) + overlay, _ = _overlay("claude-sonnet", claude_ids=["claude-sonnet"]) compat = overlay["providers"]["databricks-claude"]["compat"] assert compat["supportsEagerToolInputStreaming"] is False @@ -132,15 +165,13 @@ def test_openai_and_gemini_have_no_compat_flags(self): class TestRenderOverlayAuthAndModels: def test_token_in_api_key(self): - overlay, _ = _overlay( - "claude-sonnet", token="mytoken", claude_models={"sonnet": "claude-sonnet"} - ) + overlay, _ = _overlay("claude-sonnet", token="mytoken", claude_ids=["claude-sonnet"]) assert overlay["providers"]["databricks-claude"]["apiKey"] == "mytoken" def test_auth_header_flag_set_on_all_providers(self): overlay, _ = _overlay( "claude-sonnet", - claude_models={"sonnet": "claude-sonnet"}, + claude_ids=["claude-sonnet"], codex_models=["gpt-5"], gemini_models=["gemini-2"], ) @@ -148,10 +179,12 @@ def test_auth_header_flag_set_on_all_providers(self): assert overlay["providers"][name]["authHeader"] is True def test_claude_models_listed(self): - claude_models = {"opus": "claude-opus", "sonnet": "claude-sonnet"} - overlay, _ = _overlay("claude-sonnet", claude_models=claude_models) + # Pi has its own picker, so it registers every Claude id the workspace + # serves — not just the newest of each opus/sonnet/haiku family. + claude_ids = ["claude-opus-4-8", "claude-opus-4-6", "claude-sonnet-5", "claude-fable-5"] + overlay, _ = _overlay("claude-sonnet-5", claude_ids=claude_ids) ids = {m["id"] for m in overlay["providers"]["databricks-claude"]["models"]} - assert ids == {"claude-opus", "claude-sonnet"} + assert ids == set(claude_ids) def test_openai_models_listed(self): overlay, _ = _overlay("gpt-5", codex_models=["gpt-5", "gpt-5-mini"]) @@ -172,7 +205,7 @@ def test_managed_keys_include_model(self): def test_managed_keys_include_each_provider_present(self): _, keys = _overlay( "claude-sonnet", - claude_models={"sonnet": "claude-sonnet"}, + claude_ids=["claude-sonnet"], codex_models=["gpt-5"], gemini_models=["gemini-2"], ) @@ -182,7 +215,7 @@ def test_managed_keys_include_each_provider_present(self): class TestRenderOverlayModelSelector: def test_prefixes_claude_model(self): - overlay, _ = _overlay("claude-sonnet", claude_models={"sonnet": "claude-sonnet"}) + overlay, _ = _overlay("claude-sonnet", claude_ids=["claude-sonnet"]) assert overlay["model"] == "databricks-claude/claude-sonnet" def test_prefixes_openai_model(self): @@ -196,7 +229,7 @@ def test_prefixes_gemini_model(self): def test_preserves_already_prefixed_model(self): overlay, _ = _overlay( "databricks-claude/claude-sonnet", - claude_models={"sonnet": "claude-sonnet"}, + claude_ids=["claude-sonnet"], ) assert overlay["model"] == "databricks-claude/claude-sonnet" @@ -281,6 +314,7 @@ def _state(self, **overrides) -> dict: "workspace": WS, "base_urls": {"pi": _base_urls()}, "claude_models": {"sonnet": "claude-sonnet"}, + "claude_model_ids": ["claude-sonnet"], "codex_models": [], "gemini_models": [], "managed_configs": {}, @@ -288,6 +322,37 @@ def _state(self, **overrides) -> dict: state.update(overrides) return state + def test_registers_every_claude_id(self, tmp_path, monkeypatch): + pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) + state = self._state( + claude_models={"sonnet": "claude-sonnet-5"}, + claude_model_ids=["claude-sonnet-5", "claude-opus-4-6", "claude-fable-5"], + ) + with ( + patch("ucode.agents.pi.get_databricks_token", return_value="tok"), + patch("ucode.agents.pi.save_state"), + ): + pi_mod.write_tool_config(state, "claude-sonnet-5") + + config = json.loads(config_file.read_text()) + ids = {m["id"] for m in config["providers"]["databricks-claude"]["models"]} + assert ids == {"claude-sonnet-5", "claude-opus-4-6", "claude-fable-5"} + + def test_falls_back_to_family_map_on_stale_state(self, tmp_path, monkeypatch): + # State written before claude_model_ids existed still renders a config. + pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) + state = self._state(claude_models={"sonnet": "claude-sonnet"}) + state.pop("claude_model_ids") + with ( + patch("ucode.agents.pi.get_databricks_token", return_value="tok"), + patch("ucode.agents.pi.save_state"), + ): + pi_mod.write_tool_config(state, "claude-sonnet") + + config = json.loads(config_file.read_text()) + ids = {m["id"] for m in config["providers"]["databricks-claude"]["models"]} + assert ids == {"claude-sonnet"} + def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatch): pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) @@ -391,6 +456,125 @@ def test_pre_existing_settings_are_backed_up_before_first_write(self, tmp_path, assert merged["defaultProvider"] == "databricks-claude" assert merged["theme"] == "Default Dark" + def _write(self, pi_mod, state, model="claude-sonnet", **kwargs): + with ( + patch("ucode.agents.pi.get_databricks_token", return_value="tok"), + patch("ucode.agents.pi.save_state"), + ): + pi_mod.write_tool_config(state, model, token="tok", **kwargs) + + def test_preserves_user_model_choice(self, tmp_path, monkeypatch): + # Pi writes settings.json when the user picks a model in its selector. + # ucode must not reset that on the next launch. + pi_mod, _, settings_file, _ = self._setup(tmp_path, monkeypatch) + settings_file.write_text( + json.dumps({"defaultProvider": "databricks-claude", "defaultModel": "claude-opus"}), + encoding="utf-8", + ) + state = self._state(claude_model_ids=["claude-sonnet", "claude-opus"]) + + self._write(pi_mod, state) + + settings = json.loads(settings_file.read_text()) + assert settings["defaultModel"] == "claude-opus" + + def test_repairs_pin_to_provider_no_longer_registered(self, tmp_path, monkeypatch): + # `databricks-openai` disappears on a workspace with no Responses models; + # a settings pin naming it would leave Pi unable to resolve a model. + pi_mod, _, settings_file, _ = self._setup(tmp_path, monkeypatch) + settings_file.write_text( + json.dumps( + {"defaultProvider": "databricks-openai", "defaultModel": "system.ai.gpt-oss-120b"} + ), + encoding="utf-8", + ) + + self._write(pi_mod, self._state()) + + settings = json.loads(settings_file.read_text()) + assert settings["defaultProvider"] == "databricks-claude" + assert settings["defaultModel"] == "claude-sonnet" + + def test_repairs_pin_to_model_no_longer_offered(self, tmp_path, monkeypatch): + pi_mod, _, settings_file, _ = self._setup(tmp_path, monkeypatch) + settings_file.write_text( + json.dumps({"defaultProvider": "databricks-claude", "defaultModel": "claude-retired"}), + encoding="utf-8", + ) + + self._write(pi_mod, self._state()) + + assert json.loads(settings_file.read_text())["defaultModel"] == "claude-sonnet" + + def test_leaves_a_non_ucode_provider_pin_alone(self, tmp_path, monkeypatch): + # Pi only writes defaultProvider on an explicit user selection, so a + # provider ucode doesn't manage is a deliberate choice. + pi_mod, _, settings_file, _ = self._setup(tmp_path, monkeypatch) + original = {"defaultProvider": "openrouter", "defaultModel": "some/model"} + settings_file.write_text(json.dumps(original), encoding="utf-8") + + self._write(pi_mod, self._state()) + + assert json.loads(settings_file.read_text()) == original + + def test_refresh_path_does_not_write_settings(self, tmp_path, monkeypatch): + pi_mod, _, settings_file, _ = self._setup(tmp_path, monkeypatch) + + self._write(pi_mod, self._state(), update_settings=False) + + assert not settings_file.exists() + + def test_refresh_thread_rotates_token_without_touching_settings(self, tmp_path, monkeypatch): + pi_mod, _, _, _ = self._setup(tmp_path, monkeypatch) + calls: list[dict] = [] + monkeypatch.setattr( + pi_mod, + "write_tool_config", + lambda state, model, **kwargs: ( + calls.append({"model": model, **kwargs}), + (state, "tok"), + )[1], + ) + monkeypatch.setattr(pi_mod, "TOKEN_REFRESH_INTERVAL_SECONDS", 0) + + class _FiresOnce: + def __init__(self): + self.waits = 0 + + def wait(self, _timeout): + self.waits += 1 + return self.waits > 1 # run the body once, then exit the loop + + pi_mod._refresh_forever(self._state(), None, _FiresOnce()) + + assert len(calls) == 1 + assert calls[0]["model"] == "claude-sonnet" # from default_model + assert calls[0]["force_refresh"] is True + assert calls[0]["update_settings"] is False + + def test_refresh_thread_keeps_the_requested_model(self, tmp_path, monkeypatch): + # A --model / pinned launch must survive the 30-minute rewrite. + pi_mod, _, _, _ = self._setup(tmp_path, monkeypatch) + seen: list[str] = [] + monkeypatch.setattr( + pi_mod, + "write_tool_config", + lambda state, model, **kwargs: (seen.append(model), (state, "tok"))[1], + ) + monkeypatch.setattr(pi_mod, "TOKEN_REFRESH_INTERVAL_SECONDS", 0) + + class _FiresOnce: + def __init__(self): + self.waits = 0 + + def wait(self, _timeout): + self.waits += 1 + return self.waits > 1 + + pi_mod._refresh_forever(self._state(), "claude-fable-5", _FiresOnce()) + + assert seen == ["claude-fable-5"] + class TestValidateAllToolsPiRollback: def test_failed_pi_validation_rolls_back_settings(self, tmp_path, monkeypatch): diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 576fc22..165d678 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -10,6 +10,7 @@ from ucode.agents import ( DEFAULT_TOOL, TOOL_SPECS, + available_models_for_tool, check_gateway_endpoint, configure_selected_tools, default_model_for_tool, @@ -93,7 +94,14 @@ def test_claude_unavailable_when_no_models(self): assert check_gateway_endpoint({}, "claude") is False def test_codex_available(self): - assert check_gateway_endpoint({"codex_models": ["model-a"]}, "codex") is True + assert check_gateway_endpoint({"codex_models": ["system.ai.gpt-5"]}, "codex") is True + + def test_codex_unavailable_when_no_model_parses_as_gpt(self): + # A non-empty list isn't enough: codex pins the id verbatim and + # default_model drops anything that isn't `gpt-`. Reporting + # available here would fail at launch instead of at configure time. + state = {"codex_models": ["system.ai.gpt-oss-120b", "system.ai.gpt-oss-20b"]} + assert check_gateway_endpoint(state, "codex") is False def test_gemini_available(self): assert check_gateway_endpoint({"gemini_models": ["gemini-2"]}, "gemini") is True @@ -189,7 +197,8 @@ def test_codex_default_model_used_when_no_explicit(self): assert model == "databricks-gpt-5" def test_explicit_model_used_when_provided(self): - _, model = resolve_launch_model("claude", {}, "my-model") + state = {"claude_models": {"sonnet": "s4"}, "claude_model_ids": ["s4", "my-model"]} + _, model = resolve_launch_model("claude", state, "my-model") assert model == "my-model" def test_default_model_used_when_no_explicit(self): @@ -201,6 +210,87 @@ def test_raises_when_no_models_available(self): with pytest.raises(RuntimeError, match="No models available"): resolve_launch_model("claude", {}, None) + def test_explicit_model_must_be_available(self): + state = {"gemini_models": ["gemini-3"]} + with pytest.raises(RuntimeError, match="not available for Gemini"): + resolve_launch_model("gemini", state, "bogus") + + def test_pin_used_when_no_explicit_model(self): + state = { + "gemini_models": ["gemini-3", "gemini-2"], + "model_overrides": {"gemini": "gemini-2"}, + } + _, model = resolve_launch_model("gemini", state, None) + assert model == "gemini-2" + + def test_explicit_model_beats_pin(self): + state = { + "gemini_models": ["gemini-3", "gemini-2"], + "model_overrides": {"gemini": "gemini-2"}, + } + _, model = resolve_launch_model("gemini", state, "gemini-3") + assert model == "gemini-3" + + def test_stale_pin_raises_at_launch(self): + state = {"gemini_models": ["gemini-3"], "model_overrides": {"gemini": "gemini-retired"}} + with pytest.raises(RuntimeError, match="ucode models unpin gemini"): + resolve_launch_model("gemini", state, None) + + def test_stale_pin_falls_back_during_configure(self): + # Re-running configure is how a user recovers, so it must not hard-fail. + state = {"gemini_models": ["gemini-3"], "model_overrides": {"gemini": "gemini-retired"}} + _, model = resolve_launch_model("gemini", state, None, strict_pin=False) + assert model == "gemini-3" + + def test_pi_accepts_provider_qualified_selector(self): + state = {"claude_model_ids": ["system.ai.claude-fable-5"]} + _, model = resolve_launch_model("pi", state, "databricks-claude/system.ai.claude-fable-5") + assert model == "databricks-claude/system.ai.claude-fable-5" + + def test_pi_external_model_passes_availability_and_launches(self): + # The reported bug: an external serving-endpoint model must survive + # `ensure_model_available` so `--model` and pins reach launch. + state = {"external_models": ["azure-foundry-gpt-5-6-tera"]} + _, model = resolve_launch_model("pi", state, "azure-foundry-gpt-5-6-tera") + assert model == "azure-foundry-gpt-5-6-tera" + + def test_copilot_external_model_passes_availability(self): + state = {"external_models": ["azure-foundry-gpt-5-6-tera"]} + _, model = resolve_launch_model("copilot", state, "azure-foundry-gpt-5-6-tera") + assert model == "azure-foundry-gpt-5-6-tera" + + +class TestAvailableModelsForTool: + def test_pi_spans_claude_codex_and_gemini(self): + state = { + "claude_model_ids": ["c1", "c2"], + "codex_models": ["g1"], + "gemini_models": ["m1"], + } + assert available_models_for_tool("pi", state) == ["c1", "c2", "g1", "m1"] + + def test_copilot_spans_claude_and_codex(self): + state = {"claude_model_ids": ["c1"], "codex_models": ["g1"], "gemini_models": ["m1"]} + assert available_models_for_tool("copilot", state) == ["c1", "g1"] + + def test_pi_includes_external_models(self): + state = {"claude_model_ids": ["c1"], "external_models": ["azure-foundry-gpt-5-6-tera"]} + assert available_models_for_tool("pi", state) == ["c1", "azure-foundry-gpt-5-6-tera"] + + def test_copilot_includes_external_models(self): + state = {"codex_models": ["g1"], "external_models": ["azure-foundry-gpt-5-6-tera"]} + assert available_models_for_tool("copilot", state) == ["g1", "azure-foundry-gpt-5-6-tera"] + + def test_opencode_flattens_its_buckets(self): + state = {"opencode_models": {"anthropic": ["c1"], "oss": ["o1"], "external": ["e1"]}} + assert sorted(available_models_for_tool("opencode", state)) == ["c1", "e1", "o1"] + + def test_falls_back_to_family_map_when_ids_absent(self): + assert available_models_for_tool("claude", {"claude_models": {"opus": "o1"}}) == ["o1"] + + def test_empty_state_yields_nothing(self): + assert available_models_for_tool("pi", {}) == [] + class TestResolveProviderModels: _STATE = {"workspace": "https://ws.databricks.com", "profile": None} diff --git a/tests/test_cli.py b/tests/test_cli.py index 3c5d404..1f34fe6 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -10,6 +10,7 @@ from typer.testing import CliRunner from ucode.cli import app +from ucode.databricks import DiscoveredModels _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") @@ -37,6 +38,8 @@ def no_state_writes(): patch("ucode.agents.claude.save_state"), patch("ucode.agents.gemini.save_state"), patch("ucode.agents.opencode.save_state"), + patch("ucode.agents.copilot.save_state"), + patch("ucode.agents.pi.save_state"), ): yield @@ -442,6 +445,131 @@ def test_skipped_when_already_configured(self): mock_auto.assert_not_called() +_PI_STATE = {**MINIMAL_STATE, "available_tools": [*TOOLS, "pi"]} + + +def _patch_pi_launch(state: dict | None = None): + """_patch_launch, but with a state where `pi` is already configured.""" + state = state or _PI_STATE + return [ + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch( + "ucode.cli.resolve_launch_model", + return_value=(state, "databricks-claude-sonnet-4"), + ), + patch("ucode.cli.configure_tool", return_value=state), + patch("ucode.cli.launch_agent"), + ] + + +class TestModelFlag: + def test_flag_reaches_resolve_launch_model(self): + p = _patch_pi_launch() + with p[0], p[1], p[2], p[3], p[4] as mock_resolve, p[5], p[6]: + result = runner.invoke(app, ["pi", "--model", "system.ai.claude-fable-5"]) + assert result.exit_code == 0, result.output + assert mock_resolve.call_args[0][2] == "system.ai.claude-fable-5" + + def test_resolved_model_reaches_launch(self): + p = _patch_pi_launch() + with p[0], p[1], p[2], p[3], p[4], p[5], p[6] as mock_launch: + result = runner.invoke(app, ["pi", "--model", "databricks-claude-sonnet-4"]) + assert result.exit_code == 0, result.output + assert mock_launch.call_args[1]["model"] == "databricks-claude-sonnet-4" + + def test_plain_launch_passes_no_model_to_the_agent(self): + # None tells pi "use your default, don't disturb the selector choice". + p = _patch_pi_launch() + with p[0], p[1], p[2], p[3], p[4], p[5], p[6] as mock_launch: + result = runner.invoke(app, ["pi"]) + assert result.exit_code == 0, result.output + assert mock_launch.call_args[1]["model"] is None + + def test_pinned_model_reaches_launch_without_a_flag(self): + state = {**_PI_STATE, "model_overrides": {"pi": "databricks-claude-sonnet-4"}} + p = _patch_pi_launch(state) + with p[0], p[1], p[2], p[3], p[4], p[5], p[6] as mock_launch: + result = runner.invoke(app, ["pi"]) + assert result.exit_code == 0, result.output + assert mock_launch.call_args[1]["model"] == "databricks-claude-sonnet-4" + assert "(pinned)" in _strip_ansi(result.output) + + def test_model_and_provider_are_mutually_exclusive(self): + patches = _patch_launch("codex") + with patches[0], patches[1], patches[2], patches[3], patches[4], patches[5], patches[6]: + result = runner.invoke(app, ["codex", "--model", "m", "--provider", "main.a.b"]) + assert result.exit_code == 1 + assert "not both" in _strip_ansi(result.output) + + def test_claude_has_no_model_flag(self): + result = runner.invoke(app, ["claude", "--help"]) + assert "--model" not in _strip_ansi(result.output) + + +class TestModelsCommands: + _STATE = { + **MINIMAL_STATE, + "claude_model_ids": ["databricks-claude-sonnet-4", "databricks-claude-fable-5"], + } + + def test_list_marks_auto_and_pinned(self): + state = {**self._STATE, "model_overrides": {"pi": "databricks-claude-fable-5"}} + with patch("ucode.cli.load_state", return_value=state): + result = runner.invoke(app, ["models", "list", "--agent", "pi"]) + assert result.exit_code == 0, result.output + out = _strip_ansi(result.output) + assert "databricks-claude-fable-5 (pinned)" in out + assert "(auto)" not in out # a pin suppresses the auto marker + + def test_list_without_pin_marks_the_default(self): + with patch("ucode.cli.load_state", return_value=self._STATE): + result = runner.invoke(app, ["models", "list", "--agent", "pi"]) + assert result.exit_code == 0, result.output + assert "databricks-claude-sonnet-4 (auto)" in _strip_ansi(result.output) + + def test_pin_persists_the_override(self): + with ( + patch("ucode.cli.load_state", return_value=dict(self._STATE)), + patch("ucode.cli.set_model_override", wraps=lambda s, t, m: s) as mock_set, + ): + result = runner.invoke(app, ["models", "pin", "pi", "databricks-claude-fable-5"]) + assert result.exit_code == 0, result.output + mock_set.assert_called_once_with(self._STATE, "pi", "databricks-claude-fable-5") + + def test_pin_rejects_an_unavailable_model(self): + with patch("ucode.cli.load_state", return_value=dict(self._STATE)): + result = runner.invoke(app, ["models", "pin", "pi", "bogus"]) + assert result.exit_code == 1 + out = _strip_ansi(result.output) + assert "not available for Pi" in out + assert "databricks-claude-fable-5" in out # the error lists what is + + def test_pin_rejects_claude(self): + with patch("ucode.cli.load_state", return_value=dict(self._STATE)): + result = runner.invoke(app, ["models", "pin", "claude", "some-model"]) + assert result.exit_code == 1 + assert "/model" in _strip_ansi(result.output) + + def test_unpin_clears_the_override(self): + state = {**self._STATE, "model_overrides": {"pi": "databricks-claude-fable-5"}} + with ( + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.set_model_override", wraps=lambda s, t, m: s) as mock_set, + ): + result = runner.invoke(app, ["models", "unpin", "pi"]) + assert result.exit_code == 0, result.output + mock_set.assert_called_once_with(state, "pi", None) + + def test_unpin_is_a_noop_when_nothing_pinned(self): + with patch("ucode.cli.load_state", return_value=dict(self._STATE)): + result = runner.invoke(app, ["models", "unpin", "pi"]) + assert result.exit_code == 0 + assert "No model was pinned" in _strip_ansi(result.output) + + class TestPassthroughArgs: @pytest.mark.parametrize( "tool,extra_args", @@ -450,7 +578,10 @@ class TestPassthroughArgs: ("claude", ["--resume"]), ("codex", ["--full-auto"]), ("gemini", ["--debug"]), - ("opencode", ["--model", "my-model"]), + # `--model` is a ucode flag now, so reaching the agent's own + # `--model` requires the `--` separator. + ("opencode", ["--", "--model", "my-model"]), + ("claude", ["--model", "my-model"]), # claude_cmd declares no --model ("claude", ["-r", "--some-flag", "value"]), ], ) @@ -468,7 +599,24 @@ def test_extra_args_forwarded(self, tool, extra_args): result = runner.invoke(app, [tool, *extra_args]) assert result.exit_code == 0, result.output forwarded = mock_launch.call_args[0][2] - assert forwarded == extra_args + assert forwarded == [a for a in extra_args if a != "--"] + + def test_ucode_model_flag_is_not_forwarded_to_the_agent(self): + # Declaring --model on the launch commands intercepts it. Before this, + # `ucode opencode --model X` reached the opencode binary verbatim. + patches = _patch_launch("opencode") + with ( + patches[0], + patches[1], + patches[2], + patches[3], + patches[4], + patches[5], + patches[6] as mock_launch, + ): + result = runner.invoke(app, ["opencode", "--model", "my-model"]) + assert result.exit_code == 0, result.output + assert mock_launch.call_args[0][2] == [] def test_no_extra_args_passes_empty_list(self): patches = _patch_launch("claude") @@ -1045,8 +1193,8 @@ def _stub_deps(monkeypatch, *, pat_token, existing_state=None): monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") monkeypatch.setattr(cli_mod, "ensure_ai_gateway_v2", lambda w, t: None) - monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], None)) - monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) + monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: DiscoveredModels()) + monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, [], None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) @@ -1117,24 +1265,30 @@ def test_uc_models_used_without_legacy_fallback(self, monkeypatch): monkeypatch.setattr( cli_mod, "discover_model_services", - lambda w, t: ( - {"opus": "system.ai.claude-opus-4-8"}, - ["system.ai.gpt-5"], - [], - [], - None, + lambda w, t: DiscoveredModels( + claude_models={"opus": "system.ai.claude-opus-4-8"}, + claude_ids=["system.ai.claude-opus-4-8", "system.ai.claude-fable-5"], + codex_models=["system.ai.gpt-5"], ), ) legacy_called: list[str] = [] monkeypatch.setattr( cli_mod, "discover_claude_models", - lambda w, t: legacy_called.append("claude") or ({}, None), + lambda w, t: legacy_called.append("claude") or ({}, [], None), ) state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") assert state["claude_models"] == {"opus": "system.ai.claude-opus-4-8"} + # The full id list is persisted for agents that register every model. + assert state["claude_model_ids"] == [ + "system.ai.claude-opus-4-8", + "system.ai.claude-fable-5", + ] + # opencode takes anthropic[0] as its default, so it stays on the + # family map — widening it here would move the default onto fable. + assert state["opencode_models"]["anthropic"] == ["system.ai.claude-opus-4-8"] assert state["codex_models"] == ["system.ai.gpt-5"] assert legacy_called == [] assert "uc_enabled" not in state @@ -1143,13 +1297,16 @@ def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): # No UC model-services: each family falls back to the legacy listing. cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") monkeypatch.setattr( - cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], "no model services") + cli_mod, + "discover_model_services", + lambda w, t: DiscoveredModels(reason="no model services"), ) monkeypatch.setattr( cli_mod, "discover_claude_models", lambda w, t: ( {"opus": "databricks-claude-opus-4-8", "sonnet": "databricks-claude-sonnet-4-6"}, + ["databricks-claude-opus-4-8", "databricks-claude-sonnet-4-6"], None, ), ) @@ -1160,6 +1317,10 @@ def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): "opus": "databricks-claude-opus-4-8", "sonnet": "databricks-claude-sonnet-4-6", } + assert state["claude_model_ids"] == [ + "databricks-claude-opus-4-8", + "databricks-claude-sonnet-4-6", + ] class TestConfigureSkipValidate: @@ -1221,8 +1382,8 @@ def _stub_external_deps(monkeypatch): monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") monkeypatch.setattr(cli_mod, "ensure_ai_gateway_v2", lambda w, t: None) - monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], None)) - monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) + monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: DiscoveredModels()) + monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, [], None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index f59aa1d..a968782 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -100,6 +100,16 @@ def test_returns_anthropic_gemini_and_oss(self): assert urls["gemini"] == f"{WS}/ai-gateway/gemini/v1beta" assert urls["oss"] == f"{WS}/ai-gateway/mlflow/v1" + def test_external_uses_classic_serving_endpoints_path(self): + assert build_opencode_base_urls(WS)["external"] == f"{WS}/serving-endpoints" + + +class TestBuildExternalServingBaseUrl: + def test_points_at_classic_serving_endpoints(self): + # External models are reached off the unified gateway; the openai + # chat-completions provider appends `/chat/completions`. + assert db_mod.build_external_serving_base_url(WS) == f"{WS}/serving-endpoints" + class TestBuildSharedBaseUrls: def test_contains_all_tools(self): @@ -129,15 +139,58 @@ def test_selects_opus_4_8_when_advertised(self, monkeypatch): } monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) - models, reason = db_mod.discover_claude_models(WS, "token") + models, ids, reason = db_mod.discover_claude_models(WS, "token") assert reason is None assert models["opus"] == "databricks-claude-opus-4-8" + assert set(ids) == { + "databricks-claude-opus-4-7", + "databricks-claude-opus-4-8", + "databricks-claude-sonnet-4-6", + } + + def test_two_digit_minor_version_beats_single_digit(self, monkeypatch): + payload = { + "data": [ + {"id": "databricks-claude-opus-4-8"}, + {"id": "databricks-claude-opus-4-10"}, + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + models, _ids, _reason = db_mod.discover_claude_models(WS, "token") -def _model_service(model_id: str) -> dict: - """A model-services entry whose `name` strips to `model_id`.""" - return {"name": f"model-services/{model_id}"} + assert models["opus"] == "databricks-claude-opus-4-10" + + def test_familyless_model_reported_in_ids_only(self, monkeypatch): + # A Claude model outside opus/sonnet/haiku has no family env var, but + # agents with their own picker (pi) must still see it. + payload = {"data": [{"id": "databricks-claude-fable-5"}]} + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, ids, reason = db_mod.discover_claude_models(WS, "token") + + assert reason is None + assert models == {} + assert ids == ["databricks-claude-fable-5"] + + +ANTHROPIC = "anthropic/v1/messages" +RESPONSES = "openai/v1/responses" +GEMINI = "gemini/v1/generateContent" +MLFLOW_CHAT = "mlflow/v1/chat/completions" + + +def _model_service(model_id: str, api_types: list[str] | None = None) -> dict: + """A model-services entry whose `name` strips to `model_id`. + + Omitting `api_types` models a workspace whose API predates + `supported_api_types`, exercising the name-rule fallback. + """ + entry: dict = {"name": f"model-services/{model_id}"} + if api_types is not None: + entry["supported_api_types"] = api_types + return entry class TestModelTokenLimits: @@ -157,7 +210,130 @@ def test_uncapped_model_returns_none(self): assert db_mod.model_token_limits("system.ai.kimi-k2-7-code") is None +class TestDiscoverModelServicesByCapability: + """Bucketing driven by each service's advertised `supported_api_types`.""" + + def test_buckets_by_advertised_api_type(self, monkeypatch): + payload = { + "model_services": [ + _model_service("system.ai.claude-opus-4-8", [MLFLOW_CHAT, ANTHROPIC]), + _model_service("system.ai.claude-sonnet-5", [MLFLOW_CHAT, ANTHROPIC]), + _model_service("system.ai.gpt-5-6", [RESPONSES]), + _model_service("system.ai.gemini-3-5-flash", [GEMINI]), + _model_service("system.ai.glm-5-2", [MLFLOW_CHAT]), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + found = db_mod.discover_model_services(WS, "token") + + assert found.reason is None + # Grouped by family prefix, newest-first within each family. + assert found.claude_ids == ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"] + assert found.codex_models == ["system.ai.gpt-5-6"] + assert found.gemini_models == ["system.ai.gemini-3-5-flash"] + assert found.oss_models == ["system.ai.glm-5-2"] + + def test_claude_ids_newest_first_within_a_family(self, monkeypatch): + payload = { + "model_services": [ + _model_service("system.ai.claude-opus-4-6", [ANTHROPIC]), + _model_service("system.ai.claude-opus-4-8", [ANTHROPIC]), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + found = db_mod.discover_model_services(WS, "token") + + assert found.claude_ids == ["system.ai.claude-opus-4-8", "system.ai.claude-opus-4-6"] + + def test_gpt_oss_is_not_a_codex_model(self, monkeypatch): + # gpt-oss-* only speaks mlflow chat-completions. Bucketing it as a + # Responses model points pi's databricks-openai provider (and claude's + # web_search MCP) at /ai-gateway/codex/v1, which rejects it. + payload = { + "model_services": [ + _model_service("system.ai.gpt-oss-120b", [MLFLOW_CHAT]), + _model_service("system.ai.gpt-oss-20b", [MLFLOW_CHAT]), + _model_service("system.ai.claude-opus-4-8", [MLFLOW_CHAT, ANTHROPIC]), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + found = db_mod.discover_model_services(WS, "token") + + assert found.codex_models == [] + # Not allowlisted for opencode either, so it lands nowhere. + assert found.oss_models == [] + assert found.claude_ids == ["system.ai.claude-opus-4-8"] + + def test_claude_model_is_not_also_bucketed_as_oss(self, monkeypatch): + # Every Claude model advertises mlflow chat-completions too; precedence + # must keep it out of the oss bucket. + payload = { + "model_services": [ + _model_service("system.ai.claude-opus-4-8", [MLFLOW_CHAT, ANTHROPIC]), + _model_service("system.ai.glm-5-2", [MLFLOW_CHAT]), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + found = db_mod.discover_model_services(WS, "token") + + assert found.oss_models == ["system.ai.glm-5-2"] + assert found.claude_ids == ["system.ai.claude-opus-4-8"] + + def test_returns_every_claude_id_and_newest_per_family(self, monkeypatch): + # claude-fable-5 belongs to no opus/sonnet/haiku family, so it is + # reachable only through the full id list. + payload = { + "model_services": [ + _model_service(f"system.ai.{m}", [ANTHROPIC]) + for m in ( + "claude-fable-5", + "claude-haiku-4-5", + "claude-opus-4-6", + "claude-opus-4-8", + "claude-sonnet-4-5", + "claude-sonnet-5", + ) + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + found = db_mod.discover_model_services(WS, "token") + + assert len(found.claude_ids) == 6 + assert "system.ai.claude-fable-5" in found.claude_ids + assert found.claude_models == { + "opus": "system.ai.claude-opus-4-8", + "sonnet": "system.ai.claude-sonnet-5", + "haiku": "system.ai.claude-haiku-4-5", + } + + def test_responses_model_included_regardless_of_name(self, monkeypatch): + # Capability wins over the name rule when the payload advertises one. + payload = {"model_services": [_model_service("system.ai.o4-turbo", [RESPONSES])]} + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + assert db_mod.discover_model_services(WS, "token").codex_models == ["system.ai.o4-turbo"] + + class TestDiscoverModelServices: + """Name-rule fallback: workspaces whose payload omits supported_api_types.""" + def test_buckets_families_by_name(self, monkeypatch): payload = { "model_services": [ @@ -176,19 +352,85 @@ def test_buckets_families_by_name(self, monkeypatch): db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) ) - claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") + found = db_mod.discover_model_services(WS, "token") - assert reason is None + assert found.reason is None # Newest opus wins; sonnet bucketed; haiku absent. - assert claude == { + assert found.claude_models == { "opus": "system.ai.claude-opus-4-8", "sonnet": "system.ai.claude-sonnet-4-6", } - assert codex == ["system.ai.gpt-5"] + assert found.codex_models == ["system.ai.gpt-5"] # Gemini ordered newest-first via the shared sort key. - assert gemini[0] == "system.ai.gemini-3-5-flash" + assert found.gemini_models[0] == "system.ai.gemini-3-5-flash" # kimi and glm are the allowlisted OSS families; llama is not. - assert oss == ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code"] + assert found.oss_models == ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code"] + + def test_gpt_oss_excluded_without_capability_data(self, monkeypatch): + # The name rule requires a version digit after `gpt-`, so a workspace + # with no supported_api_types still keeps gpt-oss out of codex. + payload = { + "model_services": [ + _model_service("system.ai.gpt-oss-120b"), + _model_service("system.ai.gpt-5-6"), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + found = db_mod.discover_model_services(WS, "token") + + assert found.codex_models == ["system.ai.gpt-5-6"] + + def test_codex_models_sorted_newest_first(self, monkeypatch): + payload = { + "model_services": [ + _model_service("system.ai.gpt-4-1"), + _model_service("system.ai.gpt-5-6"), + _model_service("system.ai.gpt-5-2-codex"), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + found = db_mod.discover_model_services(WS, "token") + + # Alphabetical order would put gpt-4-1 first; pi and copilot take [0]. + assert found.codex_models[0] == "system.ai.gpt-5-6" + + def test_two_digit_minor_version_beats_single_digit(self, monkeypatch): + # Lexicographic reverse-sort ranks `4-8` above `4-10`. + payload = { + "model_services": [ + _model_service("system.ai.claude-opus-4-8"), + _model_service("system.ai.claude-opus-4-10"), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + found = db_mod.discover_model_services(WS, "token") + + assert found.claude_models["opus"] == "system.ai.claude-opus-4-10" + + def test_single_segment_version_buckets_into_family(self, monkeypatch): + # `claude-sonnet-5` has no dotted minor; it must still be a sonnet. + payload = { + "model_services": [ + _model_service("system.ai.claude-sonnet-4-6"), + _model_service("system.ai.claude-sonnet-5"), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + found = db_mod.discover_model_services(WS, "token") + + assert found.claude_models["sonnet"] == "system.ai.claude-sonnet-5" def test_oss_allowlist_drops_unsupported_families(self, monkeypatch): # Only kimi/glm are allowlisted; other families are dropped. @@ -206,11 +448,11 @@ def test_oss_allowlist_drops_unsupported_families(self, monkeypatch): db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) ) - claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") + found = db_mod.discover_model_services(WS, "token") - assert reason is None - assert (claude, codex, gemini) == ({}, [], []) - assert oss == ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code"] + assert found.reason is None + assert (found.claude_models, found.codex_models, found.gemini_models) == ({}, [], []) + assert found.oss_models == ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code"] def test_paginates_via_next_page_token(self, monkeypatch): pages = { @@ -231,21 +473,23 @@ def fake_get(url, token, timeout=10): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - claude, codex, _, _, reason = db_mod.discover_model_services(WS, "token") + found = db_mod.discover_model_services(WS, "token") - assert reason is None - assert codex == ["system.ai.gpt-5"] - assert claude == {"opus": "system.ai.claude-opus-4-8"} + assert found.reason is None + assert found.codex_models == ["system.ai.gpt-5"] + assert found.claude_models == {"opus": "system.ai.claude-opus-4-8"} def test_http_failure_returns_reason(self, monkeypatch): monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token, timeout=10: (None, "HTTP 500 Server Error") ) - claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") + found = db_mod.discover_model_services(WS, "token") - assert (claude, codex, gemini, oss) == ({}, [], [], []) - assert reason == "HTTP 500 Server Error" + assert found.claude_models == {} + assert (found.claude_ids, found.codex_models, found.gemini_models) == ([], [], []) + assert found.oss_models == [] + assert found.reason == "HTTP 500 Server Error" def test_no_matching_families_reports_sample(self, monkeypatch): payload = {"model_services": [_model_service("system.ai.llama-4-maverick")]} @@ -253,10 +497,12 @@ def test_no_matching_families_reports_sample(self, monkeypatch): db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) ) - claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") + found = db_mod.discover_model_services(WS, "token") - assert (claude, codex, gemini, oss) == ({}, [], [], []) - assert reason is not None and "llama-4-maverick" in reason + assert found.claude_models == {} + assert (found.claude_ids, found.codex_models, found.gemini_models) == ([], [], []) + assert found.oss_models == [] + assert found.reason is not None and "llama-4-maverick" in found.reason def test_ignores_non_system_ai_schemas(self, monkeypatch): # The metastore listing returns services from every schema; only @@ -274,13 +520,13 @@ def test_ignores_non_system_ai_schemas(self, monkeypatch): db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) ) - claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") + found = db_mod.discover_model_services(WS, "token") - assert reason is None - assert codex == ["system.ai.gpt-5"] - assert claude == {} # temp.erni.claude-* must not be bucketed - assert gemini == [] - assert oss == [] + assert found.reason is None + assert found.codex_models == ["system.ai.gpt-5"] + assert found.claude_models == {} # temp.erni.claude-* must not be bucketed + assert found.gemini_models == [] + assert found.oss_models == [] def test_requests_bounded_page_size(self, monkeypatch): # The endpoint 499s without a bounded page_size, so every request must @@ -289,16 +535,31 @@ def test_requests_bounded_page_size(self, monkeypatch): def fake_get(url, token, timeout=10): urls.append(url) - return {"model_services": [_model_service("system.ai.gpt-5")]}, None + return {"model_services": [_model_service("system.ai.gpt-5", [RESPONSES])]}, None monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - ids, reason = db_mod.list_model_services(WS, "token") + services, reason = db_mod.list_model_services(WS, "token") - assert ids == ["system.ai.gpt-5"] + assert services == {"system.ai.gpt-5": [RESPONSES]} assert reason is None assert all("page_size=" in u for u in urls) + def test_missing_api_types_yields_empty_capability_list(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token, timeout=10: ( + {"model_services": [_model_service("system.ai.x")]}, + None, + ), + ) + + services, reason = db_mod.list_model_services(WS, "token") + + assert services == {"system.ai.x": []} + assert reason is None + def test_retries_page_before_giving_up(self, monkeypatch): payload = {"model_services": [_model_service("system.ai.gpt-5")]} calls = {"n": 0} @@ -311,13 +572,93 @@ def flaky_get(url, token, timeout=10): monkeypatch.setattr(db_mod, "_http_get_json", flaky_get) - ids, reason = db_mod.list_model_services(WS, "token") + services, reason = db_mod.list_model_services(WS, "token") assert reason is None - assert ids == ["system.ai.gpt-5"] + assert list(services) == ["system.ai.gpt-5"] assert calls["n"] == 3 # two failures, third succeeds +def _serving_endpoint( + name: str, + *, + endpoint_type: str = "EXTERNAL_MODEL", + task: str = "llm/v1/chat", + ready: str | None = "READY", +) -> dict: + """A `/api/2.0/serving-endpoints` list entry matching the live API shape.""" + entry: dict = {"name": name, "endpoint_type": endpoint_type, "task": task} + if ready is not None: + entry["state"] = {"ready": ready} + return entry + + +class TestDiscoverExternalServingModels: + """Surfacing external-model chat serving endpoints reached via /serving-endpoints.""" + + def test_surfaces_only_ready_external_chat_endpoints(self, monkeypatch): + payload = { + "endpoints": [ + _serving_endpoint("azure-foundry-gpt-5-6-tera"), + _serving_endpoint("azure-foundry-gpt-5-5"), + # Same task, but a Databricks-managed foundation model — excluded. + _serving_endpoint("databricks-gpt-oss-120b", endpoint_type="FOUNDATION_MODEL_API"), + # A ready chat endpoint of a non-external type (a user's own custom + # model). Not reached by the OpenAI-external route — excluded. This + # pins the filter to `== EXTERNAL_MODEL`, not merely `!= foundation`. + _serving_endpoint("my-custom-chat", endpoint_type="CUSTOM_MODEL"), + # External, but an embeddings endpoint — excluded. + _serving_endpoint("azure-foundry-embed", task="llm/v1/embeddings"), + # External chat, but not ready — excluded. + _serving_endpoint("azure-foundry-warming", ready="NOT_READY"), + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_external_serving_models(WS, "token") + + assert reason is None + assert models == ["azure-foundry-gpt-5-5", "azure-foundry-gpt-5-6-tera"] + + def test_foundation_endpoint_is_excluded_near_miss(self, monkeypatch): + # Near-miss: an entry identical to a surfaced one except endpoint_type. + payload = {"endpoints": [_serving_endpoint("x", endpoint_type="FOUNDATION_MODEL_API")]} + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_external_serving_models(WS, "token") + + assert models == [] + assert reason == "no ready external-model chat serving endpoints found" + + def test_non_chat_task_is_excluded_near_miss(self, monkeypatch): + # Near-miss: external endpoint that speaks a non-chat task. + payload = {"endpoints": [_serving_endpoint("x", task="llm/v1/completions")]} + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, _ = db_mod.discover_external_serving_models(WS, "token") + + assert models == [] + + def test_missing_state_is_treated_as_available(self, monkeypatch): + payload = {"endpoints": [_serving_endpoint("x", ready=None)]} + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_external_serving_models(WS, "token") + + assert models == ["x"] + assert reason is None + + def test_http_failure_returns_reason(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token: (None, "HTTP 500 Server Error") + ) + + models, reason = db_mod.discover_external_serving_models(WS, "token") + + assert models == [] + assert reason == "HTTP 500 Server Error" + + class TestListModelProviderServices: _PAYLOAD = { "model_provider_services": [ diff --git a/tests/test_e2e_uc.py b/tests/test_e2e_uc.py index c716dcb..69b9daa 100644 --- a/tests/test_e2e_uc.py +++ b/tests/test_e2e_uc.py @@ -17,18 +17,20 @@ from ucode.cli import configure_shared_state from ucode.databricks import ( discover_model_services, + is_responses_model_id, list_mcp_services, ) from ucode.state import load_state def _has_uc_models(workspace: str, token: str) -> bool: - claude, codex, gemini, oss, _reason = discover_model_services(workspace, token) - return bool(claude or codex or gemini or oss) + found = discover_model_services(workspace, token) + return bool(found.claude_ids or found.codex_models or found.gemini_models or found.oss_models) def _all_resolved_model_ids(state: dict) -> list[str]: ids: list[str] = list((state.get("claude_models") or {}).values()) + ids += state.get("claude_model_ids") or [] ids += state.get("codex_models") or [] ids += state.get("gemini_models") or [] ids += state.get("oss_models") or [] @@ -43,18 +45,19 @@ def _all_resolved_model_ids(state: dict) -> list[str]: class TestDiscoverModelServicesE2E: def test_returns_only_system_ai_models(self, e2e_workspace, e2e_token): - claude, codex, gemini, oss, reason = discover_model_services(e2e_workspace, e2e_token) - if not (claude or codex or gemini or oss): - pytest.skip(f"No system.ai.* model services on workspace: {reason}") + found = discover_model_services(e2e_workspace, e2e_token) + if not _has_uc_models(e2e_workspace, e2e_token): + pytest.skip(f"No system.ai.* model services on workspace: {found.reason}") non_system = sorted( { m for m in _all_resolved_model_ids( { - "claude_models": claude, - "codex_models": codex, - "gemini_models": gemini, - "oss_models": oss, + "claude_models": found.claude_models, + "claude_model_ids": found.claude_ids, + "codex_models": found.codex_models, + "gemini_models": found.gemini_models, + "oss_models": found.oss_models, } ) if not m.startswith("system.ai.") @@ -62,6 +65,21 @@ def test_returns_only_system_ai_models(self, e2e_workspace, e2e_token): ) assert not non_system, f"Non-system.ai entries leaked through: {non_system[:5]}" + def test_codex_bucket_holds_only_responses_models(self, e2e_workspace, e2e_token): + # `gpt-oss-*` speaks mlflow chat-completions, not the Responses dialect + # codex and pi's databricks-openai provider route to. + found = discover_model_services(e2e_workspace, e2e_token) + if not found.codex_models: + pytest.skip("No Responses-capable models on this workspace.") + bad = [m for m in found.codex_models if not is_responses_model_id(m)] + assert not bad, f"Non-Responses models bucketed as codex: {bad}" + + def test_claude_family_map_is_a_subset_of_claude_ids(self, e2e_workspace, e2e_token): + found = discover_model_services(e2e_workspace, e2e_token) + if not found.claude_ids: + pytest.skip("No Claude models on this workspace.") + assert set(found.claude_models.values()) <= set(found.claude_ids) + class TestListMcpServicesE2E: def test_returns_only_system_ai_mcp_services(self, e2e_workspace, e2e_token): diff --git a/tests/test_state.py b/tests/test_state.py index 0b1a67f..5602309 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -12,12 +12,14 @@ STATE_VERSION, build_agent_state, clear_state, + get_model_override, get_provider_service, hydrate_state, load_full_state, load_state, mark_tool_managed, save_state, + set_model_override, set_provider_service, ) @@ -141,6 +143,46 @@ def test_clear_when_no_state_is_noop(self): clear_state() # should not raise +class TestModelOverride: + def test_get_returns_none_when_unset(self): + assert get_model_override({}, "pi") is None + assert get_model_override({"model_overrides": {}}, "pi") is None + + def test_set_and_get_roundtrip(self): + state = set_model_override({}, "pi", "system.ai.claude-fable-5") + assert state["model_overrides"]["pi"] == "system.ai.claude-fable-5" + assert get_model_override(state, "pi") == "system.ai.claude-fable-5" + assert get_model_override(state, "gemini") is None + + def test_set_none_clears_entry_and_key(self): + state = set_model_override({}, "pi", "m1") + state = set_model_override(state, "pi", None) + assert get_model_override(state, "pi") is None + # Drop the empty container entirely rather than leaving {}. + assert "model_overrides" not in state + + def test_clearing_one_tool_keeps_the_other(self): + state = set_model_override({}, "pi", "m1") + state = set_model_override(state, "gemini", "m2") + state = set_model_override(state, "pi", None) + assert get_model_override(state, "pi") is None + assert get_model_override(state, "gemini") == "m2" + + def test_survives_hydrate_state(self): + state = hydrate_state({"workspace": FAKE_WS, "model_overrides": {"pi": "m1"}}) + assert get_model_override(state, "pi") == "m1" + + def test_build_agent_state_reports_the_pin(self): + state = hydrate_state( + { + "workspace": FAKE_WS, + "claude_models": {"opus": "o1"}, + "model_overrides": {"pi": "system.ai.claude-fable-5"}, + } + ) + assert state["agents"]["pi"]["model"] == "system.ai.claude-fable-5" + + class TestProviderService: def test_get_returns_none_when_unset(self): assert get_provider_service({}, "claude") is None