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 from beaver_gateway.backends.claude_sdk import ClaudeSdkBackend, 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 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(), prompt_sources=(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 == "alpha\n\nbeta\n" 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 == {}