diff --git a/README.md b/README.md index 701325cb..c4ecbc71 100644 --- a/README.md +++ b/README.md @@ -348,9 +348,9 @@ Use the tag shown by `ollama list` with the `ollama/` prefix. `OLLAMA_BASE_URL` In **Admin UI → Model Config**, add **Fallback Models** in the order FCC should try them if the selected model fails. Each fallback has its own reasoning effort: **Inherit**, **Off**, **From client**, **Low**, **Medium**, **High**, **X-High**, or **Max**. -Choose a model and effort in the bottom row, then click **+**. Edit either dropdown in place, drag the handle to change the order, or click the red trash button to remove a fallback. +Choose a model, effort, and API key number in the bottom row, then click **+**. **Key 1**, **Key 2**, and so on refer to the current numbered keys in that provider's settings. Existing fallbacks use **Key 1**. Edit the selections in place, drag the handle to change the order, or click the red trash button to remove a fallback. -**Inherit** keeps the original request's resolved reasoning setting. **From client** uses the effort or token budget sent by your coding agent, bypassing the global or tier setting. A fallback matching the selected provider/model is skipped. +**Inherit** keeps the original request's resolved reasoning setting. **From client** uses the effort or token budget sent by your coding agent, bypassing the global or tier setting. A fallback matching the selected provider, model, and key is skipped. You can select the same model with another key number. diff --git a/e2e/conftest.py b/e2e/conftest.py index 4e0bf652..2d633428 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -277,7 +277,7 @@ def admin_base_url( } async def fixture_provider( - provider_id: str, _settings: Settings, _admission_registry + provider_id: str, _settings: Settings, _admission_registry, _api_key_number ) -> BaseProvider: if provider_id not in providers: raise AssertionError(f"Missing browser fixture provider: {provider_id}") @@ -288,7 +288,7 @@ def admin_base_url( runtime_factory=lambda snapshot, admission_registry: ProviderRuntime( snapshot, admission_registry, - dict(providers), + {(provider_id, 1): provider for provider_id, provider in providers.items()}, provider_constructor=fixture_provider, ), ) diff --git a/e2e/test_admin_models.py b/e2e/test_admin_models.py index ca67c8a0..6a59b3a6 100644 --- a/e2e/test_admin_models.py +++ b/e2e/test_admin_models.py @@ -1,6 +1,7 @@ """Rendered model-fallback controls for the local Admin UI.""" import json +from typing import cast import pytest from playwright.sync_api import Locator, Page, expect @@ -25,7 +26,7 @@ def _refresh_openrouter_models(page: Page) -> None: page.get_by_role("button", name="Model Config", exact=True).click() -def _fallbacks(editor: Locator) -> list[dict[str, str]]: +def _fallbacks(editor: Locator) -> list[dict[str, str | int]]: value = editor.locator("#field-MODEL_FALLBACKS").input_value() return json.loads(value) if value else [] @@ -93,7 +94,7 @@ def test_model_routing_saves_independent_model_and_effort_choices( assert values["MODEL_SONNET"] == "open_router/vendor/model-b" assert values["REASONING_SONNET"] == "low" assert json.loads(values["MODEL_FALLBACKS"]) == [ - {"model": "open_router/vendor/model-a", "reasoning": "off"} + {"model": "open_router/vendor/model-a", "reasoning": "off", "api_key_number": 1} ] expect(page.locator("#dirtyState")).to_have_text("No changes") page.reload() @@ -106,7 +107,7 @@ def test_model_routing_saves_independent_model_and_effort_choices( ) expect(sonnet.locator("#field-REASONING_SONNET")).to_have_value("low") assert _fallbacks(fallback) == [ - {"model": "open_router/vendor/model-a", "reasoning": "off"} + {"model": "open_router/vendor/model-a", "reasoning": "off", "api_key_number": 1} ] expect(page.locator("#dirtyState")).to_have_text("No changes") @@ -165,8 +166,12 @@ def test_fallback_editor_adds_filters_reorders_and_removes_models( add_input.fill("groq/custom-model") editor.get_by_role("button", name="Add fallback", exact=True).click() assert _fallbacks(editor) == [ - {"model": "open_router/vendor/model-b", "reasoning": "low"}, - {"model": "groq/custom-model", "reasoning": "inherit"}, + { + "model": "open_router/vendor/model-b", + "reasoning": "low", + "api_key_number": 1, + }, + {"model": "groq/custom-model", "reasoning": "inherit", "api_key_number": 1}, ] expect(page.locator("#dirtyState")).to_have_text("1 unsaved change") expect(page.locator("#applyButton")).to_be_enabled() @@ -183,12 +188,16 @@ def test_fallback_editor_adds_filters_reorders_and_removes_models( rows.nth(0), target_position={"x": 10, "y": 5} ) assert _fallbacks(editor) == [ - {"model": "groq/custom-model", "reasoning": "inherit"}, - {"model": "open_router/vendor/model-b", "reasoning": "low"}, + {"model": "groq/custom-model", "reasoning": "inherit", "api_key_number": 1}, + { + "model": "open_router/vendor/model-b", + "reasoning": "low", + "api_key_number": 1, + }, ] editor.get_by_role("button", name="Remove groq/custom-model", exact=True).click() assert _fallbacks(editor) == [ - {"model": "open_router/vendor/model-b", "reasoning": "low"} + {"model": "open_router/vendor/model-b", "reasoning": "low", "api_key_number": 1} ] editor.get_by_role( "button", name="Remove open_router/vendor/model-b", exact=True @@ -220,7 +229,7 @@ def test_fallback_editor_rejects_blank_and_duplicate_rows( add_input.fill("groq/custom-model") add_button.click() expect(page.locator("#messageArea")).to_have_text( - "That fallback model is already in the list." + "That model and API key are already in the fallback list." ) expect(editor.locator(".model-list-row")).to_have_count(1) @@ -317,7 +326,7 @@ def test_fallback_editor_pairs_match_existing_controls_and_save_effort( ) ).to_have_value("max") assert _fallbacks(editor) == [ - {"model": "open_router/vendor/model-b", "reasoning": "max"} + {"model": "open_router/vendor/model-b", "reasoning": "max", "api_key_number": 1} ] expect(page.locator("#dirtyState")).to_have_text("No changes") @@ -329,15 +338,15 @@ def test_fallback_drag_and_keyboard_reordering_save_model_effort_pairs( _open_models(page, admin_base_url) editor = page.locator('.field[data-key="MODEL_FALLBACKS"]') entries = [ - {"model": "groq/model-a", "reasoning": "low"}, - {"model": "groq/model-b", "reasoning": "client"}, - {"model": "groq/model-c", "reasoning": "max"}, - {"model": "groq/model-d", "reasoning": "high"}, + {"model": "groq/model-a", "reasoning": "low", "api_key_number": 1}, + {"model": "groq/model-b", "reasoning": "client", "api_key_number": 1}, + {"model": "groq/model-c", "reasoning": "max", "api_key_number": 1}, + {"model": "groq/model-d", "reasoning": "high", "api_key_number": 1}, ] for entry in entries: - editor.locator("#field-MODEL_FALLBACKS-add").fill(entry["model"]) + editor.locator("#field-MODEL_FALLBACKS-add").fill(cast(str, entry["model"])) editor.get_by_label("Fallback reasoning effort", exact=True).select_option( - entry["reasoning"] + cast(str, entry["reasoning"]) ) editor.get_by_role("button", name="Add fallback", exact=True).click() rows = editor.locator(".model-list-row") @@ -510,7 +519,7 @@ def test_editable_fallbacks_keep_focus_and_live_model_suggestions( model.fill("open_router/new-model") expect(model).to_be_focused() assert _fallbacks(editor) == [ - {"model": "open_router/new-model", "reasoning": "inherit"} + {"model": "open_router/new-model", "reasoning": "inherit", "api_key_number": 1} ] for effort in ("low", "high", "client"): editor.get_by_label( @@ -532,7 +541,7 @@ def test_editable_fallbacks_keep_focus_and_live_model_suggestions( expect(model).to_have_value("new-model") editor.get_by_role("option", name="open_router/new-model", exact=True).click() assert _fallbacks(editor) == [ - {"model": "open_router/new-model", "reasoning": "client"} + {"model": "open_router/new-model", "reasoning": "client", "api_key_number": 1} ] registered = page.evaluate("state.modelComboboxes.size") for _ in range(3): diff --git a/e2e/test_fallback_key_controls.py b/e2e/test_fallback_key_controls.py new file mode 100644 index 00000000..eb9a0795 --- /dev/null +++ b/e2e/test_fallback_key_controls.py @@ -0,0 +1,211 @@ +"""Fallback rows retain numbered credential choices through actual Admin edits.""" + +import json + +import pytest +from playwright.sync_api import Page, expect + +from e2e.provider_support import open_provider +from e2e.test_admin_models import _fallbacks, _open_models +from e2e.test_api_key_lists import add_key + +CUSTOM_ID = "custom_" + "a" * 32 + + +@pytest.mark.parametrize( + "admin_base_url", + [{"GROQ_API_KEY": '["first","second","third"]'}], + indirect=True, +) +def test_numbered_rows_edit_reorder_save_and_reopen(page: Page, admin_base_url: str): + page.set_viewport_size({"width": 1280, "height": 1200}) + _open_models(page, admin_base_url) + editor = page.locator('.field[data-key="MODEL_FALLBACKS"]') + model = editor.locator("#field-MODEL_FALLBACKS-add") + key = editor.get_by_label("Fallback API key", exact=True) + effort = editor.get_by_label("Fallback reasoning effort", exact=True) + for number, reasoning in ((2, "low"), (3, "high")): + model.fill("groq/model-a") + expect(key.locator("option")).to_have_text(["Key 1", "Key 2", "Key 3"]) + key.select_option(str(number)) + effort.select_option(reasoning) + editor.get_by_role("button", name="Add fallback", exact=True).click() + expect(key).to_have_value("1") + model.fill("groq/model-a") + key.select_option("2") + editor.get_by_role("button", name="Add fallback", exact=True).click() + expect(editor.locator(".model-list-row")).to_have_count(2) + expect(page.locator("#messageArea")).not_to_have_text("") + first = editor.get_by_role("combobox", name="Fallback model 1", exact=True) + first.fill("groq/model-b") + expect(first).to_be_focused() + expect(editor.get_by_label("API key for groq/model-b", exact=True)).to_have_value( + "2" + ) + handle = editor.get_by_role("button", name="Move groq/model-a", exact=True) + handle.focus() + handle.press("ArrowUp") + expected = [ + {"model": "groq/model-a", "reasoning": "high", "api_key_number": 3}, + {"model": "groq/model-b", "reasoning": "low", "api_key_number": 2}, + ] + assert _fallbacks(editor) == expected + editor.locator(".model-list-row").first.get_by_role( + "button", name="Move groq/model-a", exact=True + ).drag_to( + editor.locator(".model-list-row").last, target_position={"x": 10, "y": 38} + ) + expected.reverse() + assert _fallbacks(editor) == expected + with page.expect_response("**/admin/api/config/apply") as saved: + page.locator("#applyButton").click() + assert saved.value.json()["applied"] is True + expect(page.locator("#dirtyState")).to_have_text("No changes") + page.reload() + page.get_by_role("button", name="Model Config", exact=True).click() + assert _fallbacks(editor) == expected + expect(editor.get_by_label("API key for groq/model-a", exact=True)).to_have_value( + "3" + ) + expect(editor.get_by_label("API key for groq/model-b", exact=True)).to_have_value( + "2" + ) + page.screenshot(path="untracked/fallback-api-key-controls.png", full_page=True) + + +@pytest.mark.parametrize( + "admin_base_url", + [ + { + "OPENAI_API_KEY": '["oa-first","oa-second"]', + "ZAI_API_KEY": '["zai-first","zai-second","zai-third"]', + "OPENCODE_API_KEY": '["oc-first","oc-second"]', + "FCC_CUSTOM_PROVIDERS": json.dumps( + [ + { + "provider_id": CUSTOM_ID, + "display_name": "Gateway", + "base_url": "https://gateway.example/v1", + "api_keys": ["custom-first", "custom-second"], + } + ] + ), + } + ], + indirect=True, +) +def test_options_follow_shared_custom_and_default_authentication( + page: Page, admin_base_url: str +): + _open_models(page, admin_base_url) + editor = page.locator('.field[data-key="MODEL_FALLBACKS"]') + model = editor.locator("#field-MODEL_FALLBACKS-add") + key = editor.get_by_label("Fallback API key", exact=True) + for provider, count in ( + ("openai_api", 2), + ("zai", 3), + ("zai_api", 3), + ("opencode_zen", 2), + ("opencode_go", 2), + (CUSTOM_ID, 2), + ): + model.fill(provider + "/model") + expect(key.locator("option")).to_have_text( + [f"Key {index}" for index in range(1, count + 1)] + ) + expect(key).to_be_enabled() + for provider in ("openai", "github_copilot", "vertex", "lmstudio"): + model.fill(provider + "/model") + expect(key.locator("option")).to_have_text(["Default authentication"]) + expect(key).to_be_disabled() + expect(key).to_have_value("1") + + +@pytest.mark.parametrize( + "admin_base_url", + [{"GROQ_API_KEY": '["first","second"]'}], + indirect=True, +) +def test_changed_provider_resets_key_without_resetting_reasoning( + page: Page, admin_base_url: str +): + _open_models(page, admin_base_url) + editor = page.locator('.field[data-key="MODEL_FALLBACKS"]') + add = editor.locator("#field-MODEL_FALLBACKS-add") + add.fill("groq/model") + editor.get_by_label("Fallback API key", exact=True).select_option("2") + editor.get_by_label("Fallback reasoning effort", exact=True).select_option("max") + editor.get_by_role("button", name="Add fallback", exact=True).click() + model = editor.get_by_role("combobox", name="Fallback model 1", exact=True) + model.fill("open_router/other") + expect(model).to_be_focused() + assert _fallbacks(editor) == [ + {"model": "open_router/other", "reasoning": "max", "api_key_number": 1} + ] + expect( + editor.get_by_label("API key for open_router/other", exact=True) + ).to_have_value("1") + add.fill("groq/model") + editor.get_by_label("Fallback API key", exact=True).select_option("2") + add.fill("lmstudio/local") + expect(editor.get_by_label("Fallback API key", exact=True)).to_have_value("1") + + +@pytest.mark.parametrize( + "admin_base_url", + [ + { + "MODEL_FALLBACKS": json.dumps( + [{"model": "groq/model", "api_key_number": 2}] + ), + "GROQ_API_KEY": '["first","second"]', + } + ], + indirect=True, +) +def test_process_owned_key_selector_is_locked(page: Page, admin_base_url: str): + _open_models(page, admin_base_url) + editor = page.locator('.field[data-key="MODEL_FALLBACKS"]') + expect(editor.get_by_label("Fallback API key", exact=True)).to_be_disabled() + key = editor.get_by_label("API key for groq/model", exact=True) + expect(key).to_be_disabled() + expect(key).to_have_value("2") + expect(page.locator("#dirtyState")).to_have_text("No changes") + + +def test_provider_reload_keeps_draft_and_marks_removed_number_unavailable( + page: Page, admin_base_url: str +): + _open_models(page, admin_base_url) + page.get_by_role("button", name="Providers", exact=True).click() + dialog = open_provider(page, "nvidia_nim") + for value in ("first", "second", "third"): + add_key(dialog, value) + dialog.get_by_role("button", name="Save", exact=True).click() + expect(dialog).not_to_be_visible() + expect(page.locator("#applyButton")).to_have_text("Apply") + page.get_by_role("button", name="Model Config", exact=True).click() + editor = page.locator('.field[data-key="MODEL_FALLBACKS"]') + editor.locator("#field-MODEL_FALLBACKS-add").fill("nvidia_nim/draft") + editor.get_by_label("Fallback API key", exact=True).select_option("3") + editor.get_by_role("button", name="Add fallback", exact=True).click() + page.locator("#field-MODEL_SONNET").fill("groq/unsaved") + page.get_by_role("button", name="Providers", exact=True).click() + dialog = open_provider(page, "nvidia_nim") + dialog.get_by_role("button", name="Remove Key 3", exact=True).click() + dialog.get_by_role("button", name="Save", exact=True).click() + expect(dialog).not_to_be_visible() + expect(page.locator("#applyButton")).to_have_text("Apply") + page.get_by_role("button", name="Model Config", exact=True).click() + expect(page.locator("#field-MODEL_SONNET")).to_have_value("groq/unsaved") + key = editor.get_by_label("API key for nvidia_nim/draft", exact=True) + expect(key).to_have_value("3") + expect(key.locator("option")).to_have_text( + ["Key 1", "Key 2", "Key 3 (unavailable)"] + ) + assert _fallbacks(editor)[0]["api_key_number"] == 3 + key.select_option("2") + with page.expect_response("**/admin/api/config/apply") as saved: + page.locator("#applyButton").click() + assert saved.value.json()["applied"] is True + expect(page.locator("#dirtyState")).to_have_text("No changes") diff --git a/src/free_claude_code/api/admin_static/admin.css b/src/free_claude_code/api/admin_static/admin.css index ce57aa8d..08e7a1d4 100644 --- a/src/free_claude_code/api/admin_static/admin.css +++ b/src/free_claude_code/api/admin_static/admin.css @@ -709,14 +709,15 @@ button.startup-busy::before, .model-list-add, .model-list-row { display: grid; - grid-template-columns: 20px repeat(2, minmax(0, 1fr)) 40px; + grid-template-columns: 20px repeat(2, minmax(0, 1fr)) 172px 40px; gap: 8px; align-items: center; } .model-list-add > .model-combobox { grid-column: 2; } -.model-list-add > select { grid-column: 3; } -.model-list-add > button { grid-column: 4; } +.model-list-add > select:first-of-type { grid-column: 3; } +.model-list-add > select:last-of-type { grid-column: 4; } +.model-list-add > button { grid-column: 5; } .model-list-editor .model-list-action { width: 40px; diff --git a/src/free_claude_code/api/admin_static/admin.js b/src/free_claude_code/api/admin_static/admin.js index 4bc5c582..1e08faf8 100644 --- a/src/free_claude_code/api/admin_static/admin.js +++ b/src/free_claude_code/api/admin_static/admin.js @@ -225,6 +225,7 @@ async function load({ providersOnly = false } = {}) { state.providerChecks.clear(); renderProviders(config.provider_status); if (!providersOnly) renderSections(config.sections, config.fields); + document.querySelectorAll(".model-list-editor").forEach((element) => element.modelListEditor.refreshKeyOptions()); byId("configPath").textContent = config.paths.managed; void refreshLocalStatus(config); void refreshStartup(); @@ -1081,10 +1082,12 @@ class ModelListEditor { this.values = input.value ? JSON.parse(input.value) : []; this.inputId = `field-${field.key}-add`; this.rowComboboxes = []; + this.rowKeys = []; this.draggedEntry = null; this.element = document.createElement("div"); this.element.className = "model-list-editor"; + this.element.modelListEditor = this; const addRow = document.createElement("div"); addRow.className = "model-list-add"; @@ -1092,11 +1095,20 @@ class ModelListEditor { this.addInput = addCombobox.input; this.addEffort = this.effortSelect("inherit", "Fallback reasoning effort"); this.addEffort.id = `field-${field.key}-effort`; + this.addKey = this.keySelect("", 1, "Fallback API key"); + this.addProviderId = ""; + const updateAddModel = () => { + const providerId = this.providerId(this.addInput.value); + this.updateKeySelect(this.addKey, this.addInput.value, providerId === this.addProviderId ? Number(this.addKey.value) : 1); + this.addProviderId = providerId; + }; + this.addInput.addEventListener("input", updateAddModel); + this.addInput.addEventListener("change", updateAddModel); this.addButton = this.actionButton("", "Add fallback", () => this.add()); this.addButton.className = "secondary-button model-list-action"; this.addButton.innerHTML = ''; - addRow.append(addCombobox.element, this.addEffort, this.addButton); + addRow.append(addCombobox.element, this.addEffort, this.addKey, this.addButton); this.rows = document.createElement("div"); this.rows.className = "model-list-rows"; @@ -1135,13 +1147,16 @@ class ModelListEditor { showMessage("Enter a full provider/model fallback.", "error"); return; } - if (this.values.some((entry) => entry.model === value)) { - showMessage("That fallback model is already in the list.", "error"); + const apiKeyNumber = Number(this.addKey.value); + if (this.values.some((entry) => entry.model === value && (entry.api_key_number ?? 1) === apiKeyNumber)) { + showMessage("That model and API key are already in the fallback list.", "error"); return; } - this.values.push({ model: value, reasoning: this.addEffort.value }); + this.values.push({ model: value, reasoning: this.addEffort.value, api_key_number: apiKeyNumber }); this.addInput.value = ""; this.addEffort.value = "inherit"; + this.addProviderId = ""; + this.updateKeySelect(this.addKey, "", 1); showMessage(""); this.sync(); this.renderRows(); @@ -1174,6 +1189,7 @@ class ModelListEditor { state.modelComboboxes.delete(combobox); }); this.rowComboboxes = []; + this.rowKeys = []; this.rows.replaceChildren(); this.values.forEach((entry, index) => { @@ -1187,6 +1203,12 @@ class ModelListEditor { entry.reasoning = effort.value; this.sync(); }); + const key = this.keySelect(entry.model, entry.api_key_number ?? 1, `API key for ${entry.model}`); + this.rowKeys.push({ entry, select: key }); + key.addEventListener("change", () => { + entry.api_key_number = Number(key.value); + this.sync(); + }); const handle = this.actionButton("⠿", `Move ${entry.model}`); handle.classList.add("model-list-drag"); @@ -1210,7 +1232,11 @@ class ModelListEditor { remove.className = "danger-button model-list-action"; remove.innerHTML = ''; const updateModel = () => { + const previousProvider = this.providerId(entry.model); entry.model = combobox.input.value.trim(); + if (this.providerId(entry.model) !== previousProvider) entry.api_key_number = 1; + this.updateKeySelect(key, entry.model, entry.api_key_number ?? 1); + key.setAttribute("aria-label", `API key for ${entry.model}`); effort.setAttribute("aria-label", `Reasoning effort for ${entry.model}`); handle.setAttribute("aria-label", `Move ${entry.model}`); remove.setAttribute("aria-label", `Remove ${entry.model}`); @@ -1219,7 +1245,7 @@ class ModelListEditor { combobox.input.addEventListener("input", updateModel); combobox.input.addEventListener("change", updateModel); - row.append(handle, combobox.element, effort, remove); + row.append(handle, combobox.element, effort, key, remove); this.rows.appendChild(row); }); } @@ -1268,6 +1294,48 @@ class ModelListEditor { return select; } + providerId(model) { + return model.includes("/") ? model.trim().split("/", 1)[0] : ""; + } + + keyCount(model) { + const providerId = this.providerId(model); + const custom = state.config.custom_providers.find((provider) => provider.provider_id === providerId); + if (custom) return custom.api_keys.entries.length; + const provider = state.config.provider_status.find((item) => item.provider_id === providerId); + const credential = provider?.settings_keys.map((key) => state.fields.get(key)).find((field) => field?.type === "api_key_list"); + return credential ? credential.value.entries.length : null; + } + + keySelect(model, value, label) { + const select = window.FccFormControls.configure(document.createElement("select")); + select.setAttribute("aria-label", label); + this.updateKeySelect(select, model, value); + return select; + } + + updateKeySelect(select, model, value) { + const count = this.keyCount(model); + select.replaceChildren(); + if (count === null) { + select.appendChild(option(1, this.providerId(model) ? "Default authentication" : "Key 1")); + } else { + for (let number = 1; number <= Math.max(1, count); number++) select.appendChild(option(number, `Key ${number}`)); + } + if (value > Math.max(1, count ?? 1)) { + const unavailable = option(value, `Key ${value} (unavailable)`); + unavailable.disabled = true; + select.appendChild(unavailable); + } + select.value = String(value); + select.disabled = this.field.locked || count === null || (count === 0 && value === 1); + } + + refreshKeyOptions() { + this.updateKeySelect(this.addKey, this.addInput.value, Number(this.addKey.value)); + this.rowKeys.forEach(({ entry, select }) => this.updateKeySelect(select, entry.model, entry.api_key_number ?? 1)); + } + actionButton(text, label, action = null) { const button = document.createElement("button"); button.type = "button"; diff --git a/src/free_claude_code/api/dependencies.py b/src/free_claude_code/api/dependencies.py index b848470a..041cc286 100644 --- a/src/free_claude_code/api/dependencies.py +++ b/src/free_claude_code/api/dependencies.py @@ -24,13 +24,14 @@ def get_settings(services: ApiServices = Depends(get_services)) -> Settings: async def resolve_provider( provider_type: str, + api_key_number: int = 1, *, lease: RequestRuntimeLease, ) -> ProviderPort: """Resolve a provider through one retained generation.""" - should_log_init = not lease.is_provider_cached(provider_type) + should_log_init = not lease.is_provider_cached(provider_type, api_key_number) try: - provider = await lease.resolve_provider(provider_type) + provider = await lease.resolve_provider(provider_type, api_key_number) except UnknownProviderError: logger.error( "Unknown provider_type: '{}'. Supported: {}", @@ -39,7 +40,7 @@ async def resolve_provider( ) raise if should_log_init: - logger.info("Provider initialized: {}", provider_type) + logger.info("Provider initialized: {} Key {}", provider_type, api_key_number) return provider diff --git a/src/free_claude_code/api/routes.py b/src/free_claude_code/api/routes.py index 03436ed4..599d7666 100644 --- a/src/free_claude_code/api/routes.py +++ b/src/free_claude_code/api/routes.py @@ -51,7 +51,9 @@ router = APIRouter() def _provider_resolver(lease: RequestRuntimeLease) -> ProviderResolver: - return lambda provider_type: resolve_provider(provider_type, lease=lease) + return lambda provider_type, api_key_number: resolve_provider( + provider_type, api_key_number, lease=lease + ) async def _create_messages_response( diff --git a/src/free_claude_code/application/execution.py b/src/free_claude_code/application/execution.py index 27c00983..34c7b48e 100644 --- a/src/free_claude_code/application/execution.py +++ b/src/free_claude_code/application/execution.py @@ -131,6 +131,8 @@ class ProviderExecutor: "wire_api": wire_api, "from_provider_model_ref": failed.provider_model_ref, "to_provider_model_ref": selected.provider_model_ref, + "from_api_key_number": failed.api_key_number, + "to_api_key_number": selected.api_key_number, "candidate_index": candidate_index, "candidate_count": candidate_count, "failure_kind": failure.kind.value, @@ -170,6 +172,7 @@ class ProviderExecutor: "request_id": request_id, "wire_api": wire_api, "selected_provider_model_ref": selected.provider_model_ref, + "api_key_number": selected.api_key_number, "candidate_index": candidate_index, "candidate_count": candidate_count, } @@ -191,7 +194,9 @@ class ProviderExecutor: continuation: ContinuationSeed | None, ) -> AsyncIterator[str]: target = candidate.target - provider = await self._provider_resolver(target.provider_id) + provider = await self._provider_resolver( + target.provider_id, target.api_key_number + ) return provider.stream_native_messages( routed.request.with_model(target.provider_model), request_id=request_id, @@ -252,7 +257,9 @@ class ProviderExecutor: continuation: ContinuationSeed | None, ) -> AsyncIterator[str]: target = candidate.target - provider = await self._provider_resolver(target.provider_id) + provider = await self._provider_resolver( + target.provider_id, target.api_key_number + ) request = ( primary_request if index == 0 @@ -311,7 +318,9 @@ class ProviderExecutor: continuation: ContinuationSeed | None, ) -> AsyncIterator[str]: target = candidate.target - provider = await self._provider_resolver(target.provider_id) + provider = await self._provider_resolver( + target.provider_id, target.api_key_number + ) request = ( primary_request if index == 0 diff --git a/src/free_claude_code/application/ports.py b/src/free_claude_code/application/ports.py index 1cba6e3f..3cc6712c 100644 --- a/src/free_claude_code/application/ports.py +++ b/src/free_claude_code/application/ports.py @@ -56,7 +56,7 @@ class ProviderPort(Protocol): ) -> AsyncIterator[str]: ... -ProviderResolver = Callable[[str], Awaitable[ProviderPort]] +ProviderResolver = Callable[[str, int], Awaitable[ProviderPort]] ModelInfoLookup = Callable[[str, str], ProviderModelInfo | None] @@ -69,11 +69,13 @@ class RequestRuntimeLease(Protocol): @property def settings(self) -> Settings: ... - def is_provider_cached(self, provider_id: str) -> bool: ... + def is_provider_cached(self, provider_id: str, api_key_number: int = 1) -> bool: ... async def wait_for_token_estimation(self) -> None: ... - async def resolve_provider(self, provider_id: str) -> ProviderPort: ... + async def resolve_provider( + self, provider_id: str, api_key_number: int = 1 + ) -> ProviderPort: ... def model_info( self, provider_id: str, model_id: str diff --git a/src/free_claude_code/application/routing.py b/src/free_claude_code/application/routing.py index 63991185..efac57cf 100644 --- a/src/free_claude_code/application/routing.py +++ b/src/free_claude_code/application/routing.py @@ -48,6 +48,7 @@ class ProviderModelTarget: provider_id: str provider_model: str provider_model_ref: str + api_key_number: int = 1 @dataclass(frozen=True, slots=True) @@ -148,25 +149,34 @@ class ModelRouter: reasoning_preference=reasoning_preference, ) - def _target_from_ref(self, provider_model_ref: str) -> ProviderModelTarget: + def _target_from_ref( + self, provider_model_ref: str, api_key_number: int = 1 + ) -> ProviderModelTarget: return self._target( parse_provider_type(provider_model_ref), parse_model_name(provider_model_ref), + api_key_number, ) - def _target(self, provider_id: str, provider_model: str) -> ProviderModelTarget: + def _target( + self, provider_id: str, provider_model: str, api_key_number: int = 1 + ) -> ProviderModelTarget: self._validate_provider_id(provider_id) return ProviderModelTarget( provider_id=provider_id, provider_model=provider_model, provider_model_ref=f"{provider_id}/{provider_model}", + api_key_number=api_key_number, ) def _fallback_targets( self, primary: ProviderModelTarget ) -> tuple[ProviderModelCandidate, ...]: configured = tuple( - ProviderModelCandidate(self._target_from_ref(entry.model), entry.reasoning) + ProviderModelCandidate( + self._target_from_ref(entry.model, entry.api_key_number), + entry.reasoning, + ) for entry in self._settings.model_fallbacks or () ) return tuple( diff --git a/src/free_claude_code/config/api_keys.py b/src/free_claude_code/config/api_keys.py index 5fa75a1e..510d2b5a 100644 --- a/src/free_claude_code/config/api_keys.py +++ b/src/free_claude_code/config/api_keys.py @@ -1,4 +1,4 @@ -"""Ordered secret credentials and the single credential used at runtime.""" +"""Ordered secret credentials and explicit credential selection.""" import json from typing import Annotated @@ -43,10 +43,15 @@ type ApiKeyList = Annotated[tuple[SecretStr, ...], BeforeValidator(normalize_api def first_api_key(keys: ApiKeyList) -> str | None: - """Select the only credential supplied to provider and voice resources.""" + """Select the first credential for default provider and voice resources.""" return keys[0].get_secret_value() if keys else None +def selected_api_key(keys: ApiKeyList, number: int) -> str | None: + """Select the one-based numbered credential, without substituting another.""" + return keys[number - 1].get_secret_value() if 1 <= number <= len(keys) else None + + def api_key_values(keys: ApiKeyList) -> list[str]: """Extract plaintext for explicit persistence and internal edit resolution.""" return [key.get_secret_value() for key in keys] diff --git a/src/free_claude_code/config/model_refs.py b/src/free_claude_code/config/model_refs.py index 313da77d..7b706764 100644 --- a/src/free_claude_code/config/model_refs.py +++ b/src/free_claude_code/config/model_refs.py @@ -9,6 +9,7 @@ from pydantic import ( BaseModel, BeforeValidator, ConfigDict, + Field, StringConstraints, TypeAdapter, field_validator, @@ -21,12 +22,13 @@ RETIRED_PROVIDER_IDS = frozenset({"github_models"}) class ModelFallback(BaseModel): - """One ordered fallback's model and reasoning preference.""" + """One ordered fallback's model, reasoning preference, and key number.""" model_config = ConfigDict(frozen=True, extra="forbid", hide_input_in_errors=True) model: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] reasoning: ReasoningPreference = ReasoningPreference.INHERIT + api_key_number: int = Field(default=1, strict=True, ge=1) @field_validator("model") @classmethod diff --git a/src/free_claude_code/config/settings.py b/src/free_claude_code/config/settings.py index 06b01dd7..9e6c2db5 100644 --- a/src/free_claude_code/config/settings.py +++ b/src/free_claude_code/config/settings.py @@ -19,6 +19,7 @@ from .model_refs import ModelFallback, ModelFallbacks from .nim import NimSettings from .provider_catalog import ( BEDROCK_DEFAULT_BASE, + PROVIDER_CATALOG, SUPPORTED_PROVIDER_IDS, ) from .reasoning import ReasoningPreference @@ -98,10 +99,28 @@ class Settings(BaseModel): refs = value if isinstance(value, tuple) else (value,) if value else () for ref in refs: model_ref = ref.model if isinstance(ref, ModelFallback) else ref - if model_ref.partition("/")[0] not in self.provider_ids: + provider_id = model_ref.partition("/")[0] + if provider_id not in self.provider_ids: raise ValueError( f"{field.upper()}: Invalid provider in model reference" ) + if isinstance(ref, ModelFallback) and ref.api_key_number > 1: + definition = self.custom_provider(provider_id) + descriptor = PROVIDER_CATALOG.get(provider_id) + if definition is not None: + keys = definition.api_keys + elif descriptor is not None and descriptor.credential_attr: + keys = getattr(self, descriptor.credential_attr) + else: + raise ValueError( + f"MODEL_FALLBACKS: {provider_id} has no API key list." + ) + if ref.api_key_number > len(keys): + raise ValueError( + f"MODEL_FALLBACKS: {ref.model} selects Key " + f"{ref.api_key_number}, but only {len(keys)} keys are " + "configured. Update the fallback selection first." + ) return self # ==================== OpenAI Platform API ==================== @@ -846,9 +865,11 @@ class Settings(BaseModel): ) -> tuple[ModelFallback, ...] | None: if value is None: return None - refs = tuple(entry.model for entry in value) + refs = tuple((entry.model, entry.api_key_number) for entry in value) if len(refs) != len(set(refs)): - raise ValueError("MODEL_FALLBACKS must not contain duplicate model refs.") + raise ValueError( + "MODEL_FALLBACKS must not contain duplicate model/key pairs." + ) return value @model_validator(mode="after") diff --git a/src/free_claude_code/providers/admission.py b/src/free_claude_code/providers/admission.py index ed5b81c3..00b7e2e4 100644 --- a/src/free_claude_code/providers/admission.py +++ b/src/free_claude_code/providers/admission.py @@ -5,12 +5,12 @@ import math import random import time import uuid -from collections import deque from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from datetime import UTC, datetime from email.utils import parsedate_to_datetime from enum import StrEnum +from functools import partial from typing import TypeVar from loguru import logger @@ -20,6 +20,7 @@ from free_claude_code.core.stream_delivery import ( current_stream_delivery, ) from free_claude_code.core.trace import trace_event +from free_claude_code.providers.admission_capacity import ProviderCapacityLimiter from free_claude_code.providers.admission_policy import ProviderAdmissionLimits from free_claude_code.providers.failure_policy import ( ProviderFailureOverride, @@ -412,7 +413,7 @@ class ProviderAttempt: class ProviderAdmissionController: - """Coordinate one provider's rate, concurrency, and recovery state. + """Coordinate one credential's recovery within a provider capacity budget. Normal attempts pass through a strict sliding window and concurrency bulkhead. The first shared transient failure opens one recovery episode. Exactly one @@ -431,6 +432,7 @@ class ProviderAdmissionController: base_delay: float = DEFAULT_UPSTREAM_BASE_DELAY, max_delay: float = DEFAULT_UPSTREAM_MAX_DELAY, jitter: float = DEFAULT_UPSTREAM_JITTER, + capacity: ProviderCapacityLimiter | None = None, ) -> None: limits = ProviderAdmissionLimits(rate_limit, rate_window, max_concurrency) if max_attempts <= 0: @@ -447,11 +449,7 @@ class ProviderAdmissionController: self._base_delay = base_delay self._max_delay = max_delay self._jitter = jitter - self._limits = limits - self._admission_times: deque[float] = deque() - self._last_discarded_at: float | None = None - self._active_attempts = 0 - self._capacity_changed = asyncio.Event() + self._capacity = capacity or ProviderCapacityLimiter(limits) self._condition = asyncio.Condition() self._episode: _RecoveryEpisode | None = None self._next_generation = 1 @@ -497,7 +495,9 @@ class ProviderAdmissionController: slot_acquired = False claim: _AttemptClaim | None = None try: - admitted = await self._acquire_capacity(execution, permit) + admitted = await self._capacity.acquire( + partial(self._permit_is_current, execution, permit) + ) if not admitted: await self._abandon_probe_permit(execution, permit) continue @@ -896,60 +896,14 @@ class ProviderAdmissionController: def reconfigure(self, limits: ProviderAdmissionLimits) -> None: """Update capacity without resetting occupancy, history, or recovery.""" - if limits != self._limits: - self._limits = limits - self._wake_capacity() - - def _wake_capacity(self) -> None: - changed = self._capacity_changed - self._capacity_changed = asyncio.Event() - changed.set() + self._capacity.reconfigure(limits) def _notify_recovery_waiters(self) -> None: self._condition.notify_all() - self._wake_capacity() - - async def _acquire_capacity( - self, execution: ProviderExecution, permit: _GatePermit - ) -> bool: - while self._permit_is_current(execution, permit): - changed = self._capacity_changed - now = time.monotonic() - limits = self._limits - cutoff = now - limits.rate_window - while self._admission_times and self._admission_times[0] <= cutoff: - self._last_discarded_at = self._admission_times.popleft() - - ready_at = now - if self._last_discarded_at is not None: - ready_at = max(ready_at, self._last_discarded_at + limits.rate_window) - if len(self._admission_times) >= limits.rate_limit: - ready_at = max( - ready_at, - self._admission_times[-limits.rate_limit] + limits.rate_window, - ) - if self._active_attempts < limits.max_concurrency and ready_at <= now: - self._admission_times.append(now) - self._active_attempts += 1 - return True - - # Occupancy changes explicitly wake us. Avoid an expired rate timer - # spinning while all concurrency slots remain occupied. - delay = ( - max(0.0, ready_at - now) - if self._active_attempts < limits.max_concurrency - else None - ) - try: - async with asyncio.timeout(delay): - await changed.wait() - except TimeoutError: - pass - return False + self._capacity.wake_waiters() def _release_concurrency(self) -> None: - self._active_attempts -= 1 - self._wake_capacity() + self._capacity.release() def _retry_delay(self, error: Exception, attempt: int) -> float: exponent = max(0, attempt - 1) diff --git a/src/free_claude_code/providers/admission_capacity.py b/src/free_claude_code/providers/admission_capacity.py new file mode 100644 index 00000000..f1f38885 --- /dev/null +++ b/src/free_claude_code/providers/admission_capacity.py @@ -0,0 +1,71 @@ +"""Provider capacity shared by clients with independent credential recovery.""" + +import asyncio +import time +from collections import deque +from collections.abc import Callable + +from free_claude_code.providers.admission_policy import ProviderAdmissionLimits + + +class ProviderCapacityLimiter: + """Own one provider's rate history, active slots, and capacity waiters.""" + + def __init__(self, limits: ProviderAdmissionLimits) -> None: + self._limits = limits + self._admission_times: deque[float] = deque() + self._last_discarded_at: float | None = None + self._active_attempts = 0 + self._changed = asyncio.Event() + + def reconfigure(self, limits: ProviderAdmissionLimits) -> None: + """Update limits without resetting occupancy or admission history.""" + if limits != self._limits: + self._limits = limits + self.wake_waiters() + + def wake_waiters(self) -> None: + changed = self._changed + self._changed = asyncio.Event() + changed.set() + + async def acquire(self, permit_is_current: Callable[[], bool]) -> bool: + """Charge capacity only while the caller's recovery permit is valid.""" + while permit_is_current(): + changed = self._changed + now = time.monotonic() + limits = self._limits + cutoff = now - limits.rate_window + while self._admission_times and self._admission_times[0] <= cutoff: + self._last_discarded_at = self._admission_times.popleft() + + ready_at = now + if self._last_discarded_at is not None: + ready_at = max(ready_at, self._last_discarded_at + limits.rate_window) + if len(self._admission_times) >= limits.rate_limit: + ready_at = max( + ready_at, + self._admission_times[-limits.rate_limit] + limits.rate_window, + ) + if self._active_attempts < limits.max_concurrency and ready_at <= now: + self._admission_times.append(now) + self._active_attempts += 1 + return True + + # Occupancy changes wake waiters. Avoid an expired rate timer + # spinning while all concurrency slots remain occupied. + delay = ( + max(0.0, ready_at - now) + if self._active_attempts < limits.max_concurrency + else None + ) + try: + async with asyncio.timeout(delay): + await changed.wait() + except TimeoutError: + pass + return False + + def release(self) -> None: + self._active_attempts -= 1 + self.wake_waiters() diff --git a/src/free_claude_code/providers/admission_registry.py b/src/free_claude_code/providers/admission_registry.py index 772c11bf..3f3e9668 100644 --- a/src/free_claude_code/providers/admission_registry.py +++ b/src/free_claude_code/providers/admission_registry.py @@ -1,8 +1,10 @@ """Loop-owned admission policy and controller lifetimes, independent of clients.""" from collections.abc import Iterable +from hashlib import sha256 from typing import TYPE_CHECKING +from free_claude_code.providers.admission_capacity import ProviderCapacityLimiter from free_claude_code.providers.admission_policy import ProviderAdmissionLimits if TYPE_CHECKING: @@ -10,41 +12,54 @@ if TYPE_CHECKING: class ProviderAdmissionRegistry: - """One protection budget per configured provider for one runtime manager.""" + """Share provider capacity and retain recovery by the actual credential.""" def __init__(self, limits: ProviderAdmissionLimits) -> None: self._limits = limits - self._controllers: dict[str, ProviderAdmissionController] = {} + self._controllers: dict[ + str, dict[bytes | None, ProviderAdmissionController] + ] = {} + self._capacities: dict[str, ProviderCapacityLimiter] = {} self._closed = False - def get(self, provider_id: str) -> ProviderAdmissionController: + def get( + self, provider_id: str, *, api_key: str | None = None + ) -> ProviderAdmissionController: if self._closed: raise RuntimeError("Provider admission registry is closed") if provider_id not in self._controllers: + self._controllers[provider_id] = {} + self._capacities[provider_id] = ProviderCapacityLimiter(self._limits) + credential = sha256(api_key.encode()).digest() if api_key else None + controllers = self._controllers[provider_id] + if credential not in controllers: # Provider preparation loads SDKs in a worker before reaching here. from free_claude_code.providers.admission import ProviderAdmissionController - self._controllers[provider_id] = ProviderAdmissionController( + controllers[credential] = ProviderAdmissionController( provider_name=provider_id, rate_limit=self._limits.rate_limit, rate_window=self._limits.rate_window, max_concurrency=self._limits.max_concurrency, + capacity=self._capacities[provider_id], ) - return self._controllers[provider_id] + return controllers[credential] def reconfigure(self, limits: ProviderAdmissionLimits) -> None: """Publish validated limits without yielding to waiting requests.""" self._limits = limits - for controller in self._controllers.values(): - controller.reconfigure(limits) + for capacity in self._capacities.values(): + capacity.reconfigure(limits) def retain_custom(self, provider_ids: Iterable[str]) -> None: retained = set(provider_ids) for provider_id in tuple(self._controllers): if provider_id.startswith("custom_") and provider_id not in retained: del self._controllers[provider_id] + del self._capacities[provider_id] def close(self) -> None: """Forget state only after client generations have drained and closed.""" self._closed = True self._controllers.clear() + self._capacities.clear() diff --git a/src/free_claude_code/providers/runtime/config.py b/src/free_claude_code/providers/runtime/config.py index 26abe475..9b126a3e 100644 --- a/src/free_claude_code/providers/runtime/config.py +++ b/src/free_claude_code/providers/runtime/config.py @@ -1,7 +1,7 @@ """Provider configuration construction from neutral catalog metadata.""" from free_claude_code.application.errors import ApplicationUnavailableError -from free_claude_code.config.api_keys import first_api_key +from free_claude_code.config.api_keys import selected_api_key from free_claude_code.config.custom_providers import CustomProviderDefinition from free_claude_code.config.provider_catalog import ProviderDescriptor from free_claude_code.config.settings import Settings @@ -17,14 +17,24 @@ def string_setting(settings: Settings, attr_name: str | None) -> str | None: def provider_credential( - descriptor: ProviderDescriptor, settings: Settings + descriptor: ProviderDescriptor, settings: Settings, api_key_number: int = 1 ) -> str | None: """Return the configured credential for a provider descriptor.""" - if descriptor.static_credential is not None: - return descriptor.static_credential if descriptor.credential_attr: - return first_api_key(getattr(settings, descriptor.credential_attr)) - return None + credential = selected_api_key( + getattr(settings, descriptor.credential_attr), api_key_number + ) + if api_key_number != 1 and credential is None: + raise ApplicationUnavailableError( + f"{descriptor.display_name} Key {api_key_number} is not configured. " + "Update the fallback key selection in the Admin UI." + ) + return credential + if api_key_number != 1: + raise ApplicationUnavailableError( + f"{descriptor.display_name} has no selectable API key list." + ) + return descriptor.static_credential def has_provider_configuration( @@ -57,10 +67,10 @@ def require_provider_credential( def build_provider_config( - descriptor: ProviderDescriptor, settings: Settings + descriptor: ProviderDescriptor, settings: Settings, *, api_key_number: int = 1 ) -> ProviderConfig: """Build shared provider configuration for one provider descriptor.""" - credential = provider_credential(descriptor, settings) + credential = provider_credential(descriptor, settings, api_key_number) require_provider_credential(descriptor, credential) base_url = string_setting(settings, descriptor.base_url_attr) resolved_base_url = base_url or descriptor.default_base_url @@ -100,10 +110,16 @@ def build_provider_config( def build_custom_provider_config( - definition: CustomProviderDefinition, settings: Settings + definition: CustomProviderDefinition, settings: Settings, *, api_key_number: int = 1 ) -> ProviderConfig: + credential = selected_api_key(definition.api_keys, api_key_number) + if api_key_number != 1 and credential is None: + raise ApplicationUnavailableError( + f"{definition.display_name} Key {api_key_number} is not configured. " + "Update the fallback key selection in the Admin UI." + ) return ProviderConfig( - api_key=first_api_key(definition.api_keys), + api_key=credential, base_url=definition.base_url, proxy=None, http_read_timeout=settings.http_read_timeout, diff --git a/src/free_claude_code/providers/runtime/factory.py b/src/free_claude_code/providers/runtime/factory.py index 94f9cadf..085b4e5e 100644 --- a/src/free_claude_code/providers/runtime/factory.py +++ b/src/free_claude_code/providers/runtime/factory.py @@ -282,6 +282,8 @@ def prepare_provider( provider_id: str, provider_loaders: Mapping[str, Callable[[], ProviderFactory]], custom_definition: CustomProviderDefinition | None = None, + *, + api_key_number: int = 1, ) -> Callable[[Settings, ProviderAdmissionRegistry], BaseProvider]: """Load implementation modules in a worker; return a loop-owned constructor.""" @@ -296,8 +298,10 @@ def prepare_provider( def construct_custom( settings: Settings, admission_registry: ProviderAdmissionRegistry ) -> BaseProvider: - config = build_custom_provider_config(custom_definition, settings) - admission = admission_registry.get(provider_id) + config = build_custom_provider_config( + custom_definition, settings, api_key_number=api_key_number + ) + admission = admission_registry.get(provider_id, api_key=config.api_key) return CustomProvider( config, definition=custom_definition, admission=admission ) @@ -318,8 +322,10 @@ def prepare_provider( def construct( settings: Settings, admission_registry: ProviderAdmissionRegistry ) -> BaseProvider: - config = build_provider_config(descriptor, settings) - admission = admission_registry.get(provider_id) + config = build_provider_config( + descriptor, settings, api_key_number=api_key_number + ) + admission = admission_registry.get(provider_id, api_key=config.api_key) if factory is not None: provider = factory(config, settings, admission) if ( diff --git a/src/free_claude_code/providers/runtime/runtime.py b/src/free_claude_code/providers/runtime/runtime.py index a615af40..5cab02e3 100644 --- a/src/free_claude_code/providers/runtime/runtime.py +++ b/src/free_claude_code/providers/runtime/runtime.py @@ -15,13 +15,17 @@ if TYPE_CHECKING: from .factory import ProviderFactory type ProviderConstructor = Callable[ - [str, Settings, ProviderAdmissionRegistry], Awaitable[BaseProvider] + [str, Settings, ProviderAdmissionRegistry, int], Awaitable[BaseProvider] ] type ProviderLoader = Callable[[], ProviderFactory] +type ProviderClientKey = tuple[str, int] def _load_constructor( - provider_id: str, provider_loaders: Mapping[str, ProviderLoader], settings: Settings + provider_id: str, + provider_loaders: Mapping[str, ProviderLoader], + settings: Settings, + api_key_number: int, ) -> Callable[[Settings, ProviderAdmissionRegistry], BaseProvider]: from .factory import prepare_provider @@ -31,6 +35,7 @@ def _load_constructor( settings.custom_provider(provider_id) if provider_id.startswith("custom_") else None, + api_key_number=api_key_number, ) @@ -38,11 +43,18 @@ async def create_provider( provider_id: str, settings: Settings, admission_registry: ProviderAdmissionRegistry, + api_key_number: int = 1, *, provider_loaders: Mapping[str, ProviderLoader] | None = None, ) -> BaseProvider: constructor = await run_sync_owned( - partial(_load_constructor, provider_id, provider_loaders or {}, settings) + partial( + _load_constructor, + provider_id, + provider_loaders or {}, + settings, + api_key_number, + ) ) return constructor(settings, admission_registry) @@ -54,7 +66,7 @@ class ProviderRuntime: self, settings: Settings, admission_registry: ProviderAdmissionRegistry, - providers: MutableMapping[str, BaseProvider] | None = None, + providers: MutableMapping[ProviderClientKey, BaseProvider] | None = None, *, provider_constructor: ProviderConstructor = create_provider, ) -> None: @@ -62,39 +74,43 @@ class ProviderRuntime: self._admission_registry = admission_registry self._providers = providers if providers is not None else {} self._provider_constructor = provider_constructor - self._creations: dict[str, asyncio.Task[BaseProvider]] = {} + self._creations: dict[ProviderClientKey, asyncio.Task[BaseProvider]] = {} self._closing = False - def is_cached(self, provider_id: str) -> bool: - """Return whether a provider for this id is already cached.""" - return provider_id in self._providers + def is_cached(self, provider_id: str, api_key_number: int = 1) -> bool: + """Return whether this provider/key client is already cached.""" + return (provider_id, api_key_number) in self._providers - async def resolve_provider(self, provider_id: str) -> BaseProvider: + async def resolve_provider( + self, provider_id: str, api_key_number: int = 1 + ) -> BaseProvider: """Return an existing provider or create it lazily.""" if self._closing: raise ApplicationUnavailableError("Provider runtime is shutting down.") - if provider_id in self._providers: - return self._providers[provider_id] - task = self._creations.get(provider_id) + key = (provider_id, api_key_number) + if key in self._providers: + return self._providers[key] + task = self._creations.get(key) if task is None or task.done(): - task = asyncio.create_task(self._construct(provider_id)) - self._creations[provider_id] = task - task.add_done_callback(partial(self._creation_done, provider_id)) + task = asyncio.create_task(self._construct(key)) + self._creations[key] = task + task.add_done_callback(partial(self._creation_done, key)) return await asyncio.shield(task) def _creation_done( - self, provider_id: str, task: asyncio.Task[BaseProvider] + self, key: ProviderClientKey, task: asyncio.Task[BaseProvider] ) -> None: - if self._creations.get(provider_id) is task: - del self._creations[provider_id] + if self._creations.get(key) is task: + del self._creations[key] if not task.cancelled(): task.exception() # Observe failures even when every caller stopped waiting. - async def _construct(self, provider_id: str) -> BaseProvider: + async def _construct(self, key: ProviderClientKey) -> BaseProvider: + provider_id, api_key_number = key provider = await self._provider_constructor( - provider_id, self.settings, self._admission_registry + provider_id, self.settings, self._admission_registry, api_key_number ) - self._providers[provider_id] = provider + self._providers[key] = provider return provider async def cleanup(self) -> None: diff --git a/src/free_claude_code/runtime/provider_manager.py b/src/free_claude_code/runtime/provider_manager.py index e04bc9e8..bdfe8de8 100644 --- a/src/free_claude_code/runtime/provider_manager.py +++ b/src/free_claude_code/runtime/provider_manager.py @@ -103,24 +103,28 @@ class ProviderGenerationLease: def settings(self) -> Settings: return self._generation.settings - def is_provider_cached(self, provider_id: str) -> bool: - return self._generation.runtime.is_cached(provider_id) + def is_provider_cached(self, provider_id: str, api_key_number: int = 1) -> bool: + return self._generation.runtime.is_cached(provider_id, api_key_number) async def wait_for_token_estimation(self) -> None: self._manager._ensure_open() await self._wait.wait(self._manager._ensure_encoder()) - async def resolve_provider(self, provider_id: str) -> BaseProvider: + async def resolve_provider( + self, provider_id: str, api_key_number: int = 1 + ) -> BaseProvider: self._manager._ensure_open() generation = self._generation if provider_id not in generation.initialized: task = self._manager._catalog_task(generation, provider_id) await self._wait.wait(task) - if generation.runtime.is_cached(provider_id): - provider = await generation.runtime.resolve_provider(provider_id) + if generation.runtime.is_cached(provider_id, api_key_number): + provider = await generation.runtime.resolve_provider( + provider_id, api_key_number + ) else: waiting = asyncio.create_task( - generation.runtime.resolve_provider(provider_id) + generation.runtime.resolve_provider(provider_id, api_key_number) ) try: provider = await self._wait.wait(waiting) diff --git a/tests/api/model_fallback_support.py b/tests/api/model_fallback_support.py index 55521373..8e31d66c 100644 --- a/tests/api/model_fallback_support.py +++ b/tests/api/model_fallback_support.py @@ -289,6 +289,7 @@ def fallback_client( def resolve( provider_id: str, + _api_key_number: int, *, lease: object, ) -> ControlledFallbackProvider: diff --git a/tests/api/support.py b/tests/api/support.py index d729394f..a7634bbf 100644 --- a/tests/api/support.py +++ b/tests/api/support.py @@ -89,7 +89,10 @@ def create_test_app( runtime_factory=lambda snapshot, admission_registry: ProviderRuntime( snapshot, admission_registry, - dict(providers), + { + (provider_id, 1): provider + for provider_id, provider in providers.items() + }, ), connected_provider_ids=connected_provider_ids, ) diff --git a/tests/api/test_api_handlers.py b/tests/api/test_api_handlers.py index 6f6cffe8..bb871098 100644 --- a/tests/api/test_api_handlers.py +++ b/tests/api/test_api_handlers.py @@ -132,7 +132,7 @@ async def test_messages_handler_passes_routed_request_and_stream_metadata() -> N provider = FakeProvider() handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -175,7 +175,7 @@ async def test_messages_handler_startup_invalid_request_stays_http_error( provider = RejectStartupProvider() handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -251,7 +251,7 @@ async def test_messages_handler_aggregates_provider_stream_when_stream_false() - ) handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -294,7 +294,7 @@ async def test_messages_handler_returns_error_json_for_stream_false_sse_error() ) handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -365,7 +365,7 @@ async def test_messages_handler_discards_partial_stream_false_output_on_error() ) handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -423,7 +423,7 @@ async def test_messages_handler_stream_false_provider_exception_keeps_status() - provider = FailingProvider() handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -460,7 +460,7 @@ async def test_messages_handler_normalizes_safety_classifier_policy( provider = FakeProvider() handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -504,7 +504,7 @@ async def test_messages_handler_preserves_thinking_for_non_classifier() -> None: provider = FakeProvider() handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -546,7 +546,7 @@ async def test_messages_handler_tolerates_required_thinking_for_classifier() -> provider = FakeProvider() handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -589,7 +589,7 @@ async def test_messages_handler_prefers_no_thinking_without_classifier_stop_hint provider = FakeProvider() handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -627,7 +627,7 @@ async def test_messages_handler_preserves_unowned_classifier_stop_sequences() -> provider = FakeProvider() handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), web_tools=StubWebToolsClient(), ) original_stops = ["custom", "", "custom", "", "tail"] @@ -682,7 +682,8 @@ async def test_messages_handler_optimization_intercepts_before_provider_executio async def test_responses_handler_bypasses_message_only_optimizations() -> None: provider = FakeProvider() handler = ResponsesHandler( - Settings(), provider_resolver=AsyncMock(side_effect=lambda _: provider) + Settings(), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), ) with patch( @@ -706,7 +707,8 @@ async def test_responses_handler_bypasses_message_only_optimizations() -> None: async def test_responses_handler_does_not_apply_safety_classifier_policy() -> None: provider = FakeProvider() handler = ResponsesHandler( - Settings(), provider_resolver=AsyncMock(side_effect=lambda _: provider) + Settings(), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), ) with patch("free_claude_code.api.handlers.messages.trace_event") as trace_mock: diff --git a/tests/api/test_classifier_provider_compatibility.py b/tests/api/test_classifier_provider_compatibility.py index de115972..bab92705 100644 --- a/tests/api/test_classifier_provider_compatibility.py +++ b/tests/api/test_classifier_provider_compatibility.py @@ -157,7 +157,9 @@ async def test_openrouter_numeric_sse_rejection_uses_classifier_correction( try: response = await MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock( + side_effect=lambda _, _api_key_number: provider + ), web_tools=StubWebToolsClient(), ).create(classifier_request()) assert isinstance(response, JSONResponse) @@ -248,7 +250,7 @@ async def test_classifier_mandatory_reasoning_still_returns_verdict( handler = MessagesHandler( settings, provider_resolver=AsyncMock( - side_effect=lambda provider_id: ( + side_effect=lambda provider_id, _api_key_number: ( primary if provider_id == "nvidia_nim" else provider ) ), @@ -309,7 +311,9 @@ async def test_handler_uses_cached_capabilities_for_the_selected_provider( try: handler = MessagesHandler( Settings(), - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock( + side_effect=lambda _, _api_key_number: provider + ), model_info_lookup=lambda provider_id, model_id: next( ( info diff --git a/tests/api/test_dependencies.py b/tests/api/test_dependencies.py index 6dc7d55b..a02b6ed9 100644 --- a/tests/api/test_dependencies.py +++ b/tests/api/test_dependencies.py @@ -83,8 +83,8 @@ async def test_resolve_provider_uses_retained_lease_and_logs_first_initializatio result = await resolve_provider("nvidia_nim", lease=lease) assert result is provider - lease.resolve_provider.assert_called_once_with("nvidia_nim") - log_info.assert_called_once_with("Provider initialized: {}", "nvidia_nim") + lease.resolve_provider.assert_called_once_with("nvidia_nim", 1) + log_info.assert_called_once_with("Provider initialized: {} Key {}", "nvidia_nim", 1) @pytest.mark.asyncio diff --git a/tests/api/test_fallback_key_requests.py b/tests/api/test_fallback_key_requests.py new file mode 100644 index 00000000..35c1f0ba --- /dev/null +++ b/tests/api/test_fallback_key_requests.py @@ -0,0 +1,150 @@ +"""Numbered fallback credentials reach actual HTTP requests across transports.""" + +import json + +import httpx +import httpx2 +import pytest +from fastapi.testclient import TestClient + +from free_claude_code.config.settings import Settings +from free_claude_code.providers.admission import ProviderAdmissionController +from tests.api.support import create_test_app +from tests.api.test_delivered_stream_recovery import text_events +from tests.providers.test_anthropic_messages_transport import _sse + +KEYS = ["synthetic-first", "synthetic-second", "synthetic-third"] +CUSTOM_ID = "custom_" + "d" * 32 + + +@pytest.fixture(params=[403, 429], ids=["permission-rejection", "rate-limit"]) +def upstream(monkeypatch, request): + calls = [] + failure_status = request.param + monkeypatch.setattr(ProviderAdmissionController, "_retry_delay", lambda *_args: 0.1) + + def payload(request): + if request.method == "GET": + return 200, {"id": "same-model", "data": []} + token = request.headers.get("x-api-key") or request.headers.get( + "authorization", "" + ).removeprefix("Bearer ") + body = json.loads(request.content) + calls.append((token, body)) + if token != "synthetic-third": + return failure_status, { + "error": { + "type": "rate_limit_error" + if failure_status == 429 + else "permission_error", + "message": "No quota", + } + } + protocol = ( + "messages" + if request.url.path.endswith("/messages") + else "responses" + if request.url.path.endswith("/responses") + else "chat" + ) + events = text_events(protocol, "selected key worked") + raw = ( + "".join(f"data: {json.dumps(event)}\n\n" for event in events).encode() + if protocol == "chat" + else _sse(*events) + ) + return 200, raw + + def response(request, library): + status, value = payload(request) + if isinstance(value, dict): + return library.Response(status, json=value) + return library.Response( + status, headers={"content-type": "text/event-stream"}, content=value + ) + + init = httpx.AsyncClient.__init__ + init2 = httpx2.AsyncClient.__init__ + + def local_http(self, *args, **kwargs): + kwargs["transport"] = httpx.MockTransport( + lambda request: response(request, httpx) + ) + init(self, *args, **kwargs) + + def local_http2(self, *args, **kwargs): + kwargs["transport"] = httpx2.MockTransport( + lambda request: response(request, httpx2) + ) + init2(self, *args, **kwargs) + + monkeypatch.setattr(httpx.AsyncClient, "__init__", local_http) + monkeypatch.setattr(httpx2.AsyncClient, "__init__", local_http2) + return calls, failure_status + + +@pytest.mark.parametrize("wire", ["messages", "responses"]) +@pytest.mark.parametrize( + ("provider", "api_format"), + [ + ("groq", "openai_chat"), + ("openai_api", "openai_responses"), + ("anthropic", "anthropic_messages"), + (CUSTOM_ID, "openai_chat"), + (CUSTOM_ID, "openai_responses"), + (CUSTOM_ID, "anthropic_messages"), + ], +) +def test_fallback_sequence_uses_each_rows_selected_key( + upstream, provider, api_format, wire +): + calls, failure_status = upstream + model = f"{provider}/same-model" + custom = ( + [ + { + "provider_id": CUSTOM_ID, + "display_name": "Gateway", + "base_url": "https://gateway.example/v1", + "api_format": api_format, + "api_keys": KEYS, + } + ] + if provider == CUSTOM_ID + else [] + ) + settings = Settings( + MODEL=model, + GROQ_API_KEY=KEYS, + OPENAI_API_KEY=KEYS, + ANTHROPIC_API_KEY=KEYS, + PROVIDER_RATE_LIMIT=1000, + PROVIDER_RATE_WINDOW=1, + FCC_CUSTOM_PROVIDERS=custom, + MODEL_FALLBACKS=[ + {"model": model, "reasoning": "off", "api_key_number": 2}, + {"model": model, "reasoning": "off", "api_key_number": 3}, + ], + ) + body = ( + {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 64} + if wire == "messages" + else {"input": "hello"} + ) + app = create_test_app(settings) + with TestClient(app) as client: + result = client.post( + f"/v1/{wire}", json={"model": model, "stream": True, **body} + ) + assert result.status_code == 200, result.text + assert "selected key worked" in result.text + tokens = [token for token, _ in calls] + assert [ + token + for index, token in enumerate(tokens) + if not index or token != tokens[index - 1] + ] == KEYS + if failure_status == 429: + assert tokens == [KEYS[0]] * 5 + [KEYS[1]] * 5 + [KEYS[2]] + assert all(body["model"] == "same-model" for _, body in calls) + assert all(key not in result.text for key in KEYS) diff --git a/tests/api/test_model_fallback_wire_controls.py b/tests/api/test_model_fallback_wire_controls.py index 02ac1be7..3e0fa139 100644 --- a/tests/api/test_model_fallback_wire_controls.py +++ b/tests/api/test_model_fallback_wire_controls.py @@ -65,7 +65,7 @@ async def test_classifier_fallback_keeps_manual_thinking_compatible(effort): handler = MessagesHandler( settings, provider_resolver=AsyncMock( - side_effect=lambda provider_id: ( + side_effect=lambda provider_id, _api_key_number: ( primary if provider_id == "nvidia_nim" else fallback ) ), @@ -134,7 +134,7 @@ async def test_fallback_effort_replaces_budget_and_preserves_wire_options( ) executor = ProviderExecutor( AsyncMock( - side_effect=lambda provider_id: ( + side_effect=lambda provider_id, _api_key_number: ( primary if provider_id == "nvidia_nim" else fallback ) ), diff --git a/tests/api/test_request_lifetime.py b/tests/api/test_request_lifetime.py index 6c3bd17a..42a0aec2 100644 --- a/tests/api/test_request_lifetime.py +++ b/tests/api/test_request_lifetime.py @@ -460,11 +460,13 @@ class _Lease: def model_info(self, provider_id: str, model_id: str) -> ProviderModelInfo | None: return None - def is_provider_cached(self, provider_id: str) -> bool: + def is_provider_cached(self, provider_id: str, api_key_number: int = 1) -> bool: del provider_id return True - async def resolve_provider(self, provider_id: str) -> ProviderPort: + async def resolve_provider( + self, provider_id: str, api_key_number: int = 1 + ) -> ProviderPort: assert provider_id == "nvidia_nim" return self._provider diff --git a/tests/api/test_request_outcomes.py b/tests/api/test_request_outcomes.py index 7e7f181f..c3b5e8f7 100644 --- a/tests/api/test_request_outcomes.py +++ b/tests/api/test_request_outcomes.py @@ -136,7 +136,7 @@ def test_fallback_records_selected_provider_without_logging_failed_attempt(outco with ( patch( "free_claude_code.api.routes.resolve_provider", - side_effect=lambda name, **kwargs: providers[name], + side_effect=lambda name, _api_key_number, **kwargs: providers[name], ), TestClient(create_test_app(settings)) as client, ): diff --git a/tests/api/test_safe_logging.py b/tests/api/test_safe_logging.py index 8b70f2ea..5ddbe294 100644 --- a/tests/api/test_safe_logging.py +++ b/tests/api/test_safe_logging.py @@ -27,7 +27,9 @@ async def test_create_message_skips_full_payload_debug_log_by_default(caplog): mock_provider.stream_messages = fake_stream service = MessagesHandler( settings, - provider_resolver=AsyncMock(side_effect=lambda _: mock_provider), + provider_resolver=AsyncMock( + side_effect=lambda _, _api_key_number: mock_provider + ), web_tools=StubWebToolsClient(), ) @@ -57,7 +59,9 @@ async def test_create_message_logs_full_payload_when_opt_in(caplog): mock_provider.stream_messages = fake_stream service = MessagesHandler( settings, - provider_resolver=AsyncMock(side_effect=lambda _: mock_provider), + provider_resolver=AsyncMock( + side_effect=lambda _, _api_key_number: mock_provider + ), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -123,7 +127,9 @@ async def test_create_message_unexpected_error_default_logs_exclude_exception_te mock_provider.stream_messages = stream_boom service = MessagesHandler( settings, - provider_resolver=AsyncMock(side_effect=lambda _: mock_provider), + provider_resolver=AsyncMock( + side_effect=lambda _, _api_key_number: mock_provider + ), web_tools=StubWebToolsClient(), ) request = MessagesRequest( @@ -160,7 +166,9 @@ async def test_create_message_unexpected_error_terminal_json_ignores_status_code mock_provider.stream_messages = stream_boom service = MessagesHandler( settings, - provider_resolver=AsyncMock(side_effect=lambda _: mock_provider), + provider_resolver=AsyncMock( + side_effect=lambda _, _api_key_number: mock_provider + ), web_tools=StubWebToolsClient(), ) request = MessagesRequest( diff --git a/tests/api/test_web_server_tools.py b/tests/api/test_web_server_tools.py index 621dbc1a..598b87c2 100644 --- a/tests/api/test_web_server_tools.py +++ b/tests/api/test_web_server_tools.py @@ -139,7 +139,7 @@ def _local_tool_body( } ) - async def unexpected_provider(_provider_id: str): + async def unexpected_provider(_provider_id: str, _api_key_number: int): pytest.fail("Forced local tools must not resolve a provider") service = WebToolService( @@ -384,7 +384,7 @@ def _automatic_search_service( effective_settings = settings or Settings() return MessagesHandler( effective_settings, - provider_resolver=AsyncMock(side_effect=lambda _: provider), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: provider), model_router=FixedProviderModelRouter( effective_settings, _PROVIDER_IDS[0], @@ -608,7 +608,7 @@ async def test_service_rejects_forced_server_tool_when_local_handler_is_disabled assert settings.enable_web_server_tools is False service = MessagesHandler( settings, - provider_resolver=AsyncMock(side_effect=lambda _: MagicMock()), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: MagicMock()), model_router=FixedProviderModelRouter(settings, provider_id), web_tools=StubWebToolsClient(), ) @@ -1916,7 +1916,7 @@ async def test_service_rejects_listed_server_tools_for_every_provider( settings = Settings.model_validate({"ENABLE_WEB_SERVER_TOOLS": False}) service = MessagesHandler( settings, - provider_resolver=AsyncMock(side_effect=lambda _: MagicMock()), + provider_resolver=AsyncMock(side_effect=lambda _, _api_key_number: MagicMock()), model_router=FixedProviderModelRouter(settings, provider_id), web_tools=StubWebToolsClient(), ) diff --git a/tests/application/test_execution.py b/tests/application/test_execution.py index 5d5314e4..bd0c9813 100644 --- a/tests/application/test_execution.py +++ b/tests/application/test_execution.py @@ -391,7 +391,7 @@ async def test_fallback_receives_its_own_cached_capability_in_stream( fallback = MetadataProvider(["verdict"]) providers = {"provider": primary, "fallback": fallback} - async def resolve(name): + async def resolve(name, _api_key_number): resolved.add(name) return providers[name] @@ -450,7 +450,7 @@ def _executor_stream( request_id: str, ) -> AsyncIterator[str]: executor = ProviderExecutor( - AsyncMock(side_effect=lambda _provider_id: provider), + AsyncMock(side_effect=lambda _provider_id, _api_key_number: provider), token_counter=lambda _messages, _system, _tools: 17, progress_timeout_seconds=timeout_seconds, ) @@ -467,7 +467,7 @@ async def test_executor_routes_native_responses_without_messages_conversion() -> routed = _routed_responses_request() request = routed.request executor = ProviderExecutor( - AsyncMock(side_effect=lambda _provider_id: provider), + AsyncMock(side_effect=lambda _provider_id, _api_key_number: provider), progress_timeout_seconds=60.0, responses_token_counter=lambda _request: 23, ) @@ -502,7 +502,7 @@ async def test_executor_uses_structural_provider_port_and_defers_stream_startup( routed = _routed_request() request = routed.request executor = ProviderExecutor( - AsyncMock(side_effect=lambda _provider_id: provider), + AsyncMock(side_effect=lambda _provider_id, _api_key_number: provider), progress_timeout_seconds=60.0, token_counter=lambda _messages, _system, _tools: 17, ) @@ -546,7 +546,7 @@ async def test_primary_success_never_resolves_fallback() -> None: fallback = FakeProvider() resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return {"provider": primary, "fallback": fallback}[provider_id] @@ -570,7 +570,7 @@ async def test_retryable_preframe_failure_selects_fallback_after_closing_primary primary = ControlledProvider([_execution_failure("primary overloaded")]) fallback = ControlledProvider(["fallback-frame"]) - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: if provider_id == "fallback": assert primary.stream_close_calls == 1 order.append("fallback-resolved") @@ -622,6 +622,8 @@ async def test_retryable_preframe_failure_selects_fallback_after_closing_primary "wire_api": "messages", "from_provider_model_ref": "provider/provider-model", "to_provider_model_ref": "fallback/fallback-model", + "from_api_key_number": 1, + "to_api_key_number": 1, "candidate_index": 2, "candidate_count": 2, "failure_kind": "overloaded", @@ -649,7 +651,9 @@ async def test_fallback_chain_preserves_exact_last_failure() -> None: "second": ControlledProvider([second]), } executor = ProviderExecutor( - AsyncMock(side_effect=providers.__getitem__), + AsyncMock( + side_effect=lambda provider_id, _api_key_number: providers[provider_id] + ), progress_timeout_seconds=60.0, ) stream = executor.stream_messages( @@ -674,7 +678,7 @@ async def test_multiple_retryable_failures_walk_fallbacks_in_order() -> None: } resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return providers[provider_id] @@ -699,7 +703,7 @@ async def test_empty_primary_completion_does_not_select_fallback() -> None: fallback = FakeProvider() resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return {"provider": primary, "fallback": fallback}[provider_id] @@ -721,7 +725,7 @@ async def test_unexpected_primary_failure_does_not_select_fallback() -> None: fallback = FakeProvider() resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return {"provider": primary, "fallback": fallback}[provider_id] @@ -745,7 +749,7 @@ async def test_invalid_none_stream_remains_terminal() -> None: fallback = FakeProvider() resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return {"provider": primary, "fallback": fallback}[provider_id] @@ -778,7 +782,7 @@ async def test_nonretryable_provider_failure_selects_fallback_before_first_frame fallback = ControlledProvider(["fallback-frame"]) resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return {"provider": primary, "fallback": fallback}[provider_id] @@ -815,7 +819,9 @@ async def test_primary_canonical_startup_failure_selects_fallback() -> None: fallback = ControlledProvider(["fallback-frame"]) providers = {"provider": primary, "fallback": fallback} executor = ProviderExecutor( - AsyncMock(side_effect=providers.__getitem__), + AsyncMock( + side_effect=lambda provider_id, _api_key_number: providers[provider_id] + ), progress_timeout_seconds=60.0, ) @@ -844,7 +850,9 @@ async def test_nonretryable_stream_construction_failure_selects_fallback() -> No fallback = ControlledProvider(["fallback-frame"]) providers = {"provider": primary, "fallback": fallback} executor = ProviderExecutor( - AsyncMock(side_effect=providers.__getitem__), + AsyncMock( + side_effect=lambda provider_id, _api_key_number: providers[provider_id] + ), progress_timeout_seconds=60.0, ) stream = executor.stream_messages( @@ -864,7 +872,7 @@ async def test_failure_after_first_frame_never_selects_fallback() -> None: fallback = FakeProvider() resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return {"provider": primary, "fallback": fallback}[provider_id] @@ -890,7 +898,9 @@ async def test_lazy_fallback_startup_application_error_stops_unchanged() -> None fallback = ApplicationErrorValidationProvider(startup_error) providers = {"provider": primary, "fallback": fallback} executor = ProviderExecutor( - AsyncMock(side_effect=providers.__getitem__), + AsyncMock( + side_effect=lambda provider_id, _api_key_number: providers[provider_id] + ), progress_timeout_seconds=60.0, ) stream = executor.stream_messages( @@ -937,7 +947,9 @@ async def test_candidate_requests_are_isolated_from_provider_mutation() -> None: providers = {"provider": primary, "fallback": fallback} routed = _routed_request(_target("fallback", "model")) executor = ProviderExecutor( - AsyncMock(side_effect=providers.__getitem__), + AsyncMock( + side_effect=lambda provider_id, _api_key_number: providers[provider_id] + ), progress_timeout_seconds=60.0, ) @@ -960,7 +972,7 @@ async def test_closing_executor_stream_closes_provider_stream_once() -> None: provider = FakeProvider() routed = _routed_request() executor = ProviderExecutor( - AsyncMock(side_effect=lambda _provider_id: provider), + AsyncMock(side_effect=lambda _provider_id, _api_key_number: provider), progress_timeout_seconds=60.0, token_counter=lambda _messages, _system, _tools: 17, ) @@ -983,7 +995,7 @@ async def test_stream_construction_failure_remains_deferred_to_iteration() -> No fallback = FakeProvider() resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return {"provider": provider, "fallback": fallback}[provider_id] @@ -1010,7 +1022,7 @@ async def test_executor_validation_is_deferred_until_iteration() -> None: provider = FailingValidationProvider() token_counter = MagicMock(return_value=17) executor = ProviderExecutor( - AsyncMock(side_effect=lambda _provider_id: provider), + AsyncMock(side_effect=lambda _provider_id, _api_key_number: provider), progress_timeout_seconds=60.0, token_counter=token_counter, ) @@ -1033,7 +1045,7 @@ async def test_executor_validation_is_deferred_until_iteration() -> None: def test_executor_rejects_invalid_progress_timeout(timeout_seconds: float) -> None: with pytest.raises(ValueError, match="finite and positive"): ProviderExecutor( - AsyncMock(side_effect=lambda _provider_id: FakeProvider()), + AsyncMock(side_effect=lambda _provider_id, _api_key_number: FakeProvider()), progress_timeout_seconds=timeout_seconds, ) @@ -1086,7 +1098,7 @@ async def test_application_progress_timeout_never_resolves_fallback() -> None: fallback = FakeProvider() resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return {"provider": primary, "fallback": fallback}[provider_id] @@ -1117,7 +1129,7 @@ async def test_provider_cleanup_cannot_delay_fallback_past_progress_deadline() - fallback = FakeProvider() resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return {"provider": primary, "fallback": fallback}[provider_id] @@ -1220,7 +1232,9 @@ async def test_fallback_transition_does_not_reset_shared_progress_deadline() -> fallback = ControlledProvider([fallback_wait]) providers = {"provider": primary, "fallback": fallback} executor = ProviderExecutor( - AsyncMock(side_effect=providers.__getitem__), + AsyncMock( + side_effect=lambda provider_id, _api_key_number: providers[provider_id] + ), progress_timeout_seconds=60, ) stream = executor.stream_messages( @@ -1293,7 +1307,7 @@ async def test_cancelling_progress_wait_remains_cancellation() -> None: fallback = FakeProvider() resolved_ids: list[str] = [] - async def resolve(provider_id: str) -> FakeProvider: + async def resolve(provider_id: str, _api_key_number: int) -> FakeProvider: resolved_ids.append(provider_id) return {"provider": provider, "fallback": fallback}[provider_id] diff --git a/tests/application/test_request_conversion.py b/tests/application/test_request_conversion.py index 2446d56c..c8ea3976 100644 --- a/tests/application/test_request_conversion.py +++ b/tests/application/test_request_conversion.py @@ -57,7 +57,8 @@ async def test_executor_converts_chat_request_once(wire): client=client, ) executor = ProviderExecutor( - AsyncMock(side_effect=lambda _: provider), progress_timeout_seconds=60 + AsyncMock(side_effect=lambda _, _api_key_number: provider), + progress_timeout_seconds=60, ) builder_name = ( "_build_request_body" @@ -103,7 +104,8 @@ async def test_executor_converts_selected_opencode_route_once(wire, egress, warm await provider.list_model_infos() with patch.object(transport, name, wraps=getattr(transport, name)) as build: executor = ProviderExecutor( - AsyncMock(side_effect=lambda _: provider), progress_timeout_seconds=60 + AsyncMock(side_effect=lambda _, _api_key_number: provider), + progress_timeout_seconds=60, ) stream = getattr(executor, f"stream_{wire}")( routed, raw_log_payload=dict, request_id="opencode-conversion-once" @@ -142,7 +144,8 @@ async def test_executor_converts_subscription_request_once(wire): provider._responses, name, wraps=getattr(provider._responses, name) ) as build: executor = ProviderExecutor( - AsyncMock(side_effect=lambda _: provider), progress_timeout_seconds=60 + AsyncMock(side_effect=lambda _, _api_key_number: provider), + progress_timeout_seconds=60, ) stream = getattr(executor, f"stream_{wire}")( routed, raw_log_payload=dict, request_id="subscription-conversion-once" diff --git a/tests/application/test_routing.py b/tests/application/test_routing.py index 0822cc5c..50c7edcb 100644 --- a/tests/application/test_routing.py +++ b/tests/application/test_routing.py @@ -383,6 +383,33 @@ def test_fallback_effort_does_not_change_primary_identity(requested): ] == [("groq/backup", ReasoningPreference.LOW)] +@pytest.mark.parametrize("requested", ["claude-opus-4", "groq/primary"]) +def test_same_model_fallbacks_keep_alternate_keys_and_their_efforts(requested): + settings = Settings( + MODEL="groq/primary", + MODEL_OPUS="groq/primary", + GROQ_API_KEY=["synthetic-first", "synthetic-second", "synthetic-third"], + MODEL_FALLBACKS=[ + {"model": "groq/primary", "reasoning": "off", "api_key_number": 1}, + {"model": "groq/primary", "reasoning": "low", "api_key_number": 2}, + {"model": "groq/primary", "reasoning": "high", "api_key_number": 3}, + ], + ) + route = ModelRouter(settings).resolve(requested) + assert route.primary.api_key_number == 1 + assert [ + ( + candidate.target.provider_model_ref, + candidate.target.api_key_number, + candidate.reasoning_preference, + ) + for candidate in route.fallbacks + ] == [ + ("groq/primary", 2, ReasoningPreference.LOW), + ("groq/primary", 3, ReasoningPreference.HIGH), + ] + + def test_direct_gateway_route_preserves_original_and_deduplicates_canonical_target( settings, ): diff --git a/tests/application/test_web_tools.py b/tests/application/test_web_tools.py index db910c60..b50a2747 100644 --- a/tests/application/test_web_tools.py +++ b/tests/application/test_web_tools.py @@ -16,7 +16,7 @@ from free_claude_code.core.web_tools import WebFetchResult from tests.web_tools_support import StubWebToolsClient -async def _no_provider(_provider_id): +async def _no_provider(_provider_id, _api_key_number): pytest.fail("This request must not resolve a provider") diff --git a/tests/cli/test_doctor.py b/tests/cli/test_doctor.py index 26d8654a..28106732 100644 --- a/tests/cli/test_doctor.py +++ b/tests/cli/test_doctor.py @@ -275,6 +275,8 @@ def test_fallback_model_efforts_are_json_serializable_in_doctor_report(): ] settings = Settings(MODEL_FALLBACKS=entries) report = json.loads(json.dumps(diagnostics.settings_report(settings, {}))) - assert report["model_routing"]["model_fallbacks"] == entries + assert report["model_routing"]["model_fallbacks"] == [ + entry | {"api_key_number": 1} for entry in entries + ] assert report["providers"]["groq"]["referenced"] is True assert report["providers"]["open_router"]["referenced"] is True diff --git a/tests/config/test_fallback_key_settings.py b/tests/config/test_fallback_key_settings.py new file mode 100644 index 00000000..6984e72e --- /dev/null +++ b/tests/config/test_fallback_key_settings.py @@ -0,0 +1,229 @@ +"""Fallback key numbers refer to each provider's effective numbered list.""" + +import json +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from free_claude_code.config.admin.persistence import prepare_admin_update +from free_claude_code.config.admin.values import load_config_response +from free_claude_code.config.loader import ( + ManagedConfigSnapshot, + compose_settings_snapshot, +) +from free_claude_code.config.model_refs import configured_chat_model_refs +from free_claude_code.config.settings import Settings + +KEYS = ["synthetic-first", "synthetic-second", "synthetic-third"] +REVISION_KEY = b"test-fallback-revision" +CUSTOM_ID = "custom_" + "a" * 32 + + +def snapshot(managed, process=None): + process = process or {} + effective = compose_settings_snapshot(managed, process) + return ManagedConfigSnapshot( + effective.settings, effective.sources, managed, process, Path("unused.env") + ) + + +def admin_field(state, key): + fields = load_config_response(state, revision_key=REVISION_KEY)["fields"] + assert isinstance(fields, list) + for field in fields: + assert isinstance(field, dict) + if field["key"] == key: + return field + raise AssertionError(f"Missing Admin field: {key}") + + +@pytest.mark.parametrize( + "value", + ["groq/one", ["groq/one"], [{"model": "groq/one", "reasoning": "high"}]], +) +def test_existing_fallback_formats_use_key_one(value): + settings = Settings(MODEL_FALLBACKS=value) + assert settings.model_fallbacks is not None + assert settings.model_fallbacks[0].api_key_number == 1 + + +def test_explicit_key_number_and_effort_survive_json_input(): + settings = Settings.model_validate( + { + "GROQ_API_KEY": KEYS, + "MODEL_FALLBACKS": '[{"model":"groq/one","reasoning":"high","api_key_number":3}]', + } + ) + assert settings.model_fallbacks is not None + assert settings.model_fallbacks[0].model_dump(mode="json") == { + "model": "groq/one", + "reasoning": "high", + "api_key_number": 3, + } + + +def test_same_model_can_use_distinct_keys_without_duplicate_discovery_refs(): + settings = Settings( + MODEL="groq/one", + GROQ_API_KEY=KEYS, + MODEL_FALLBACKS=[ + {"model": "groq/one", "api_key_number": 1}, + {"model": "groq/one", "api_key_number": 2}, + {"model": "groq/one", "api_key_number": 3}, + ], + ) + assert settings.model_fallbacks is not None + assert [entry.api_key_number for entry in settings.model_fallbacks] == [1, 2, 3] + assert [ref.model_ref for ref in configured_chat_model_refs(settings)] == [ + "groq/one" + ] + + +def test_different_efforts_do_not_allow_duplicate_model_key_pairs(): + with pytest.raises(ValidationError, match="duplicate"): + Settings( + GROQ_API_KEY=KEYS, + MODEL_FALLBACKS=[ + {"model": "groq/one", "reasoning": "off", "api_key_number": 2}, + {"model": "groq/one", "reasoning": "high", "api_key_number": 2}, + ], + ) + + +@pytest.mark.parametrize("number", [0, -1, True, 1.5, "2"]) +def test_key_number_requires_a_positive_integer(number): + with pytest.raises(ValidationError) as rejected: + Settings( + GROQ_API_KEY=KEYS, + MODEL_FALLBACKS=[{"model": "groq/one", "api_key_number": number}], + ) + errors = rejected.value.errors(include_input=False) + assert errors[0]["loc"] == ("MODEL_FALLBACKS", 0, "api_key_number") + assert errors[0]["type"] in {"greater_than_equal", "int_type"} + + +def test_unavailable_key_is_rejected_without_disclosing_keys(): + with pytest.raises(ValidationError, match="Key 4") as rejected: + Settings( + GROQ_API_KEY=KEYS, + MODEL_FALLBACKS=[{"model": "groq/one", "api_key_number": 4}], + ) + assert all(key not in str(rejected.value) for key in KEYS) + + +@pytest.mark.parametrize("provider", ["openai", "github_copilot", "vertex", "ollama"]) +def test_other_authentication_has_only_the_default_selection(provider): + default = Settings( + MODEL_FALLBACKS=[{"model": f"{provider}/one", "api_key_number": 1}] + ) + assert default.model_fallbacks is not None + assert default.model_fallbacks[0].api_key_number == 1 + with pytest.raises(ValidationError, match="API key list"): + Settings(MODEL_FALLBACKS=[{"model": f"{provider}/one", "api_key_number": 2}]) + + +@pytest.mark.parametrize("provider", ["zai", "zai_api", "opencode_zen", "opencode_go"]) +def test_shared_credential_owners_supply_the_same_numbered_keys(provider): + settings = Settings( + ZAI_API_KEY=KEYS, + OPENCODE_API_KEY=KEYS, + MODEL_FALLBACKS=[{"model": f"{provider}/one", "api_key_number": 3}], + ) + assert settings.model_fallbacks is not None + assert settings.model_fallbacks[0].api_key_number == 3 + + +def test_custom_provider_key_selection_and_unauthenticated_default(): + definition = { + "provider_id": CUSTOM_ID, + "display_name": "Gateway", + "base_url": "https://gateway.example/v1", + "api_keys": KEYS, + } + selected = Settings( + FCC_CUSTOM_PROVIDERS=[definition], + MODEL_FALLBACKS=[{"model": f"{CUSTOM_ID}/one", "api_key_number": 2}], + ) + assert selected.model_fallbacks is not None + assert selected.model_fallbacks[0].api_key_number == 2 + anonymous = Settings( + FCC_CUSTOM_PROVIDERS=[definition | {"api_keys": []}], + MODEL_FALLBACKS=[{"model": f"{CUSTOM_ID}/one"}], + ) + assert anonymous.model_fallbacks is not None + assert anonymous.model_fallbacks[0].api_key_number == 1 + with pytest.raises(ValidationError, match="Key 2"): + Settings( + FCC_CUSTOM_PROVIDERS=[definition | {"api_keys": []}], + MODEL_FALLBACKS=[{"model": f"{CUSTOM_ID}/one", "api_key_number": 2}], + ) + + +def test_admin_save_reload_preserves_the_selected_number_and_clean_state(): + entries = [{"model": "groq/one", "reasoning": "low", "api_key_number": 3}] + original = snapshot({"GROQ_API_KEY": json.dumps(KEYS)}) + saved = prepare_admin_update( + {"MODEL_FALLBACKS": entries}, + original, + original.settings, + revision_key=REVISION_KEY, + ) + assert saved.valid, saved.errors + assert json.loads(saved.target_values["MODEL_FALLBACKS"]) == entries + loaded = snapshot(saved.target_values) + field = admin_field(loaded, "MODEL_FALLBACKS") + value = field["value"] + assert isinstance(value, str) + assert json.loads(value) == entries + unchanged = prepare_admin_update( + {"MODEL_FALLBACKS": field["value"]}, + loaded, + loaded.settings, + revision_key=REVISION_KEY, + ) + assert unchanged.valid and unchanged.changed_keys == () + assert unchanged.pending_fields == () + + +def test_process_fallbacks_select_from_process_owned_keys(): + state = snapshot( + { + "GROQ_API_KEY": json.dumps(KEYS), + "MODEL_FALLBACKS": '[{"model":"groq/one","api_key_number":3}]', + }, + { + "GROQ_API_KEY": '["process-first","process-second"]', + "MODEL_FALLBACKS": '[{"model":"groq/one","api_key_number":2}]', + }, + ) + assert state.settings.model_fallbacks is not None + assert state.settings.model_fallbacks[0].api_key_number == 2 + field = admin_field(state, "MODEL_FALLBACKS") + assert field["locked"] is True + + +def test_removing_a_selected_key_fails_the_prospective_update(): + original = snapshot( + { + "GROQ_API_KEY": json.dumps(KEYS), + "MODEL_FALLBACKS": '[{"model":"groq/one","api_key_number":3}]', + } + ) + field = admin_field(original, "GROQ_API_KEY") + value = field["value"] + assert isinstance(value, dict) + prepared = prepare_admin_update( + { + "GROQ_API_KEY": { + "revision": value["revision"], + "entries": [{"retain": 0}, {"retain": 1}], + } + }, + original, + original.settings, + revision_key=REVISION_KEY, + ) + assert not prepared.valid + assert "Key 3" in " ".join(prepared.errors) + assert len(original.settings.groq_api_keys) == 3 diff --git a/tests/config/test_model_fallbacks.py b/tests/config/test_model_fallbacks.py index 52dee4f3..9f4a6e9a 100644 --- a/tests/config/test_model_fallbacks.py +++ b/tests/config/test_model_fallbacks.py @@ -116,6 +116,14 @@ def test_admin_fallback_records_round_trip_without_restart_or_false_changes(as_j {"model": "groq/one", "reasoning": "low"}, {"model": "open_router/vendor/two", "reasoning": "inherit"}, ] + canonical = [ + {"model": "groq/one", "reasoning": "low", "api_key_number": 1}, + { + "model": "open_router/vendor/two", + "reasoning": "inherit", + "api_key_number": 1, + }, + ] original = _snapshot() saved = prepare_admin_update( {"MODEL_FALLBACKS": json.dumps(entries) if as_json else entries}, @@ -125,11 +133,11 @@ def test_admin_fallback_records_round_trip_without_restart_or_false_changes(as_j ) assert saved.valid, saved.errors assert saved.pending_fields == () - assert json.loads(saved.target_values["MODEL_FALLBACKS"]) == entries + assert json.loads(saved.target_values["MODEL_FALLBACKS"]) == canonical loaded = _snapshot(saved.target_values) field = _fallback_field(loaded) value = field["value"] - assert isinstance(value, str) and json.loads(value) == entries + assert isinstance(value, str) and json.loads(value) == canonical assert field["source"] == ConfigSource.MANAGED.value options = field["options"] assert isinstance(options, list) @@ -174,9 +182,9 @@ def test_empty_process_fallbacks_override_managed_records_and_remain_locked(empt ) assert result.valid and result.settings is not None assert result.settings.model_fallbacks is None - assert json.loads(result.target_values["MODEL_FALLBACKS"]) == json.loads( - managed["MODEL_FALLBACKS"] - ) + assert json.loads(result.target_values["MODEL_FALLBACKS"]) == [ + {"model": "groq/one", "reasoning": "max", "api_key_number": 1} + ] def test_retired_provider_normalization_preserves_surviving_efforts(): diff --git a/tests/providers/test_admission_reconfiguration.py b/tests/providers/test_admission_reconfiguration.py index 1d0b025e..9d40efbb 100644 --- a/tests/providers/test_admission_reconfiguration.py +++ b/tests/providers/test_admission_reconfiguration.py @@ -7,7 +7,7 @@ import pytest from pydantic import ValidationError from free_claude_code.config.settings import Settings -from free_claude_code.providers import admission as admission_module +from free_claude_code.providers import admission_capacity as capacity_module from free_claude_code.providers.admission import ( ProviderAdmissionController, ProviderOperationKind, @@ -51,7 +51,7 @@ async def test_limit_changes_preserve_occupied_slots_and_wake_waiters(): async def test_waiting_for_concurrency_does_not_spend_rate_capacity(monkeypatch): clock = SimpleNamespace(now=0.0) monkeypatch.setattr( - admission_module, "time", SimpleNamespace(monotonic=lambda: clock.now) + capacity_module, "time", SimpleNamespace(monotonic=lambda: clock.now) ) controller = ProviderAdmissionController( provider_name="test", rate_limit=1, rate_window=10, max_concurrency=1 @@ -88,7 +88,7 @@ async def test_rate_count_increase_wakes_waiter_without_clearing_history(): async def test_longer_window_waits_only_until_discarded_history_expires(monkeypatch): clock = SimpleNamespace(now=0.0) monkeypatch.setattr( - admission_module, "time", SimpleNamespace(monotonic=lambda: clock.now) + capacity_module, "time", SimpleNamespace(monotonic=lambda: clock.now) ) controller = ProviderAdmissionController( provider_name="test", rate_limit=1, rate_window=2 @@ -115,7 +115,7 @@ async def test_shorter_window_wakes_waiter_and_longer_window_uses_retained_histo ): clock = SimpleNamespace(now=0.0) monkeypatch.setattr( - admission_module, "time", SimpleNamespace(monotonic=lambda: clock.now) + capacity_module, "time", SimpleNamespace(monotonic=lambda: clock.now) ) controller = ProviderAdmissionController( provider_name="test", rate_limit=1, rate_window=60 diff --git a/tests/providers/test_credential_admission.py b/tests/providers/test_credential_admission.py new file mode 100644 index 00000000..a6611b96 --- /dev/null +++ b/tests/providers/test_credential_admission.py @@ -0,0 +1,132 @@ +"""Selected credentials isolate recovery while sharing provider capacity.""" + +import asyncio + +import httpx +import pytest + +from free_claude_code.config.settings import Settings +from free_claude_code.providers.admission import ProviderOperationKind +from free_claude_code.providers.admission_policy import ProviderAdmissionLimits +from free_claude_code.providers.admission_registry import ProviderAdmissionRegistry +from free_claude_code.providers.groq import GroqProvider +from free_claude_code.providers.runtime import ProviderRuntime + + +@pytest.mark.asyncio +@pytest.mark.parametrize("protection", ["rate", "concurrency"]) +async def test_different_credentials_share_capacity_and_limit_updates(protection): + settings = Settings( + GROQ_API_KEY=["synthetic-first", "synthetic-second"], + PROVIDER_RATE_LIMIT=1 if protection == "rate" else 1000, + PROVIDER_RATE_WINDOW=60, + PROVIDER_MAX_CONCURRENCY=1, + ) + registry = ProviderAdmissionRegistry( + ProviderAdmissionLimits.from_settings(settings) + ) + runtime = ProviderRuntime(settings, registry) + pending = None + first_attempt = second_attempt = None + try: + first = await runtime.resolve_provider("groq", 1) + second = await runtime.resolve_provider("groq", 2) + assert isinstance(first, GroqProvider) + assert isinstance(second, GroqProvider) + first_attempt = await first._admission.start_execution().open_attempt( + ProviderOperationKind.GENERATION + ) + await first_attempt.accept() + if protection == "rate": + await first_attempt.aclose() + pending = asyncio.create_task( + second._admission.start_execution().open_attempt( + ProviderOperationKind.GENERATION + ) + ) + await asyncio.sleep(0) + assert not pending.done(), f"another credential bypassed provider {protection}" + registry.reconfigure(ProviderAdmissionLimits(1000, 60, 2)) + second_attempt = await asyncio.wait_for(pending, 1) + await second_attempt.accept() + finally: + if pending is not None: + pending.cancel() + outcomes = await asyncio.gather(pending, return_exceptions=True) + if not isinstance(outcomes[0], BaseException): + await outcomes[0].aclose() + if second_attempt is not None: + await second_attempt.aclose() + if first_attempt is not None: + await first_attempt.aclose() + await runtime.cleanup() + registry.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("new_first", ["synthetic-second", "replacement-first"]) +async def test_recovery_follows_credentials_across_key_reorder_and_replacement( + new_first, +): + old_settings = Settings( + GROQ_API_KEY=["synthetic-first", "synthetic-second"], + PROVIDER_RATE_LIMIT=1000, + ) + new_settings = Settings( + GROQ_API_KEY=[new_first, "synthetic-first"], PROVIDER_RATE_LIMIT=1000 + ) + registry = ProviderAdmissionRegistry( + ProviderAdmissionLimits.from_settings(old_settings) + ) + old_runtime = ProviderRuntime(old_settings, registry) + new_runtime = ProviderRuntime(new_settings, registry) + pending = None + old_attempt = healthy_attempt = None + try: + old_first = await old_runtime.resolve_provider("groq", 1) + healthy = await new_runtime.resolve_provider("groq", 1) + same_credential = await new_runtime.resolve_provider("groq", 2) + assert isinstance(old_first, GroqProvider) + assert isinstance(healthy, GroqProvider) + assert isinstance(same_credential, GroqProvider) + old_attempt = await old_first._admission.start_execution().open_attempt( + ProviderOperationKind.GENERATION + ) + response = httpx.Response( + 503, + headers={"retry-after": "60"}, + request=httpx.Request("POST", "https://provider.example"), + ) + await old_attempt.fail( + httpx.HTTPStatusError( + "unavailable", request=response.request, response=response + ) + ) + await old_attempt.aclose() + pending = asyncio.create_task( + same_credential._admission.start_execution().open_attempt( + ProviderOperationKind.GENERATION + ) + ) + await asyncio.sleep(0) + assert not pending.done(), "moving a credential bypassed its recovery" + healthy_attempt = await asyncio.wait_for( + healthy._admission.start_execution().open_attempt( + ProviderOperationKind.GENERATION + ), + 1, + ) + await healthy_attempt.accept() + finally: + if pending is not None: + pending.cancel() + outcomes = await asyncio.gather(pending, return_exceptions=True) + if not isinstance(outcomes[0], BaseException): + await outcomes[0].aclose() + if healthy_attempt is not None: + await healthy_attempt.aclose() + if old_attempt is not None: + await old_attempt.aclose() + await old_runtime.cleanup() + await new_runtime.cleanup() + registry.close() diff --git a/tests/providers/test_history_transports.py b/tests/providers/test_history_transports.py index c74820ee..644dc6c7 100644 --- a/tests/providers/test_history_transports.py +++ b/tests/providers/test_history_transports.py @@ -500,7 +500,11 @@ async def test_successful_fallback_preserves_raw_state_and_gets_unmodified_input routed = _routed_request(_target("open_router", "fallback-model")) original = routed.request.model_dump() executor = ProviderExecutor( - AsyncMock(side_effect=lambda name: primary if name == "provider" else provider), + AsyncMock( + side_effect=lambda name, _api_key_number: ( + primary if name == "provider" else provider + ) + ), progress_timeout_seconds=10, ) try: diff --git a/tests/providers/test_lmstudio.py b/tests/providers/test_lmstudio.py index dc28b1f2..880ee034 100644 --- a/tests/providers/test_lmstudio.py +++ b/tests/providers/test_lmstudio.py @@ -435,7 +435,7 @@ async def test_primary_context_validation_counts_toward_progress_timeout( fallback = FakeProvider() if wire == "messages" else ResponsesFakeProvider() providers = {"provider": provider, "fallback": fallback} executor = ProviderExecutor( - AsyncMock(side_effect=lambda name: providers[name]), + AsyncMock(side_effect=lambda name, _api_key_number: providers[name]), progress_timeout_seconds=0.05, ) route = _routed_request if wire == "messages" else _routed_responses_request @@ -910,7 +910,7 @@ async def test_context_rejection_before_first_event_allows_fallback(wire): fallback = FakeProvider() if wire == "messages" else ResponsesFakeProvider() providers = {"provider": provider, "fallback": fallback} executor = ProviderExecutor( - AsyncMock(side_effect=lambda name: providers[name]), + AsyncMock(side_effect=lambda name, _api_key_number: providers[name]), progress_timeout_seconds=2, ) route = _routed_request if wire == "messages" else _routed_responses_request diff --git a/tests/providers/test_model_discovery.py b/tests/providers/test_model_discovery.py index 43f49dbc..a5adbd53 100644 --- a/tests/providers/test_model_discovery.py +++ b/tests/providers/test_model_discovery.py @@ -77,7 +77,9 @@ def _manager( return ProviderRuntimeManager( settings, runtime_factory=lambda snapshot, admission_registry: ProviderRuntime( - snapshot, admission_registry, dict(providers) + snapshot, + admission_registry, + {(provider_id, 1): provider for provider_id, provider in providers.items()}, ), ) diff --git a/tests/providers/test_opencode.py b/tests/providers/test_opencode.py index 9deb2fb2..05138041 100644 --- a/tests/providers/test_opencode.py +++ b/tests/providers/test_opencode.py @@ -1119,7 +1119,7 @@ async def test_candidate_fallback_resolves_each_opencode_transport( reasoning_preference=ReasoningPreference.INHERIT, ) executor = ProviderExecutor( - AsyncMock(side_effect=lambda _provider_id: provider), + AsyncMock(side_effect=lambda _provider_id, _api_key_number: provider), progress_timeout_seconds=60.0, ) diff --git a/tests/providers/test_provider_api_key_selection.py b/tests/providers/test_provider_api_key_selection.py new file mode 100644 index 00000000..ff55195a --- /dev/null +++ b/tests/providers/test_provider_api_key_selection.py @@ -0,0 +1,118 @@ +"""A generation owns distinct, coalesced clients for numbered credentials.""" + +import asyncio + +import pytest + +from free_claude_code.application.errors import ApplicationUnavailableError +from free_claude_code.config.custom_providers import CustomProviderDefinition +from free_claude_code.config.provider_catalog import PROVIDER_CATALOG +from free_claude_code.config.settings import Settings +from free_claude_code.providers.admission_policy import ProviderAdmissionLimits +from free_claude_code.providers.admission_registry import ProviderAdmissionRegistry +from free_claude_code.providers.groq import GroqProvider +from free_claude_code.providers.runtime import ProviderRuntime +from free_claude_code.providers.runtime.config import ( + build_custom_provider_config, + build_provider_config, +) +from tests.api.support import ApiTestRuntime + +KEYS = ["synthetic-first", "synthetic-second", "synthetic-third"] + + +@pytest.mark.parametrize("provider", ["groq", "gemini", "anthropic", "openai_api"]) +def test_builtin_provider_configuration_selects_the_numbered_key(provider): + settings = Settings( + GROQ_API_KEY=KEYS, + GEMINI_API_KEY=KEYS, + ANTHROPIC_API_KEY=KEYS, + OPENAI_API_KEY=KEYS, + ) + config = build_provider_config( + PROVIDER_CATALOG[provider], settings, api_key_number=3 + ) + assert config.api_key == "synthetic-third" + + +def test_custom_provider_configuration_selects_the_numbered_key(): + definition = CustomProviderDefinition( + provider_id="custom_" + "b" * 32, + display_name="Gateway", + base_url="https://gateway.example/v1", + api_keys=KEYS, + ) + config = build_custom_provider_config(definition, Settings(), api_key_number=2) + assert config.api_key == "synthetic-second" + + +def test_missing_selected_key_never_falls_back_to_first(): + settings = Settings(GROQ_API_KEY=KEYS[:1]) + with pytest.raises(ApplicationUnavailableError, match="Key 2"): + build_provider_config(PROVIDER_CATALOG["groq"], settings, api_key_number=2) + + +@pytest.mark.asyncio +async def test_concurrent_resolutions_coalesce_by_provider_and_key(): + settings = Settings(GROQ_API_KEY=KEYS) + runtime = ProviderRuntime( + settings, + ProviderAdmissionRegistry(ProviderAdmissionLimits.from_settings(settings)), + ) + try: + first, second, second_again, third = await asyncio.gather( + runtime.resolve_provider("groq", 1), + runtime.resolve_provider("groq", 2), + runtime.resolve_provider("groq", 2), + runtime.resolve_provider("groq", 3), + ) + assert isinstance(first, GroqProvider) + assert isinstance(second, GroqProvider) + assert isinstance(third, GroqProvider) + assert second is second_again + assert await runtime.resolve_provider("groq") is first + assert first is not second and second is not third + assert [provider._config.api_key for provider in (first, second, third)] == KEYS + assert runtime.is_cached("groq") and runtime.is_cached("groq", 2) + assert first._admission is not second._admission + assert second._admission is not third._admission + finally: + await runtime.cleanup() + assert not runtime.is_cached("groq") + assert not runtime.is_cached("groq", 2) + assert not runtime.is_cached("groq", 3) + assert all(provider._client.is_closed() for provider in (first, second, third)) + + +@pytest.mark.asyncio +async def test_retained_generation_keeps_its_selected_key_after_a_key_edit(): + old_settings = Settings(MODEL="groq/one", GROQ_API_KEY=KEYS) + new_settings = Settings( + MODEL="groq/one", GROQ_API_KEY=["synthetic-first", "replacement-second"] + ) + manager = ApiTestRuntime(old_settings) + old_lease = await manager.acquire() + new_lease = None + try: + old_client = await old_lease.resolve_provider("groq", 2) + assert isinstance(old_client, GroqProvider) + + async def commit(): + return None + + await manager.replace(new_settings, commit=commit) + new_lease = await manager.acquire() + new_client = await new_lease.resolve_provider("groq", 2) + assert isinstance(new_client, GroqProvider) + assert old_client._config.api_key == "synthetic-second" + assert new_client._config.api_key == "replacement-second" + assert not old_client._client.is_closed() + assert await old_lease.resolve_provider("groq", 2) is old_client + await old_lease.release() + assert old_client._client.is_closed() + assert not new_client._client.is_closed() + finally: + await old_lease.release() + if new_lease is not None: + await new_lease.release() + await manager.close() diff --git a/tests/providers/test_provider_runtime.py b/tests/providers/test_provider_runtime.py index e7b9185c..243becd5 100644 --- a/tests/providers/test_provider_runtime.py +++ b/tests/providers/test_provider_runtime.py @@ -2,7 +2,7 @@ import asyncio import gc import subprocess import sys -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest @@ -1281,6 +1281,7 @@ async def test_create_provider_instantiates_each_builtin(): rate_limit=7, rate_window=11, max_concurrency=3, + capacity=ANY, ) admission_factory.reset_mock() @@ -1310,7 +1311,7 @@ async def test_provider_creation_retries_after_shared_failure(cancelled): provider = MagicMock(cleanup=AsyncMock()) attempts = 0 - async def construct(_id, _settings, _admission_registry): + async def construct(_id, _settings, _admission_registry, _api_key_number): nonlocal attempts attempts += 1 attempt = attempts @@ -1366,7 +1367,7 @@ async def test_finished_creation_callback_cannot_forget_a_new_retry(): attempts = 0 retry = None - async def construct(_id, _settings, _admission_registry): + async def construct(_id, _settings, _admission_registry, _api_key_number): nonlocal attempts, retry attempts += 1 if attempts == 1: @@ -1411,7 +1412,7 @@ async def test_cancelled_acquisition_keeps_construction_available_to_next_caller provider = MagicMock(cleanup=AsyncMock()) attempts = 0 - async def construct(_id, _settings, _admission_registry): + async def construct(_id, _settings, _admission_registry, _api_key_number): nonlocal attempts attempts += 1 entered.set() @@ -1453,7 +1454,7 @@ async def test_abandoned_creation_failure_has_no_unretrieved_exception(): previous_handler = loop.get_exception_handler() errors = [] - async def construct(_id, _settings, _admission_registry): + async def construct(_id, _settings, _admission_registry, _api_key_number): entered.set() await release.wait() failed.set() @@ -1495,7 +1496,7 @@ async def test_shutdown_drains_construction_that_finishes_during_cancellation(): entered = asyncio.Event() provider = MagicMock(cleanup=AsyncMock()) - async def construct(_id, _settings, _admission_registry): + async def construct(_id, _settings, _admission_registry, _api_key_number): entered.set() try: await asyncio.Event().wait() @@ -1611,7 +1612,7 @@ async def test_provider_runtime_cleanup_runs_all_even_if_one_fails() -> None: ProviderAdmissionRegistry( ProviderAdmissionLimits.from_settings(_make_settings()) ), - {"a": p1, "b": p2}, + {("a", 1): p1, ("b", 1): p2}, ) with pytest.raises(RuntimeError, match="first"): @@ -1652,7 +1653,7 @@ async def test_cancelled_cleanup_retains_current_and_unvisited_providers() -> No ProviderAdmissionRegistry( ProviderAdmissionLimits.from_settings(_make_settings()) ), - {"first": first, "second": second, "third": third}, + {("first", 1): first, ("second", 1): second, ("third", 1): third}, ) cleanup_task = asyncio.create_task(runtime.cleanup()) await second_started.wait() @@ -1688,7 +1689,7 @@ async def test_provider_runtime_cleanup_exceptiongroup_on_multiple_failures() -> ProviderAdmissionRegistry( ProviderAdmissionLimits.from_settings(_make_settings()) ), - {"x": p1, "y": p2}, + {("x", 1): p1, ("y", 1): p2}, ) with pytest.raises(ExceptionGroup) as exc_info: diff --git a/tests/runtime/test_anthropic_provider_lifecycle.py b/tests/runtime/test_anthropic_provider_lifecycle.py index 60b27732..812d7a29 100644 --- a/tests/runtime/test_anthropic_provider_lifecycle.py +++ b/tests/runtime/test_anthropic_provider_lifecycle.py @@ -38,7 +38,7 @@ async def test_workspace_replacement_retires_http_and_capabilities_with_lease(): transport=httpx.MockTransport(handle), ) owners.append(owner) - return ProviderRuntime(settings, registry, {"anthropic": owner}) + return ProviderRuntime(settings, registry, {("anthropic", 1): owner}) async def commit(): return None diff --git a/tests/runtime/test_application_runtime.py b/tests/runtime/test_application_runtime.py index 58b0984a..8df5d3dd 100644 --- a/tests/runtime/test_application_runtime.py +++ b/tests/runtime/test_application_runtime.py @@ -243,7 +243,7 @@ def _runtime_with_admin_provider( runtime_factory=lambda snapshot, admission_registry: ProviderRuntime( snapshot, admission_registry, - {"nvidia_nim": provider}, + {("nvidia_nim", 1): provider}, ), ) return ApplicationRuntime( diff --git a/tests/runtime/test_integration_startup.py b/tests/runtime/test_integration_startup.py index 7f2b5214..b75ebe47 100644 --- a/tests/runtime/test_integration_startup.py +++ b/tests/runtime/test_integration_startup.py @@ -52,7 +52,7 @@ def runtime(): return_value=frozenset({ProviderModelInfo("one")}) ) - async def construct(_id, _settings, _admission_registry): + async def construct(_id, _settings, _admission_registry, _api_key_number): return provider manager = ProviderRuntimeManager( diff --git a/tests/runtime/test_progressive_startup.py b/tests/runtime/test_progressive_startup.py index 69a9e2aa..6ae55247 100644 --- a/tests/runtime/test_progressive_startup.py +++ b/tests/runtime/test_progressive_startup.py @@ -48,7 +48,7 @@ async def test_http_and_admin_serve_while_selected_catalog_is_held(): provider.list_model_infos = AsyncMock(side_effect=list_models) - async def construct(_provider_id, _settings, _admission_registry): + async def construct(_provider_id, _settings, _admission_registry, _api_key_number): return provider manager = ProviderRuntimeManager( @@ -97,7 +97,7 @@ async def test_independent_provider_and_cancelled_waiter_share_initialization(): return_value=frozenset({ProviderModelInfo("one")}) ) - async def construct(provider_id, _settings, _admission_registry): + async def construct(provider_id, _settings, _admission_registry, _api_key_number): return slow if provider_id == "groq" else fast manager = ProviderRuntimeManager( @@ -158,7 +158,7 @@ async def test_failed_discovery_construction_recovers_in_same_generation(refresh ) attempts = 0 - async def construct(_id, _settings, _admission_registry): + async def construct(_id, _settings, _admission_registry, _api_key_number): nonlocal attempts attempts += 1 if attempts == 1: @@ -200,7 +200,7 @@ async def test_recovery_wait_uses_remaining_budget_without_orphaning_waiter( attempts = 0 construction = None - async def construct(_id, _settings, _admission_registry): + async def construct(_id, _settings, _admission_registry, _api_key_number): nonlocal attempts, construction attempts += 1 if attempts == 1: @@ -291,7 +291,9 @@ async def test_late_retired_discovery_cannot_replace_current_metadata(): provider.list_model_infos = AsyncMock(side_effect=models) - async def construct(_provider_id, _settings, _admission_registry): + async def construct( + _provider_id, _settings, _admission_registry, _api_key_number + ): return provider return ProviderRuntime( @@ -339,7 +341,7 @@ async def test_complete_catalog_publication_preserves_prior_file_until_all_attem provider.list_model_infos = AsyncMock(side_effect=models) - async def construct(_provider_id, _settings, _admission_registry): + async def construct(_provider_id, _settings, _admission_registry, _api_key_number): return provider manager = ProviderRuntimeManager( @@ -389,7 +391,7 @@ async def test_catalog_wait_follows_replacement( provider.list_model_infos = AsyncMock(side_effect=models) - async def construct(_id, _settings, _admission_registry): + async def construct(_id, _settings, _admission_registry, _api_key_number): return provider return ProviderRuntime( @@ -461,7 +463,7 @@ async def test_catalog_waiter_exit_preserves_shared_discovery( provider.list_model_infos = AsyncMock(side_effect=models) - async def construct(_id, _settings, _admission_registry): + async def construct(_id, _settings, _admission_registry, _api_key_number): return provider manager = ProviderRuntimeManager( diff --git a/tests/runtime/test_provider_manager.py b/tests/runtime/test_provider_manager.py index 58921de1..486e6cfe 100644 --- a/tests/runtime/test_provider_manager.py +++ b/tests/runtime/test_provider_manager.py @@ -25,10 +25,12 @@ class FakeRuntime(ProviderRuntime): self.provider = MagicMock() self.provider.list_model_infos = AsyncMock(return_value=frozenset()) - def is_cached(self, provider_id: str) -> bool: + def is_cached(self, provider_id: str, api_key_number: int = 1) -> bool: return provider_id == "cached" - async def resolve_provider(self, provider_id: str) -> BaseProvider: + async def resolve_provider( + self, provider_id: str, api_key_number: int = 1 + ) -> BaseProvider: return cast(BaseProvider, self.provider) async def cleanup(self) -> None: