llm-connect/tests/test_claude_code.py

from __future__ import annotations

from types import SimpleNamespace

from llm_connect.claude_code import ClaudeCodeAdapter
from llm_connect.config import LLMConfig
from llm_connect.models import RunConfig


def test_execute_prompt_passes_json_schema_to_claude_cli(monkeypatch):
    calls: dict[str, object] = {}

    def fake_run(cmd, input, capture_output, text, timeout):  # noqa: ANN001
        calls["cmd"] = cmd
        calls["input"] = input
        calls["capture_output"] = capture_output
        calls["text"] = text
        calls["timeout"] = timeout
        # With --output-format json the CLI returns an envelope.
        envelope = {
            "type": "result",
            "result": '{"summary":"ok","recommendations":[]}',
        }
        import json as _json
        return SimpleNamespace(returncode=0, stdout=_json.dumps(envelope), stderr="")

    monkeypatch.setattr("llm_connect.claude_code.subprocess.run", fake_run)
    adapter = ClaudeCodeAdapter(cli_path="/custom/claude")

    response = adapter.execute_prompt(
        "Produce a report.",
        RunConfig(
            timeout_seconds=42,
            model_params={"json_schema": {"type": "object"}},
        ),
    )

    assert calls["cmd"] == [
        "/custom/claude",
        "--print",
        "--json-schema",
        '{"type":"object"}',
        "--output-format",
        "json",
    ]
    assert calls["input"] == "Produce a report."
    assert calls["timeout"] == 42
    # Envelope's result field carries the schema-enforced JSON; the adapter
    # unwraps it before returning to the caller.
    assert response.content == '{"summary":"ok","recommendations":[]}'


def test_execute_prompt_unwraps_cli_json_envelope_result_field(monkeypatch):
    """With --output-format json the CLI wraps the model payload in an
    envelope. The adapter unwraps the textual result so the caller still
    sees the model's structured-output JSON, not the envelope."""
    def fake_run(cmd, input, capture_output, text, timeout):  # noqa: ANN001
        envelope = {
            "type": "result",
            "result": '{"summary":"ok","recommendations":[]}',
            "total_cost_usd": 0.001,
        }
        import json as _json
        return SimpleNamespace(
            returncode=0,
            stdout=_json.dumps(envelope),
            stderr="",
        )

    monkeypatch.setattr("llm_connect.claude_code.subprocess.run", fake_run)
    adapter = ClaudeCodeAdapter(cli_path="/custom/claude")

    response = adapter.execute_prompt(
        "Produce a report.",
        RunConfig(model_params={"json_schema": {"type": "object"}}),
    )

    assert response.content == '{"summary":"ok","recommendations":[]}'


def test_execute_prompt_prefers_json_field_over_prose_preamble(monkeypatch):
    """When the model adds a prose preamble in the envelope's primary text
    field but the schema-enforced JSON is in a different field, the adapter
    must find and return the JSON, not the preamble."""
    def fake_run(cmd, input, capture_output, text, timeout):  # noqa: ANN001
        envelope = {
            "type": "result",
            "result": "Triage report generated and returned via structured output. Key signals: healthy.",
            "structured_result": '{"summary":"healthy","recommendations":[]}',
            "total_cost_usd": 0.002,
        }
        import json as _json
        return SimpleNamespace(returncode=0, stdout=_json.dumps(envelope), stderr="")

    monkeypatch.setattr("llm_connect.claude_code.subprocess.run", fake_run)
    adapter = ClaudeCodeAdapter(cli_path="/custom/claude")

    response = adapter.execute_prompt(
        "Long triage prompt.",
        RunConfig(model_params={"json_schema": {"type": "object"}}),
    )

    assert response.content == '{"summary":"healthy","recommendations":[]}'


def test_execute_prompt_skips_envelope_metadata_keys(monkeypatch):
    """Metadata keys like `type`, `model`, `usage` must never be returned as
    the model payload, even if their values look JSON-like."""
    def fake_run(cmd, input, capture_output, text, timeout):  # noqa: ANN001
        envelope = {
            "type": '{"this":"is_metadata"}',  # decoy
            "usage": {"input_tokens": 5},      # decoy dict
            "result": '{"summary":"ok"}',
        }
        import json as _json
        return SimpleNamespace(returncode=0, stdout=_json.dumps(envelope), stderr="")

    monkeypatch.setattr("llm_connect.claude_code.subprocess.run", fake_run)
    adapter = ClaudeCodeAdapter(cli_path="/custom/claude")

    response = adapter.execute_prompt(
        "Prompt.", RunConfig(model_params={"json_schema": {"type": "object"}})
    )

    assert response.content == '{"summary":"ok"}'


def test_execute_prompt_no_unwrap_without_json_schema(monkeypatch):
    """Without --json-schema we do not pass --output-format json, so the
    envelope unwrap path stays inert and raw stdout passes through."""
    def fake_run(cmd, input, capture_output, text, timeout):  # noqa: ANN001
        return SimpleNamespace(
            returncode=0,
            stdout='{"result":"this is just stdout, not an envelope"}',
            stderr="",
        )

    monkeypatch.setattr("llm_connect.claude_code.subprocess.run", fake_run)
    adapter = ClaudeCodeAdapter(cli_path="/custom/claude")

    response = adapter.execute_prompt("Plain prompt.", RunConfig())

    assert response.content == '{"result":"this is just stdout, not an envelope"}'


def test_claude_code_adapter_prefers_env_cli_path(monkeypatch):
    monkeypatch.setenv("LLM_CONNECT_CLAUDE_CLI_PATH", "/home/me/bin/claude")

    adapter = ClaudeCodeAdapter(
        config=LLMConfig(provider="claude-code", claude_cli_path="claude")
    )

    assert adapter._cli_path == "/home/me/bin/claude"