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": {}} 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")] ) yield StreamEvent(uuid="u", session_id="s", event=start) 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.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 == "\nalpha\n\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