myagents/tests/test_ollama.py

"""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 (
    _Handler,
    is_adapter_endpoint,
    listen_url,
    normalize_system,
    parse_num_ctx,
    runtime_num_ctx,
    upstream_target,
)


class TestNormalizeSystem:
    def test_moves_system_message_to_top_level(self) -> None:
        body = {
            "system": [{"type": "text", "text": "SYS1"}],
            "messages": [
                {"role": "user", "content": [{"type": "text", "text": "hi"}]},
                {
                    "role": "system",
                    "content": [{"type": "text", "text": "SYS2"}],
                },
            ],
        }
        out = normalize_system(body)
        assert [m["role"] for m in out["messages"]] == ["user"]
        assert [b["text"] for b in out["system"]] == ["SYS1", "SYS2"]

    def test_keeps_non_system_order(self) -> None:
        body = {
            "messages": [
                {"role": "user", "content": "a"},
                {
                    "role": "system",
                    "content": [{"type": "text", "text": "s"}],
                },
                {"role": "user", "content": "b"},
            ]
        }
        out = normalize_system(body)
        assert [m["role"] for m in out["messages"]] == ["user", "user"]
        assert [m["content"] for m in out["messages"]] == ["a", "b"]

    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:
        text = ollama._sysd_unit_text()
        assert "ExecStart=" in text
        assert "-m myagents.ollama_adapter" in text
        assert "WantedBy=default.target" in text

    def test_sysd_start_enables_and_reloads(
        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))
            return subprocess.CompletedProcess(cmd, 0, stdout="")

        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", "daemon-reload"] in calls
        assert [
            "systemctl",
            "--user",
            "enable",
            "--now",
            ollama.UNIT_NAME,
        ] 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]] = []

        def fake_run(cmd, **kwargs):
            kwargs.pop("capture_output", None)
            kwargs.pop("text", None)
            calls.append(list(cmd))
            return subprocess.CompletedProcess(cmd, 0)

        monkeypatch.setattr(ollama.subprocess, "run", fake_run)
        ollama._sysd_stop()

        assert [
            "systemctl",
            "--user",
            "disable",
            "--now",
            ollama.UNIT_NAME,
        ] in calls
        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:
        if ollama._platform() == "darwin":
            assert ollama._backend().name == "launchd"

    def test_linux_systemd(self, monkeypatch) -> None:
        monkeypatch.setattr(ollama, "_platform", lambda: "linux")
        assert ollama._backend().name == "systemd"

    def test_unsupported_platform_raises(self, monkeypatch) -> None:
        monkeypatch.setattr(ollama, "_platform", lambda: "win32")
        with pytest.raises(click.ClickException):
            ollama._backend()