diff --git a/myagents/commands/ollama.py b/myagents/commands/ollama.py index 6305f18..a6ff152 100644 --- a/myagents/commands/ollama.py +++ b/myagents/commands/ollama.py @@ -33,11 +33,10 @@ from pathlib import Path import click from rich.console import Console -from myagents.ollama_adapter import DEFAULT_PORT +from myagents.ollama_adapter import listen_url console = Console() -_HEALTH = f"http://127.0.0.1:{DEFAULT_PORT}/health" _UID = os.getuid() _LOG_DIR = Path.home() / ".xiaohe" / "logs" @@ -45,12 +44,20 @@ _LOG_DIR = Path.home() / ".xiaohe" / "logs" def health_ok() -> bool: """True when the adapter endpoint responds.""" try: - with urllib.request.urlopen(_HEALTH, timeout=1) as resp: + with urllib.request.urlopen(f"{listen_url()}/health", timeout=1) as resp: return resp.status == 200 except OSError: return False +def _wait_healthy(*, attempts: int = 20, delay: float = 0.5) -> bool: + for _ in range(attempts): + if health_ok(): + return True + time.sleep(delay) + return False + + # ── macOS: launchd ─────────────────────────────────────────────────────── LABEL = "team.xiaohe.ollama-adapter" @@ -102,7 +109,7 @@ def _mac_write() -> None: def _mac_start() -> None: if _mac_loaded(): - return + _mac_stop() result = subprocess.run( ["launchctl", "bootstrap", f"gui/{_UID}", str(PLIST_PATH)], capture_output=True, @@ -171,6 +178,7 @@ def _sysd_write() -> None: def _sysd_start() -> None: + already = _sysd_state()[1] == "active" _sysd_write() # Best-effort linger so the service survives logout (needs no login # session); harmless when it fails on systems without logind. @@ -180,6 +188,8 @@ def _sysd_start() -> None: capture_output=True, ) _systemctl("enable", "--now", UNIT_NAME, check=True) + if already: + _systemctl("restart", UNIT_NAME, check=True) def _sysd_stop() -> None: @@ -238,11 +248,7 @@ def ensure_running() -> bool: backend = _backend() backend.write() backend.start() - for _ in range(20): # up to ~10s - if health_ok(): - return True - time.sleep(0.5) - return False + return _wait_healthy() @click.group("ollama", invoke_without_command=True) @@ -255,17 +261,14 @@ def ollama_cmd(ctx: click.Context) -> None: @ollama_cmd.command("up") def ollama_up() -> None: - """Install and start the platform service, waiting for health.""" - if health_ok(): - console.print("[green]ollama-adapter already running.[/green]") - return + """Install (or bounce) the platform service, waiting for health.""" backend = _backend() backend.write() backend.start() - if ensure_running(): + if _wait_healthy(): console.print( f"[bold green]ollama-adapter up[/bold green] " - f"(http://127.0.0.1:{DEFAULT_PORT} -> Ollama, {backend.name})" + f"({listen_url()} -> Ollama, {backend.name})" ) else: raise click.ClickException( @@ -297,9 +300,9 @@ def ollama_status() -> None: """Show service and health state.""" up = health_ok() endpoint = ( - f"[green]up[/green] http://127.0.0.1:{DEFAULT_PORT}" + f"[green]up[/green] {listen_url()}" if up - else f"[red]down[/red] http://127.0.0.1:{DEFAULT_PORT}" + else f"[red]down[/red] {listen_url()}" ) console.print(f" endpoint : {endpoint}") try: diff --git a/myagents/ollama_adapter.py b/myagents/ollama_adapter.py index a7bae2b..e70a4cc 100644 --- a/myagents/ollama_adapter.py +++ b/myagents/ollama_adapter.py @@ -1,17 +1,15 @@ -"""Anthropic -> Ollama normalization adapter for Claude Code. +"""Anthropic -> Ollama proxy with system-message hoisting. -Claude Code injects some system content (e.g. the agent-types list) as -``role: "system"`` messages inside ``messages[]``, mid-conversation. Qwen-family -chat templates require a leading system message, so Ollama's ``/v1/messages`` -responds HTTP 500 ("System message must be at the beginning"). This adapter -moves any in-messages system content into the top-level Anthropic ``system`` -field before forwarding to Ollama. +Ollama already speaks ``/v1/messages``. Claude Code still injects +``role: "system"`` turns inside ``messages[]`` (agent-types, hooks, skills). +Qwen-family chat templates require a leading system message, so Ollama +returns HTTP 500. This process is a thin HTTP proxy: hoist those turns +into the top-level Anthropic ``system`` field, then forward every other +request unchanged. -Runs standalone: - - python -m myagents.ollama_adapter # 127.0.0.1:8199 - -Stdlib only (no fastapi/uvicorn) to keep myagents a lightweight CLI package. +Listen: ``XIAOHE_OLLAMA_ADAPTER_HOST`` / ``XIAOHE_OLLAMA_ADAPTER_PORT`` +(default 127.0.0.1:8199). Upstream: ``OLLAMA_HOST`` (default +http://127.0.0.1:11434), same as the ollama CLI (scheme optional). """ from __future__ import annotations @@ -20,141 +18,251 @@ import contextlib import json import logging import os -from http.client import HTTPConnection +from http.client import HTTPConnection, HTTPSConnection from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from urllib.parse import urlsplit +from urllib.request import Request, urlopen DEFAULT_HOST = os.environ.get("XIAOHE_OLLAMA_ADAPTER_HOST", "127.0.0.1") DEFAULT_PORT = int(os.environ.get("XIAOHE_OLLAMA_ADAPTER_PORT", "8199")) -# Respect OLLAMA_HOST just like the ollama CLI. -OLLAMA_BASE = os.environ.get("OLLAMA_HOST", "http://127.0.0.1:11434") +OLLAMA_DEFAULT_PORT = 11434 +UPSTREAM_TIMEOUT = int(os.environ.get("XIAOHE_OLLAMA_ADAPTER_TIMEOUT", "600")) +OLLAMA_BASE = os.environ.get("OLLAMA_HOST", f"http://127.0.0.1:{OLLAMA_DEFAULT_PORT}") +_LOOPBACK = frozenset({"127.0.0.1", "localhost", "::1", "0.0.0.0"}) _log = logging.getLogger("myagents.ollama_adapter") +def listen_url() -> str: + """URL Claude Code should use as ``ANTHROPIC_BASE_URL``.""" + return f"http://{DEFAULT_HOST}:{DEFAULT_PORT}" + + +def is_adapter_endpoint(base_url: str) -> bool: + """True when *base_url* points at this adapter's listen address.""" + raw = (base_url or "").strip() + if not raw: + return False + if "://" not in raw: + raw = f"http://{raw}" + left = urlsplit(raw) + right = urlsplit(listen_url()) + left_host = (left.hostname or "").lower() + right_host = (right.hostname or "").lower() + if left_host in _LOOPBACK and right_host in _LOOPBACK: + hosts_match = True + else: + hosts_match = left_host == right_host + left_scheme = (left.scheme or "http").lower() + left_port = left.port or (443 if left_scheme == "https" else 80) + right_port = right.port or DEFAULT_PORT + return hosts_match and left_port == right_port + + +def upstream_target(base: str | None = None) -> tuple[str, str, int]: + """Return ``(scheme, host, port)`` for an Ollama base URL or host:port.""" + raw = (base if base is not None else OLLAMA_BASE).strip() + if not raw: + raw = f"http://127.0.0.1:{OLLAMA_DEFAULT_PORT}" + if "://" not in raw: + raw = f"http://{raw}" + parts = urlsplit(raw) + scheme = (parts.scheme or "http").lower() + host = parts.hostname or "127.0.0.1" + if parts.port: + port = parts.port + elif scheme == "https": + port = 443 + else: + port = OLLAMA_DEFAULT_PORT + return scheme, host, port + + +def parse_num_ctx(parameters: str | None) -> int | None: + """Read Ollama's runtime ``num_ctx`` from a ``/api/show`` parameters blob.""" + if not parameters: + return None + for line in parameters.splitlines(): + parts = line.split() + if len(parts) >= 2 and parts[0].lower() == "num_ctx": + try: + n = int(parts[1]) + except ValueError: + continue + if n >= 1024: + return n + return None + + +def runtime_num_ctx(model: str, base: str | None = None) -> int | None: + """Ollama's allocated context (``num_ctx``), not architecture max. + + ``/api/tags`` reports ``details.context_length`` = 262144 for Qwen3.5 + even when the Modelfile capped ``num_ctx`` at 32k. Claude Code must + be told the runtime cap or it will send prompts the daemon rejects. + """ + name = (model or "").strip() + if not name: + return None + scheme, host, port = upstream_target(base) + req = Request( + f"{scheme}://{host}:{port}/api/show", + data=json.dumps({"name": name}).encode(), + headers={"content-type": "application/json"}, + method="POST", + ) + try: + with urlopen(req, timeout=2) as resp: + data = json.loads(resp.read().decode()) + except (OSError, json.JSONDecodeError, ValueError): + return None + if not isinstance(data, dict): + return None + raw = data.get("parameters") + return parse_num_ctx(raw if isinstance(raw, str) else None) + + +def _system_blocks(raw: object) -> list[dict]: + """Normalize a top-level ``system`` value or message content to blocks.""" + if isinstance(raw, str) and raw: + return [{"type": "text", "text": raw}] + if not isinstance(raw, list): + return [] + blocks: list[dict] = [] + for item in raw: + if isinstance(item, str) and item: + blocks.append({"type": "text", "text": item}) + elif ( + isinstance(item, dict) + and item.get("type") == "text" + and item.get("text") + ): + blocks.append(dict(item)) + return blocks + + def normalize_system(body: dict) -> dict: """Move ``role: "system"`` messages into the top-level ``system`` field. Mutates and returns *body*. Non-system messages keep their order. + Extra fields on text blocks (e.g. ``cache_control``) are preserved. """ - raw_sys = body.get("system") - if isinstance(raw_sys, list): - sys_blocks = [b for b in raw_sys if isinstance(b, dict)] - elif isinstance(raw_sys, str) and raw_sys: - sys_blocks = [{"type": "text", "text": raw_sys}] - else: - sys_blocks = [] - + sys_blocks = _system_blocks(body.get("system")) kept: list = [] for message in body.get("messages") or []: if isinstance(message, dict) and message.get("role") == "system": - content = message.get("content") - if isinstance(content, list): - for block in content: - if ( - isinstance(block, dict) - and block.get("type") == "text" - and block.get("text") - ): - sys_blocks.append( - {"type": "text", "text": block["text"]} - ) + sys_blocks.extend(_system_blocks(message.get("content"))) else: kept.append(message) - body["system"] = sys_blocks body["messages"] = kept return body +def _upstream_conn() -> HTTPConnection: + scheme, host, port = upstream_target() + if scheme == "https": + return HTTPSConnection(host, port, timeout=UPSTREAM_TIMEOUT) + return HTTPConnection(host, port, timeout=UPSTREAM_TIMEOUT) + + class _Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" - # -- plumbing ----------------------------------------------------------- - def _json(self, status: int, data: dict) -> None: body = json.dumps(data).encode() self.send_response(status) self.send_header("content-type", "application/json") self.send_header("content-length", str(len(body))) + self.send_header("connection", "close") + self.close_connection = True self.end_headers() self.wfile.write(body) - def _forward(self, raw: bytes, path: str) -> None: - """POST to Ollama and stream the response through unchanged.""" - parts = urlsplit(OLLAMA_BASE) - headers = { - "x-api-key": "ollama", - "content-type": "application/json", - "anthropic-version": self.headers.get( - "anthropic-version", "2023-06-01" - ), - "content-length": str(len(raw)), - } - conn = HTTPConnection( - parts.hostname or "127.0.0.1", parts.port or 80, timeout=600 - ) + def _outgoing_headers(self, raw: bytes) -> dict[str, str]: + headers: dict[str, str] = {} + if raw: + headers["content-type"] = ( + self.headers.get("content-type") or "application/json" + ) + headers["content-length"] = str(len(raw)) + for key in ( + "x-api-key", + "authorization", + "anthropic-version", + "anthropic-beta", + ): + val = self.headers.get(key) + if val: + headers[key] = val + if "x-api-key" not in headers and "authorization" not in headers: + headers["x-api-key"] = "ollama" + return headers + + def _relay_response(self, resp) -> None: + """Copy upstream status/body and always terminate the client response. + + HTTP/1.1 keep-alive without Content-Length leaves the client waiting + forever (UI stuck on "working"). + """ + self.send_response(resp.status) + content_type = resp.getheader("content-type") + if content_type: + self.send_header("content-type", content_type) + content_length = resp.getheader("content-length") + if content_length: + self.send_header("content-length", content_length) + else: + self.send_header("cache-control", "no-cache") + self.send_header("connection", "close") + self.close_connection = True + self.end_headers() try: - conn.request("POST", path, body=raw, headers=headers) + while True: + chunk = resp.read(65536) + if not chunk: + break + self.wfile.write(chunk) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + return + + def _forward(self, method: str, path: str, raw: bytes) -> None: + headers = self._outgoing_headers(raw) + conn = _upstream_conn() + try: + conn.request(method, path, body=raw or None, headers=headers) resp = conn.getresponse() except OSError as exc: _log.warning("upstream %s failed: %s", OLLAMA_BASE, exc) self._json(502, {"error": {"message": f"upstream error: {exc}"}}) return - self.send_response(resp.status) - content_type = resp.getheader("content-type") - if content_type: - self.send_header("content-type", content_type) - self.send_header("cache-control", "no-cache") - self.end_headers() - while True: - chunk = resp.read(65536) - if not chunk: - break - try: - self.wfile.write(chunk) - self.wfile.flush() - except (BrokenPipeError, ConnectionResetError): - break - conn.close() - - # -- routes ------------------------------------------------------------- + try: + self._relay_response(resp) + finally: + conn.close() def do_GET(self) -> None: - if self.path == "/health": + if self.path.split("?", 1)[0] == "/health": self._json(200, {"status": "ok", "upstream": OLLAMA_BASE}) return - if self.path == "/v1/models": - # Anthropic-style model list so /model switching works offline. - try: - from urllib.request import urlopen + self._forward("GET", self.path, b"") - with urlopen(OLLAMA_BASE + "/api/tags", timeout=5) as resp: - data = json.loads(resp.read().decode()) - names = [m["name"] for m in data.get("models", [])] - except OSError: - names = [] - self._json( - 200, - {"data": [{"id": n, "object": "model"} for n in names]}, - ) - return - self._json(404, {"error": {"message": "not found"}}) + def do_HEAD(self) -> None: + self._forward("HEAD", self.path, b"") def do_POST(self) -> None: - if self.path.split("?", 1)[0] != "/v1/messages": - self._json(404, {"error": {"message": "not found"}}) - return length = int(self.headers.get("content-length") or 0) raw = self.rfile.read(length) if length else b"" - try: - body = json.loads(raw) - except json.JSONDecodeError: - body = None - if isinstance(body, dict): - normalize_system(body) - raw = json.dumps(body).encode() - self._forward(raw, self.path) + if self.path.split("?", 1)[0] == "/v1/messages": + try: + body = json.loads(raw) + except json.JSONDecodeError: + body = None + if isinstance(body, dict): + normalize_system(body) + raw = json.dumps(body).encode() + self._forward("POST", self.path, raw) def log_message(self, format: str, *args: object) -> None: # noqa: A002 if args: @@ -166,8 +274,9 @@ def main() -> None: logging.basicConfig(level=logging.INFO) server = ThreadingHTTPServer((DEFAULT_HOST, DEFAULT_PORT), _Handler) server.daemon_threads = True + scheme, host, port = upstream_target() print( - f"ollama-adapter: http://{DEFAULT_HOST}:{DEFAULT_PORT} -> {OLLAMA_BASE}", + f"ollama-adapter: {listen_url()} -> {scheme}://{host}:{port}", flush=True, ) with contextlib.suppress(KeyboardInterrupt): diff --git a/tests/test_ollama.py b/tests/test_ollama.py index 3546a09..32f7c9b 100644 --- a/tests/test_ollama.py +++ b/tests/test_ollama.py @@ -1,12 +1,28 @@ """Tests for myagents.commands.ollama and myagents.ollama_adapter.""" +from __future__ import annotations + +import json import subprocess +import threading +from contextlib import contextmanager +from http.client import HTTPConnection +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import click import pytest +import myagents.ollama_adapter as adapter_mod from myagents.commands import ollama -from myagents.ollama_adapter import normalize_system +from myagents.ollama_adapter import ( + _Handler, + is_adapter_endpoint, + listen_url, + normalize_system, + parse_num_ctx, + runtime_num_ctx, + upstream_target, +) class TestNormalizeSystem: @@ -40,11 +56,355 @@ class TestNormalizeSystem: assert [m["role"] for m in out["messages"]] == ["user", "user"] assert [m["content"] for m in out["messages"]] == ["a", "b"] - def test_string_system_untouched(self) -> None: + def test_string_system_becomes_block(self) -> None: body = {"system": "top", "messages": [{"role": "user", "content": "a"}]} out = normalize_system(body) assert out["system"] == [{"type": "text", "text": "top"}] + def test_hoists_string_content_on_system_message(self) -> None: + body = { + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "system", "content": "agent-types"}, + {"role": "user", "content": "ping"}, + ] + } + out = normalize_system(body) + assert [m["role"] for m in out["messages"]] == ["user", "user"] + assert out["system"] == [{"type": "text", "text": "agent-types"}] + + def test_preserves_cache_control_on_text_blocks(self) -> None: + body = { + "system": [ + { + "type": "text", + "text": "SYS", + "cache_control": {"type": "ephemeral"}, + } + ], + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "more", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ], + } + out = normalize_system(body) + assert out["system"][0]["cache_control"] == {"type": "ephemeral"} + assert out["system"][1] == { + "type": "text", + "text": "more", + "cache_control": {"type": "ephemeral"}, + } + assert out["messages"] == [] + + def test_drops_empty_system_turns(self) -> None: + body = { + "messages": [ + {"role": "system", "content": ""}, + {"role": "user", "content": "hi"}, + ] + } + out = normalize_system(body) + assert out["messages"] == [{"role": "user", "content": "hi"}] + assert out["system"] == [] + + +class TestListenAndUpstream: + def test_listen_url(self) -> None: + assert listen_url() == ( + f"http://{adapter_mod.DEFAULT_HOST}:{adapter_mod.DEFAULT_PORT}" + ) + + def test_is_adapter_endpoint_loopback_aliases(self) -> None: + port = adapter_mod.DEFAULT_PORT + assert is_adapter_endpoint(f"http://127.0.0.1:{port}") + assert is_adapter_endpoint(f"http://127.0.0.1:{port}/") + assert is_adapter_endpoint(f"http://localhost:{port}") + assert not is_adapter_endpoint("http://127.0.0.1:11434") + assert not is_adapter_endpoint("") + assert not is_adapter_endpoint("https://api.deepseek.com/anthropic") + + def test_upstream_target_defaults_to_ollama_port(self) -> None: + assert upstream_target("http://127.0.0.1") == ( + "http", + "127.0.0.1", + 11434, + ) + assert upstream_target("127.0.0.1:11434") == ( + "http", + "127.0.0.1", + 11434, + ) + assert upstream_target("https://example.com") == ( + "https", + "example.com", + 443, + ) + assert upstream_target("https://example.com:8443") == ( + "https", + "example.com", + 8443, + ) + + +class TestRuntimeNumCtx: + def test_parse_num_ctx(self) -> None: + assert parse_num_ctx("num_ctx 32768\n") == 32768 + assert parse_num_ctx("temperature 1.0\nnum_ctx 65536") == 65536 + assert parse_num_ctx(None) is None + assert parse_num_ctx("") is None + assert parse_num_ctx("temperature 1.0") is None + + def test_runtime_num_ctx_reads_show(self, monkeypatch) -> None: + class _Resp: + def read(self) -> bytes: + return b'{"parameters":"num_ctx 65536\\n"}' + + def __enter__(self): + return self + + def __exit__(self, *args: object) -> None: + return None + + monkeypatch.setattr(adapter_mod, "urlopen", lambda *a, **k: _Resp()) + assert runtime_num_ctx("qwen3.5:4b-ctx64k") == 65536 + + def test_runtime_num_ctx_missing_model(self) -> None: + assert runtime_num_ctx("") is None + + +def _serve(handler: type[BaseHTTPRequestHandler]) -> ThreadingHTTPServer: + httpd = ThreadingHTTPServer(("127.0.0.1", 0), handler) + threading.Thread(target=httpd.serve_forever, daemon=True).start() + return httpd + + +@contextmanager +def _running(handler: type[BaseHTTPRequestHandler]): + httpd = _serve(handler) + try: + yield httpd + finally: + httpd.shutdown() + httpd.server_close() + + +@contextmanager +def _adapter_against(upstream_handler, monkeypatch): + with _running(upstream_handler) as upstream: + monkeypatch.setattr( + adapter_mod, + "OLLAMA_BASE", + f"http://127.0.0.1:{upstream.server_address[1]}", + ) + with _running(_Handler) as proxy: + yield proxy + + +def _recorder( + store: dict, + *, + status: int = 200, + body: bytes = b'{"ok":true}', + send_length: bool = True, +): + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def _capture(self) -> None: + n = int(self.headers.get("content-length") or 0) + raw = self.rfile.read(n) if n else b"" + store["method"] = self.command + store["path"] = self.path + store["body"] = raw + store["headers"] = {k.lower(): v for k, v in self.headers.items()} + + def _reply(self, head: bool = False) -> None: + self.send_response(status) + self.send_header("content-type", "application/json") + if send_length: + self.send_header("content-length", str(len(body))) + else: + self.send_header("connection", "close") + self.close_connection = True + self.end_headers() + if not head: + self.wfile.write(body) + + def do_GET(self) -> None: + self._capture() + self._reply() + + def do_POST(self) -> None: + self._capture() + self._reply() + + def do_HEAD(self) -> None: + self._capture() + self._reply(head=True) + + def log_message(self, format: str, *args: object) -> None: + pass + + return Handler + + +def _request( + proxy: ThreadingHTTPServer, + method: str, + path: str, + body: bytes = b"", + headers: dict[str, str] | None = None, +) -> tuple[int, bytes, str | None]: + host, port = proxy.server_address + conn = HTTPConnection(host, port, timeout=2) + hdrs = dict(headers or {}) + if body and "content-type" not in {k.lower() for k in hdrs}: + hdrs["content-type"] = "application/json" + conn.request(method, path, body=body or None, headers=hdrs) + resp = conn.getresponse() + raw = resp.read() + conn.close() + return resp.status, raw, resp.getheader("content-length") + + +class TestProxy: + def test_health_is_local(self, monkeypatch) -> None: + store: dict = {} + with _adapter_against(_recorder(store), monkeypatch) as proxy: + status, raw, _length = _request(proxy, "GET", "/health") + assert status == 200 + assert json.loads(raw)["status"] == "ok" + assert store == {} + + def test_get_models_is_forwarded(self, monkeypatch) -> None: + store: dict = {} + payload = b'{"object":"list","data":[]}' + with _adapter_against( + _recorder(store, body=payload), monkeypatch + ) as proxy: + status, raw, _length = _request(proxy, "GET", "/v1/models") + assert status == 200 + assert raw == payload + assert store["method"] == "GET" + assert store["path"] == "/v1/models" + + def test_post_messages_hoists_system_and_keeps_query( + self, monkeypatch + ) -> None: + store: dict = {} + payload = { + "model": "x", + "max_tokens": 1, + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "system", "content": "agent-types"}, + {"role": "user", "content": "ping"}, + ], + } + with _adapter_against(_recorder(store), monkeypatch) as proxy: + status, _raw, _length = _request( + proxy, + "POST", + "/v1/messages?beta=true", + json.dumps(payload).encode(), + ) + assert status == 200 + assert store["path"] == "/v1/messages?beta=true" + forwarded = json.loads(store["body"]) + assert [m["role"] for m in forwarded["messages"]] == ["user", "user"] + assert forwarded["system"] == [ + {"type": "text", "text": "agent-types"} + ] + + def test_other_post_paths_are_forwarded_unnormalized( + self, monkeypatch + ) -> None: + store: dict = {} + payload = b'{"model":"x"}' + with _adapter_against(_recorder(store), monkeypatch) as proxy: + status, _raw, _length = _request( + proxy, + "POST", + "/v1/messages/count_tokens?beta=true", + payload, + ) + assert status == 200 + assert store["path"] == "/v1/messages/count_tokens?beta=true" + assert store["body"] == payload + + def test_forwards_client_api_key(self, monkeypatch) -> None: + store: dict = {} + with _adapter_against(_recorder(store), monkeypatch) as proxy: + _request( + proxy, + "POST", + "/v1/messages", + b'{"messages":[]}', + headers={"x-api-key": "from-client"}, + ) + assert store["headers"]["x-api-key"] == "from-client" + + def test_upstream_error_is_502(self, monkeypatch) -> None: + monkeypatch.setattr(adapter_mod, "OLLAMA_BASE", "http://127.0.0.1:1") + with _running(_Handler) as proxy: + status, raw, _length = _request( + proxy, "POST", "/v1/messages", b"{}" + ) + assert status == 502 + assert "upstream error" in json.loads(raw)["error"]["message"] + + +class TestForwardTerminates: + def test_completes_when_upstream_sends_content_length( + self, monkeypatch + ) -> None: + store: dict = {} + body = b'{"id":"msg_1","type":"message"}' + with _adapter_against( + _recorder(store, body=body), monkeypatch + ) as proxy: + status, raw, length = _request( + proxy, "POST", "/v1/messages", b'{"messages":[]}' + ) + assert status == 200 + assert json.loads(raw)["id"] == "msg_1" + assert length == str(len(raw)) + + def test_completes_when_upstream_omits_content_length( + self, monkeypatch + ) -> None: + store: dict = {} + body = b'data: {"type":"message_stop"}\n\n' + with _adapter_against( + _recorder(store, body=body, send_length=False), monkeypatch + ) as proxy: + status, raw, _length = _request( + proxy, "POST", "/v1/messages?beta=true", b'{"messages":[]}' + ) + assert status == 200 + assert b"message_stop" in raw + + def test_forwards_upstream_4xx_and_completes(self, monkeypatch) -> None: + store: dict = {} + body = b'{"error":{"message":"bad request"}}' + with _adapter_against( + _recorder(store, status=400, body=body), monkeypatch + ) as proxy: + status, raw, length = _request( + proxy, "POST", "/v1/messages", b'{"messages":[]}' + ) + assert status == 400 + assert json.loads(raw)["error"]["message"] == "bad request" + assert length == str(len(raw)) + class TestSystemdUnit: def test_unit_text_smoke(self) -> None: @@ -62,7 +422,7 @@ class TestSystemdUnit: kwargs.pop("capture_output", None) kwargs.pop("text", None) calls.append(list(cmd)) - return subprocess.CompletedProcess(cmd, 0) + return subprocess.CompletedProcess(cmd, 0, stdout="") monkeypatch.setattr(ollama.subprocess, "run", fake_run) monkeypatch.setattr(ollama, "UNIT_DIR", tmp_path) @@ -80,6 +440,34 @@ class TestSystemdUnit: ] in calls assert any(c[0] == "loginctl" and "enable-linger" in c for c in calls) assert (tmp_path / ollama.UNIT_NAME).exists() + assert ["systemctl", "--user", "restart", ollama.UNIT_NAME] not in calls + + def test_sysd_start_restarts_when_already_active( + self, monkeypatch, tmp_path + ) -> None: + calls: list[list[str]] = [] + + def fake_run(cmd, **kwargs): + kwargs.pop("capture_output", None) + kwargs.pop("text", None) + calls.append(list(cmd)) + if cmd[:3] == ["systemctl", "--user", "is-active"]: + return subprocess.CompletedProcess( + cmd, 0, stdout="active\n", stderr="" + ) + if cmd[:3] == ["systemctl", "--user", "is-enabled"]: + return subprocess.CompletedProcess( + cmd, 0, stdout="enabled\n", stderr="" + ) + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + monkeypatch.setattr(ollama.subprocess, "run", fake_run) + monkeypatch.setattr(ollama, "UNIT_DIR", tmp_path) + monkeypatch.setattr(ollama, "UNIT_PATH", tmp_path / ollama.UNIT_NAME) + + ollama._sysd_start() + + assert ["systemctl", "--user", "restart", ollama.UNIT_NAME] in calls def test_sysd_stop_disables(self, monkeypatch) -> None: calls: list[list[str]] = [] @@ -103,9 +491,39 @@ class TestSystemdUnit: assert ["systemctl", "--user", "daemon-reload"] in calls +class TestMacLaunchd: + def test_mac_start_bootstraps_when_not_loaded(self, monkeypatch) -> None: + calls: list[list[str]] = [] + monkeypatch.setattr(ollama, "_mac_loaded", lambda: False) + + def fake_run(cmd, **kwargs): + calls.append(list(cmd)) + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + monkeypatch.setattr(ollama.subprocess, "run", fake_run) + ollama._mac_start() + assert calls[0][:2] == ["launchctl", "bootstrap"] + assert not any(c[1] == "bootout" for c in calls) + + def test_mac_start_rebounds_when_loaded(self, monkeypatch) -> None: + loaded = {"v": True} + calls: list[list[str]] = [] + monkeypatch.setattr(ollama, "_mac_loaded", lambda: loaded["v"]) + + def fake_run(cmd, **kwargs): + if len(cmd) > 1 and cmd[1] == "bootout": + loaded["v"] = False + calls.append(list(cmd)) + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + monkeypatch.setattr(ollama.subprocess, "run", fake_run) + ollama._mac_start() + assert calls[0][:2] == ["launchctl", "bootout"] + assert calls[1][:2] == ["launchctl", "bootstrap"] + + class TestBackendDispatch: def test_darwin_launchd(self) -> None: - # Test runs on macOS; _platform() returns the real platform. if ollama._platform() == "darwin": assert ollama._backend().name == "launchd"