beaver-gateway/tests/test_claude_sdk_backend.py

import asyncio
import tempfile
from pathlib import Path
from typing import Any

import pytest
from anthropic.types import (
    RawContentBlockDeltaEvent,
    RawContentBlockStartEvent,
    RawContentBlockStopEvent,
    RawMessageDeltaEvent,
    RawMessageStartEvent,
    RawMessageStopEvent,
)
from claude_agent_sdk import (
    AssistantMessage,
    InMemorySessionStore,
    ResultMessage,
    StreamEvent,
    TextBlock,
    ToolResultBlock,
    ToolUseBlock,
    UserMessage,
    project_key_for_directory,
)

from beaver_gateway.agents.base import ExposedMcp
from beaver_gateway.agents.claude import ClaudeAgent, ClaudeOptions, Prompts
from beaver_gateway.backends.claude_sdk import (
    ClaudeSdkBackend,
    RunnerConfig,
    UsageEvent,
    fingerprint,
)
from beaver_gateway.core.transcript import messages_from_entries
from beaver_gateway.core.turn_capture import TurnCapture


def _stream(index: int, text: str) -> list[StreamEvent]:
    def ev(event: dict[str, Any]) -> StreamEvent:
        return StreamEvent(uuid="u", session_id="s", event=event)

    return [
        ev(
            {
                "type": "content_block_start",
                "index": index,
                "content_block": {"type": "text", "text": ""},
            }
        ),
        ev(
            {
                "type": "content_block_delta",
                "index": index,
                "delta": {"type": "text_delta", "text": text},
            }
        ),
        ev({"type": "content_block_stop", "index": index}),
    ]


def _result(session_id: str) -> ResultMessage:
    return ResultMessage(
        subtype="success",
        duration_ms=10,
        duration_api_ms=5,
        is_error=False,
        num_turns=1,
        session_id=session_id,
        stop_reason="end_turn",
        total_cost_usd=0.01,
        usage={
            "input_tokens": 3,
            "output_tokens": 7,
            "cache_read_input_tokens": 100,
            "cache_creation_input_tokens": 20,
        },
    )


class FakeClient:
    instances: list["FakeClient"] = []

    def __init__(self, options: Any) -> None:
        self.options = options
        self.prompts: list[str] = []
        self.connected = False
        self.session_id = options.resume or "fresh-session"
        FakeClient.instances.append(self)

    async def connect(self) -> None:
        self.connected = True

    async def query(self, prompt: str) -> None:
        self.prompts.append(prompt)

    async def interrupt(self) -> None:
        self.interrupted = True

    async def receive_response(self):
        start = {
            "type": "message_start",
            "message": {
                "usage": {"input_tokens": 5, "cache_read_input_tokens": 20_000}
            },
        }
        yield StreamEvent(uuid="u", session_id="s", event=start)
        for e in _stream(0, "calling "):
            yield e
        tool = ToolUseBlock(id="toolu_1", name="Read", input={"path": "x"})
        yield StreamEvent(
            uuid="u",
            session_id="s",
            event={
                "type": "content_block_start",
                "index": 1,
                "content_block": {"type": "tool_use", "id": "toolu_1", "name": "Read"},
            },
        )
        yield StreamEvent(
            uuid="u", session_id="s", event={"type": "content_block_stop", "index": 1}
        )
        yield AssistantMessage(content=[TextBlock(text="calling "), tool], model="m")
        yield UserMessage(
            content=[ToolResultBlock(tool_use_id="toolu_1", content="42")]
        )
        second = {
            "type": "message_start",
            "message": {
                "usage": {
                    "input_tokens": 7,
                    "cache_read_input_tokens": 20_000,
                    "cache_creation_input_tokens": 3_000,
                }
            },
        }
        yield StreamEvent(uuid="u", session_id="s", event=second)
        for e in _stream(0, "done"):
            yield e
        yield AssistantMessage(content=[TextBlock(text="done")], model="m")
        yield StreamEvent(
            uuid="sub",
            session_id="s",
            event={"type": "content_block_stop", "index": 5},
            parent_tool_use_id="toolu_9",
        )
        yield _result(self.session_id)

    async def disconnect(self) -> None:
        self.connected = False


@pytest.fixture
def cwd() -> Path:
    return Path(tempfile.mkdtemp(prefix="beaver-sdk-"))


def _backend(
    cwd: Path, store: InMemorySessionStore, sink=None, **agent_kwargs
) -> ClaudeSdkBackend:
    FakeClient.instances.clear()
    agent = ClaudeAgent(
        name="a",
        model="claude-x",
        system_prompt="hi",
        cwd=cwd,
        options=ClaudeOptions(effort="low"),
        **agent_kwargs,
    )
    return ClaudeSdkBackend(
        agent=agent,
        mcp_internal_urls={"firefly": "http://127.0.0.1:1/mcp/firefly/"},
        session_store=store,
        mcp_tool_names={"firefly": ["list_account", "delete_account", "store_account"]},
        usage_sink=sink,
        client_factory=FakeClient,
        work_dir=cwd / "work",
    )


async def _drain(events) -> list[Any]:
    return [e async for e in events]


async def test_stream_envelope_and_index_rebase(cwd: Path) -> None:
    backend = _backend(
        cwd, InMemorySessionStore(), expose_mcps=(ExposedMcp(name="firefly"),)
    )
    capture = TurnCapture()
    events = await _drain(
        backend.complete(
            agent=backend.agent,
            messages=[{"role": "user", "content": "hi"}],
            conversation_id="conv-1",
            capture=capture,
        )
    )
    assert isinstance(events[0], RawMessageStartEvent)
    assert isinstance(events[-1], RawMessageStopEvent)
    assert isinstance(events[-2], RawMessageDeltaEvent)
    assert events[-2].usage.output_tokens == 7
    assert events[-2].usage.cache_read_input_tokens == 100
    starts = [e.index for e in events if isinstance(e, RawContentBlockStartEvent)]
    assert starts == [0, 1, 2]
    stops = [e.index for e in events if isinstance(e, RawContentBlockStopEvent)]
    assert stops == [0, 1, 2]
    deltas = [e.delta.text for e in events if isinstance(e, RawContentBlockDeltaEvent)]
    assert deltas == ["calling ", "done"]

    assert capture.session_id == "fresh-session"
    assert capture.usage is not None and capture.usage.cost_usd == 0.01
    assert capture.usage.context_tokens == 23_007
    assert capture.synthesized_messages == [
        {
            "role": "assistant",
            "content": [
                {"type": "text", "text": "calling "},
                {
                    "type": "tool_use",
                    "id": "toolu_1",
                    "name": "Read",
                    "input": {"path": "x"},
                },
            ],
        },
        {
            "role": "user",
            "content": [
                {"type": "tool_result", "tool_use_id": "toolu_1", "content": "42"}
            ],
        },
        {"role": "assistant", "content": [{"type": "text", "text": "done"}]},
    ]
    client = FakeClient.instances[0]
    assert client.prompts == ["hi"]
    opts = client.options
    assert opts.setting_sources == []
    assert opts.strict_mcp_config is True
    assert opts.permission_mode == "bypassPermissions"
    assert opts.mcp_servers == {
        "firefly": {"type": "http", "url": "http://127.0.0.1:1/mcp/firefly/"}
    }
    assert opts.session_store is not None
    assert opts.resume is None
    assert Path(opts.cli_path).exists()
    wrapper = Path(opts.cli_path).read_text()
    assert "CLAUDE_" in wrapper and "DATABASE_URL" not in wrapper


async def test_conversation_reuses_live_session(cwd: Path) -> None:
    backend = _backend(cwd, InMemorySessionStore())
    for i in range(2):
        await _drain(
            backend.complete(
                agent=backend.agent,
                messages=[{"role": "user", "content": f"turn {i}"}],
                conversation_id="conv-1",
            )
        )
    assert len(FakeClient.instances) == 1
    assert FakeClient.instances[0].prompts == ["turn 0", "turn 1"]
    assert backend.sessions["conv-1"]["turns"] == 2


async def test_history_is_seeded_into_store_and_resumed(cwd: Path) -> None:
    store = InMemorySessionStore()
    backend = _backend(cwd, store)
    history = [
        {"role": "user", "content": "first"},
        {"role": "assistant", "content": [{"type": "text", "text": "reply"}]},
        {"role": "user", "content": "second"},
    ]
    await _drain(
        backend.complete(
            agent=backend.agent, messages=history, conversation_id="conv-2"
        )
    )
    client = FakeClient.instances[0]
    assert client.options.resume is not None
    key = {
        "project_key": project_key_for_directory(str(cwd)),
        "session_id": client.options.resume,
    }
    entries = store.get_entries(key)
    assert messages_from_entries(entries) == history[:2]
    assert client.prompts == ["second"]


async def test_known_session_id_is_resumed_without_seeding(cwd: Path) -> None:
    store = InMemorySessionStore()
    backend = _backend(cwd, store)
    history = [
        {"role": "user", "content": "first"},
        {"role": "assistant", "content": "reply"},
        {"role": "user", "content": "second"},
    ]
    await _drain(
        backend.complete(
            agent=backend.agent,
            messages=history,
            conversation_id="conv-3",
            session_id="known-sid",
        )
    )
    assert FakeClient.instances[0].options.resume == "known-sid"
    assert store.size == 0


async def test_stateless_caller_hits_same_session_next_turn(cwd: Path) -> None:
    backend = _backend(cwd, InMemorySessionStore())
    first = [{"role": "user", "content": "hi"}]
    await _drain(backend.complete(agent=backend.agent, messages=first))
    follow_up = [
        *first,
        {
            "role": "assistant",
            "content": [
                {"type": "text", "text": "calling "},
                {"type": "text", "text": "done"},
            ],
        },
        {"role": "user", "content": "more"},
    ]
    assert list(backend.sessions) == [fingerprint(follow_up[:-1])]
    await _drain(backend.complete(agent=backend.agent, messages=follow_up))
    assert len(FakeClient.instances) == 1
    assert FakeClient.instances[0].prompts == ["hi", "more"]


async def test_mcp_deny_and_usage_sink(cwd: Path) -> None:
    seen: list[UsageEvent] = []

    async def sink(event: UsageEvent) -> None:
        seen.append(event)

    backend = _backend(
        cwd,
        InMemorySessionStore(),
        sink,
        expose_mcps=(ExposedMcp(name="firefly", deny=("delete_*",)),),
    )
    await _drain(
        backend.complete(
            agent=backend.agent,
            messages=[{"role": "user", "content": "x"}],
            conversation_id="c",
        )
    )
    assert FakeClient.instances[0].options.disallowed_tools == [
        "mcp__firefly__delete_account"
    ]
    assert len(seen) == 1
    assert seen[0].usage.input_tokens == 3
    assert seen[0].conversation_id == "c"
    assert seen[0].session_id == "fresh-session"


async def test_skill_sets_become_sorted_plugins(cwd: Path) -> None:
    sets = cwd / "skills"
    for name in ("zeta", "общие"):
        (sets / name / "demo").mkdir(parents=True)
        (sets / name / "demo" / "SKILL.md").write_text("---\nname: demo\n---\n")
    backend = _backend(
        cwd, InMemorySessionStore(), skill_sets=(sets / "zeta", sets / "общие")
    )
    await _drain(
        backend.complete(
            agent=backend.agent,
            messages=[{"role": "user", "content": "x"}],
            conversation_id="c",
        )
    )
    plugins = FakeClient.instances[0].options.plugins
    assert [p["type"] for p in plugins] == ["local", "local"]
    paths = [Path(p["path"]) for p in plugins]
    assert [p.name for p in paths] == ["zeta", "общие"]
    for path in paths:
        assert (path / ".claude-plugin" / "plugin.json").exists()
        assert (path / "skills" / "demo" / "SKILL.md").exists()
    assert FakeClient.instances[0].options.skills == "all"


async def test_prompt_sources_are_assembled(cwd: Path) -> None:
    (cwd / "a.md").write_text("alpha\n")
    (cwd / "b.md").write_text("\nbeta\n\n")
    backend = _backend(
        cwd,
        InMemorySessionStore(),
        prompts=Prompts(deep=(("role", cwd / "a.md"), cwd / "b.md")),
    )
    await _drain(
        backend.complete(
            agent=backend.agent,
            messages=[{"role": "user", "content": "x"}],
            conversation_id="c",
        )
    )
    assert (
        FakeClient.instances[0].options.system_prompt
        == "<role>\nalpha\n</role>\n\nbeta\n"
    )


async def test_runner_user_lands_in_wrapper(cwd: Path) -> None:
    import os
    import pwd

    me = pwd.getpwuid(os.getuid())
    backend = _backend(cwd, InMemorySessionStore())
    backend._runner = RunnerConfig(user=me.pw_name, home=cwd)
    backend._uid, backend._gid = me.pw_uid, me.pw_gid
    await _drain(
        backend.complete(
            agent=backend.agent,
            messages=[{"role": "user", "content": "x"}],
            conversation_id="c",
        )
    )
    opts = FakeClient.instances[0].options
    assert opts.env["HOME"] == str(cwd)
    assert opts.env["CLAUDE_CONFIG_DIR"] == str(cwd / ".claude")
    wrapper = Path(opts.cli_path).read_text()
    assert f"UID = {me.pw_uid}" in wrapper
    assert "os.setuid(UID)" in wrapper


async def test_close_disconnects(cwd: Path) -> None:
    backend = _backend(cwd, InMemorySessionStore())
    async with backend:
        await _drain(
            backend.complete(
                agent=backend.agent,
                messages=[{"role": "user", "content": "x"}],
                conversation_id="c",
            )
        )
    assert FakeClient.instances[0].connected is False
    assert backend.sessions == {}


class GatedClient(FakeClient):
    """Streams one delta, then waits for ``gate`` before finishing the turn."""

    gate: asyncio.Event

    async def receive_response(self):
        yield StreamEvent(
            uuid="u", session_id="s", event={"type": "message_start", "message": {}}
        )
        for e in _stream(0, "first"):
            yield e
        await GatedClient.gate.wait()
        yield AssistantMessage(content=[TextBlock(text="first")], model="m")
        yield _result(self.session_id)


async def test_deltas_reach_the_caller_before_the_turn_ends(cwd: Path) -> None:
    """The turn is not buffered: a delta is observable while the CLI still runs."""
    GatedClient.gate = asyncio.Event()
    backend = _backend(cwd, InMemorySessionStore())
    backend._factory = GatedClient
    events = backend.complete(
        agent=backend.agent,
        messages=[{"role": "user", "content": "hi"}],
        conversation_id="conv-live",
    )
    seen: list[Any] = []
    async for ev in events:
        seen.append(ev)
        if isinstance(ev, RawContentBlockDeltaEvent):
            break
    assert seen[-1].delta.text == "first"
    assert not GatedClient.gate.is_set()
    GatedClient.gate.set()
    rest = [e async for e in events]
    assert isinstance(rest[-1], RawMessageStopEvent)


async def test_policy_hook_denies_and_audits(cwd: Path) -> None:
    from beaver_gateway.core.policy import Deny, ToolCall

    def no_days(c: ToolCall):
        p = c.path()
        if c.tool == "Write" and p is not None and "📅 дни" in p.parts:
            return Deny(reason="дневник только читать")
        c.state["seen"] = c.state.get("seen", 0) + 1
        return None

    audits = []

    async def sink(a):
        audits.append(a)

    backend = _backend(cwd, InMemorySessionStore(), policy=(no_days,))
    backend._audit_sink = sink
    await _drain(
        backend.complete(
            agent=backend.agent,
            messages=[{"role": "user", "content": "hi"}],
            conversation_id="conv-1",
        )
    )
    hook = FakeClient.instances[0].options.hooks["PreToolUse"][0].hooks[0]
    denied = await hook(
        {"tool_name": "Write", "tool_input": {"file_path": "/vault/📅 дни/x.md"}},
        "t1",
        {},
    )
    assert denied["hookSpecificOutput"]["permissionDecision"] == "deny"
    assert (
        denied["hookSpecificOutput"]["permissionDecisionReason"]
        == "дневник только читать"
    )
    allowed = await hook(
        {"tool_name": "Write", "tool_input": {"file_path": "/vault/мета/бобер/x.md"}},
        "t2",
        {},
    )
    assert allowed == {}
    await hook(
        {"tool_name": "Read", "tool_input": {"file_path": "/vault/a.md"}}, "t3", {}
    )
    assert backend.live("conv-1").state["seen"] == 2
    assert [(a.tool, a.decision) for a in audits] == [
        ("Write", "deny"),
        ("Write", "allow"),
        ("Read", "allow"),
    ]
    assert audits[0].brief == "/vault/📅 дни/x.md"
    assert audits[0].kind == "deep" and audits[0].conversation == "conv-1"


async def test_no_policy_no_hooks(cwd: Path) -> None:
    backend = _backend(cwd, InMemorySessionStore())
    await _drain(
        backend.complete(
            agent=backend.agent, messages=[{"role": "user", "content": "hi"}]
        )
    )
    assert FakeClient.instances[0].options.hooks is None