diff options
| -rw-r--r-- | llamachat/providers.py | 72 | ||||
| -rwxr-xr-x | test_llamachat.py | 103 |
2 files changed, 175 insertions, 0 deletions
diff --git a/llamachat/providers.py b/llamachat/providers.py index 0d2d9f7..051a931 100644 --- a/llamachat/providers.py +++ b/llamachat/providers.py @@ -155,3 +155,75 @@ def apply_filter(provider: Provider, listed: list[str]) -> list[str]: if not needles: return list(listed) return [m for m in listed if any(n in m.lower() for n in needles)] + + +def _run_pass(cmd: list[str], timeout: int) -> str: + """Run `pass show NAME` and return its stdout.""" + result = subprocess.run( + cmd, capture_output=True, text=True, timeout=timeout, check=True + ) + return result.stdout + + +class KeyResolver: + """Resolves api_key fields, lazily and once per process. + + Lazy matters: `pass` needs the GPG key, so a session that only touches + local models must never trigger a pinentry. Resolved values stay in + memory and are never written anywhere. + """ + + def __init__(self, runner=_run_pass): + self._runner = runner + self._cache: dict[tuple[str, str], str] = {} + + def resolve(self, provider: Provider) -> str: + """The bearer token for this provider, or '' when it needs none.""" + spec = provider.api_key + if not spec: + return "" + # Keyed by the spec as well as the name: the config can be reloaded + # onto the same resolver, and keying by name alone would keep serving + # the old key after the user edits the entry to fix it, so the fix + # would look like it did nothing until a restart. + cache_key = (provider.name, spec) + if cache_key in self._cache: + return self._cache[cache_key] + + if spec.startswith("env:"): + value = os.environ.get(spec[4:], "") + if not value: + raise KeyError_( + f"{provider.name}: environment variable {spec[4:]} is not set" + ) + elif spec.startswith("pass:"): + value = self._from_pass(provider, spec[5:]) + else: + value = spec + + self._cache[cache_key] = value + return value + + def _from_pass(self, provider: Provider, entry: str) -> str: + try: + out = self._runner(["pass", "show", entry], timeout=KEY_TIMEOUT) + except subprocess.TimeoutExpired: + raise KeyError_( + f"{provider.name}: `pass show {entry}` timed out after " + f"{KEY_TIMEOUT}s. Is a pinentry waiting for input?" + ) + except subprocess.CalledProcessError as exc: + # stderr is captured, so without this it is swallowed and the user + # gets an exit code where gpg told them exactly what was wrong. + # Only stderr is quoted, never stdout, which holds the secret. + detail = (exc.stderr or "").strip().splitlines() + reason = detail[-1] if detail else f"exit status {exc.returncode}" + raise KeyError_(f"{provider.name}: `pass show {entry}` failed: {reason}") + except Exception as exc: + raise KeyError_(f"{provider.name}: `pass show {entry}` failed: {exc}") + # A password store entry keeps the secret on the first line and + # metadata below it. + first = (out or "").strip().splitlines() + if not first or not first[0].strip(): + raise KeyError_(f"{provider.name}: `pass show {entry}` returned nothing") + return first[0].strip() diff --git a/test_llamachat.py b/test_llamachat.py index 8b722ae..08f31c9 100755 --- a/test_llamachat.py +++ b/test_llamachat.py @@ -16,6 +16,7 @@ import json import os +import subprocess import sys import tempfile from pathlib import Path @@ -1064,6 +1065,107 @@ def test_model_ids_and_filtering(): print("ok model ids and filtering") +def test_key_resolution(): + """api_key is prefix-dispatched, resolved lazily and cached.""" + from llamachat import providers + + resolver = providers.KeyResolver() + + # No key configured: no Authorization header, and nothing is run. + empty = providers.Provider(name="local", base_url="http://x.example.org") + assert resolver.resolve(empty) == "" + + # A literal key is used as-is. + literal = providers.Provider( + name="p", base_url="http://x.example.org", api_key="sk-test-not-a-real-key" + ) + assert resolver.resolve(literal) == "sk-test-not-a-real-key" + + # env: reads the environment. + os.environ["LLAMACHAT_TEST_KEY"] = "from-env" + env = providers.Provider( + name="e", base_url="http://x.example.org", + api_key="env:LLAMACHAT_TEST_KEY", + ) + assert resolver.resolve(env) == "from-env" + del os.environ["LLAMACHAT_TEST_KEY"] + + # A missing env var is an error naming the provider, not a silent "". + missing = providers.Provider( + name="gone", base_url="http://x.example.org", + api_key="env:LLAMACHAT_ABSENT_VAR", + ) + try: + resolver.resolve(missing) + assert False, "a missing env var must raise" + except providers.KeyError_ as exc: + assert "gone" in str(exc) + + # pass: shells out. Substitute the runner rather than requiring gpg. + calls = [] + + def fake_run(cmd, timeout): + calls.append((cmd, timeout)) + return "line-one\nline-two\n" + + passed = providers.Provider( + name="together", base_url="http://x.example.org", + api_key="pass:api/together", + ) + cached = providers.KeyResolver(runner=fake_run) + assert cached.resolve(passed) == "line-one" # first line only + assert calls[0][0] == ["pass", "show", "api/together"] + assert calls[0][1] == providers.KEY_TIMEOUT + + # Cached: a second resolve must not shell out again. + assert cached.resolve(passed) == "line-one" + assert len(calls) == 1 + + # But the cache follows the spec, not just the name. A reloaded config + # that points the same provider at a different entry must re-resolve, + # otherwise correcting a wrong entry appears to do nothing. + moved = providers.Provider( + name="together", base_url="http://x.example.org", + api_key="pass:api/together-corrected", + ) + assert cached.resolve(moved) == "line-one" + assert len(calls) == 2 + assert calls[1][0] == ["pass", "show", "api/together-corrected"] + + # A failing pass is reported, naming the provider. + def boom(cmd, timeout): + raise OSError("pass: entry not found") + + try: + providers.KeyResolver(runner=boom).resolve(passed) + assert False, "a failing pass must raise" + except providers.KeyError_ as exc: + assert "together" in str(exc) + + # A non-zero exit quotes gpg's stderr, which is the only useful part, and + # never stdout, which is where the secret would be. + def refused(cmd, timeout): + raise subprocess.CalledProcessError( + 2, cmd, output="sk-test-not-a-real-key\n", + stderr="gpg: decryption failed: No secret key\n", + ) + + try: + providers.KeyResolver(runner=refused).resolve(passed) + assert False, "a non-zero pass exit must raise" + except providers.KeyError_ as exc: + assert "No secret key" in str(exc) + assert "sk-test-not-a-real-key" not in str(exc) + + # Empty output is a failure too: an empty key would 401 confusingly. + try: + providers.KeyResolver(runner=lambda cmd, timeout: " \n").resolve(passed) + assert False, "empty pass output must raise" + except providers.KeyError_: + pass + print("ok api key resolution") + + class _FakeResponse: """Enough of an http.client response for urlopen's context manager.""" @@ -1803,6 +1905,7 @@ if __name__ == "__main__": test_config_defaults() test_provider_parsing() test_model_ids_and_filtering() + test_key_resolution() test_search_tool_schema() test_search_results_sanitising() test_tool_call_accumulation() |