diff --git a/docs/configuration.md b/docs/configuration.md index 9c8f4e8..04d21b8 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -213,6 +213,7 @@ and `UserPromptSubmit` deny-verdicts are enforced. See [hooks.md](hooks.md). ```toml [custom_providers.local] base_url = "http://localhost:11434/v1" +api_key = "sk-..." # optional api_key_env = "LOCAL_API_KEY" # optional auth_policy = "none" # auto | required | none # headers = { X-Team = "infra" } diff --git a/src/lecode/auth.py b/src/lecode/auth.py index ce697fb..025f02d 100644 --- a/src/lecode/auth.py +++ b/src/lecode/auth.py @@ -64,8 +64,11 @@ def resolve_api_key( if value: key, source = value, "env" break - if key is None and config.llm.api_key: - key, source = config.llm.api_key, "config" + if key is None: + if custom is not None and custom.api_key: + key, source = custom.api_key, "config" + elif config.llm.api_key: + key, source = config.llm.api_key, "config" if key is None and policy == "required": raise AuthError( diff --git a/src/lecode/config/models.py b/src/lecode/config/models.py index 31897eb..db350de 100644 --- a/src/lecode/config/models.py +++ b/src/lecode/config/models.py @@ -254,6 +254,7 @@ class CustomProvider(BaseModel): model_config = ConfigDict(extra="ignore") base_url: str + api_key: str | None = None api_key_env: str | None = None headers: dict[str, str] = Field(default_factory=dict) auth_policy: AuthPolicy = "auto" diff --git a/src/lecode/setup_wizard.py b/src/lecode/setup_wizard.py index da3cc10..3a5c6f4 100644 --- a/src/lecode/setup_wizard.py +++ b/src/lecode/setup_wizard.py @@ -108,6 +108,7 @@ def import_from_pi(home: Path) -> dict[str, str]: if not base_url: return {} # anthropic/openai/… — no equivalent in lecode answers["provider"] = "custom" + answers["provider_name"] = provider answers["base_url"] = base_url key = _auth_key(auth, provider) if key: @@ -202,6 +203,7 @@ def import_from_opencode(home: Path) -> dict[str, Any]: base_url = options.get("baseURL", "") if isinstance(options, dict) else "" if base_url: answers["provider"] = "custom" + answers["provider_name"] = provider answers["base_url"] = base_url if model: answers["model"] = model @@ -235,7 +237,10 @@ def _import_summary(answers: dict[str, Any]) -> str: """One-line description of what an import found (key redacted).""" parts = [] if answers.get("provider"): - parts.append(f"provider {answers['provider']}") + provider = answers["provider"] + if provider == "custom": + provider = answers.get("provider_name") or "custom" + parts.append(f"provider {provider}") if answers.get("base_url"): parts.append(f"base_url {answers['base_url']}") if answers.get("model"): @@ -355,21 +360,34 @@ async def _ask_yes_no(session: PromptSession, message: str, default: bool = True def build_config(answers: dict[str, Any]) -> dict[str, Any]: """Assemble the raw config dict from wizard answers.""" + provider = answers["provider"] llm: dict[str, Any] = { - "provider": answers["provider"], + "provider": provider, "model": answers["model"], } - if answers.get("api_key"): - llm["api_key"] = answers["api_key"] - if answers.get("base_url"): - llm["base_url"] = answers["base_url"] + custom_providers: dict[str, Any] = {} + if provider == "custom": + name = answers.get("provider_name") or "custom" + llm["provider"] = name + entry: dict[str, Any] = {"base_url": answers.get("base_url", "")} + if answers.get("api_key"): + entry["api_key"] = answers["api_key"] + custom_providers[name] = entry + else: + if answers.get("api_key"): + llm["api_key"] = answers["api_key"] + if answers.get("base_url"): + llm["base_url"] = answers["base_url"] config: dict[str, Any] = { "schema_version": 1, "llm": llm, "notifications": {"enabled": answers["notifications"]}, } + if custom_providers: + config["custom_providers"] = custom_providers if answers.get("mcp_servers"): config["mcp"] = {"servers": answers["mcp_servers"]} + return config @@ -405,8 +423,21 @@ async def gather_answers(session: PromptSession, home: Path | None = None) -> di provider_choices: list[str] = list(PROVIDER_CHOICES) provider_default = imported.get("provider") provider = await _ask_choice(session, "Provider:", provider_choices, default=provider_default) + provider_name = imported.get("provider_name", "custom") base_url = imported.get("base_url", "") if provider == "custom": + while True: + provider_name = await _ask_text( + session, + "Provider name (an identifier for this endpoint)", + default=provider_name, + ) + if provider_name and provider_name != "openrouter": + break + if not provider_name: + print("error: a provider name is required") + else: + print("error: 'openrouter' is reserved for the built-in provider") while True: base_url = await _ask_text( session, "Base URL (OpenRouter-compatible)", default=base_url @@ -456,6 +487,7 @@ async def gather_answers(session: PromptSession, home: Path | None = None) -> di notifications = await _ask_yes_no(session, "Audio notifications?", default=True) answers: dict[str, Any] = { "provider": provider, + "provider_name": provider_name if provider == "custom" else "", "base_url": base_url, "api_key": api_key, "model": model, diff --git a/src/lecode/tui/feed.py b/src/lecode/tui/feed.py index bfc88c1..4d29db8 100644 --- a/src/lecode/tui/feed.py +++ b/src/lecode/tui/feed.py @@ -105,13 +105,10 @@ def stream_token(self, text: str, *, thinking: bool = False) -> None: if thinking: self._thinking_parts.append(text) return + self._stream_parts.append(text) if self.stream_sink is not None: - self._stream_parts.append(text) self.stream_sink(text) self._stream_printed = True - else: - self._console.print(text, end="", markup=False, highlight=False, soft_wrap=True) - self._stream_printed = True def _flush_stream(self) -> None: """Close the open stream: print the accumulated text as one whole print. @@ -119,16 +116,11 @@ def _flush_stream(self) -> None: With a sink, the live region is cleared first and the full text lands in the scrollback newline-terminated (safe under patch_stdout). """ - if not self._stream_printed: - return - if self.stream_sink is not None: - full = "".join(self._stream_parts) - self._stream_parts = [] - if self.stream_clear is not None: - self.stream_clear() - self._console.print(full, markup=False, highlight=False, soft_wrap=True) - else: - self._console.print() + full = "".join(self._stream_parts) + self._stream_parts = [] + if self.stream_sink is not None and self.stream_clear is not None and self._stream_printed: + self.stream_clear() + self._console.print(Markdown(full)) self._stream_printed = False def stream_end(self) -> None: @@ -172,8 +164,9 @@ def llm_response( elapsed_s: float = 0.0, ) -> None: """Log per-call usage and output tok/s over model-call time when available.""" - # The streamed answer text has no trailing newline yet — close it first. - self._flush_stream() + if self._streaming: + # The streamed answer text has no trailing newline yet — close it first. + self._flush_stream() line = ( f"[{self._stamp()}] ← {model} (round {turn})" f" · ↑{human_tokens(input_tokens)} in · ↓{human_tokens(output_tokens)} out" diff --git a/tests/test_auth.py b/tests/test_auth.py index 79db2e7..f62b70a 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -72,6 +72,15 @@ def test_custom_provider_falls_back_to_config_key(): assert (resolved.key, resolved.source) == ("config-key", "config") +def test_custom_provider_prefers_its_own_config_key(): + config = _config( + llm={"api_key": "global-key"}, + custom_providers={"myllm": {"base_url": "http://x", "api_key": "own-key"}}, + ) + resolved = resolve_api_key("myllm", config) + assert (resolved.key, resolved.source) == ("own-key", "config") + + def test_policy_required_raises_without_key(): config = _config(llm={"auth_policy": "required"}) with pytest.raises(AuthError, match="required"): diff --git a/tests/test_setup_wizard.py b/tests/test_setup_wizard.py index 7bd8b1d..a308f9c 100644 --- a/tests/test_setup_wizard.py +++ b/tests/test_setup_wizard.py @@ -89,21 +89,35 @@ async def test_wizard_config_file_is_owner_only(cfg_dir, clean_home): assert mode == 0o600 -async def test_wizard_custom_provider_asks_base_url(cfg_dir, clean_home): - session = FakeSession(["custom", "https://llm.local/v1", "local-key", "2", "n"]) +async def test_wizard_custom_provider_writes_named_entry(cfg_dir, clean_home): + session = FakeSession(["custom", "local", "https://llm.local/v1", "local-key", "2", "n"]) await run_wizard(session, home=clean_home) raw = tomllib.loads((cfg_dir / "config.toml").read_text()) - assert raw["llm"]["provider"] == "custom" - assert raw["llm"]["base_url"] == "https://llm.local/v1" + assert raw["llm"]["provider"] == "local" assert raw["llm"]["model"] == "openai/gpt-5-mini" # pick 2 + assert "api_key" not in raw["llm"] + assert raw["custom_providers"]["local"] == { + "base_url": "https://llm.local/v1", + "api_key": "local-key", + } -async def test_wizard_base_url_validated(cfg_dir, clean_home): - session = FakeSession(["2", "ftp://nope", "https://ok.example/v1", "", "1", ""]) +async def test_wizard_custom_provider_base_url_validated(cfg_dir, clean_home): + session = FakeSession(["2", "local", "ftp://nope", "https://ok.example/v1", "", "1", ""]) + await run_wizard(session, home=clean_home) + raw = tomllib.loads((cfg_dir / "config.toml").read_text()) + assert raw["custom_providers"]["local"]["base_url"] == "https://ok.example/v1" + assert "api_key" not in raw["custom_providers"]["local"] + + +async def test_wizard_custom_provider_name_validated(cfg_dir, clean_home): + session = FakeSession( + ["custom", "openrouter", "gemini", "https://gemini.example/v1", "k", "1", ""] + ) await run_wizard(session, home=clean_home) raw = tomllib.loads((cfg_dir / "config.toml").read_text()) - assert raw["llm"]["base_url"] == "https://ok.example/v1" - assert "api_key" not in raw["llm"] # custom tolerates an empty key + assert raw["llm"]["provider"] == "gemini" + assert "openrouter" not in raw["custom_providers"] async def test_wizard_key_required_loops_until_nonempty(cfg_dir, clean_home): @@ -136,6 +150,27 @@ def test_build_config_minimal_shape(): } +def test_build_config_named_custom_provider(): + raw = build_config( + { + "provider": "custom", + "provider_name": "gemini", + "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/", + "api_key": "gk", + "model": "gemini-2.5-flash", + "notifications": False, + } + ) + assert raw["llm"]["provider"] == "gemini" + assert raw["llm"]["model"] == "gemini-2.5-flash" + assert "api_key" not in raw["llm"] + assert "base_url" not in raw["llm"] + assert raw["custom_providers"]["gemini"] == { + "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/", + "api_key": "gk", + } + + async def test_model_menu_shows_context_and_price(cfg_dir, clean_home, capsys): session = FakeSession(["1", "sk-or-key", "2", "", "n"]) answers = await gather_answers(session, home=clean_home) @@ -394,6 +429,7 @@ def test_import_from_pi_custom_provider(tmp_path): ) assert import_from_pi(tmp_path) == { "provider": "custom", + "provider_name": "local", "base_url": "http://localhost:1234/v1", "api_key": "local-key", "model": "qwen3", @@ -436,6 +472,7 @@ def test_import_from_opencode_custom_base_url(tmp_path): ) assert import_from_opencode(tmp_path) == { "provider": "custom", + "provider_name": "corp", "base_url": "https://corp.example/v1", "model": "qwen3", }