feat(backends,storage,core,markdown,infra): claude agent sdk backend, session store, transcript seeding, runner isolation
This commit is contained in:
@@ -0,0 +1,385 @@
|
||||
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 == {}
|
||||
@@ -0,0 +1,64 @@
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from claude_agent_sdk.testing import run_session_store_conformance
|
||||
from sqlalchemy import delete
|
||||
|
||||
from beaver_gateway.storage import Database, PostgresSessionStore, TranscriptEntry
|
||||
|
||||
POSTGRES_URL = os.environ.get("BEAVER_TEST_DATABASE_URL")
|
||||
|
||||
|
||||
async def test_sqlite_conformance() -> None:
|
||||
root = Path(tempfile.mkdtemp(prefix="beaver-store-"))
|
||||
counter = 0
|
||||
|
||||
async def make_store() -> PostgresSessionStore:
|
||||
nonlocal counter
|
||||
counter += 1
|
||||
db = Database(f"sqlite:///{root / f'{counter}.db'}")
|
||||
await db.create_all()
|
||||
return PostgresSessionStore(db)
|
||||
|
||||
await run_session_store_conformance(make_store)
|
||||
|
||||
|
||||
@pytest.mark.skipif(POSTGRES_URL is None, reason="BEAVER_TEST_DATABASE_URL not set")
|
||||
async def test_postgres_conformance() -> None:
|
||||
assert POSTGRES_URL is not None
|
||||
db = Database(POSTGRES_URL)
|
||||
await db.create_all()
|
||||
|
||||
async def make_store() -> PostgresSessionStore:
|
||||
async with db.session() as session:
|
||||
await session.execute(delete(TranscriptEntry))
|
||||
await session.commit()
|
||||
return PostgresSessionStore(db)
|
||||
|
||||
try:
|
||||
await run_session_store_conformance(make_store)
|
||||
finally:
|
||||
await db.dispose()
|
||||
|
||||
|
||||
async def test_append_dedups_by_uuid_and_keeps_floats() -> None:
|
||||
db = Database(f"sqlite:///{tempfile.mkdtemp(prefix='beaver-store-')}/d.db")
|
||||
await db.create_all()
|
||||
store = PostgresSessionStore(db)
|
||||
key = {"project_key": "p", "session_id": "s"}
|
||||
batch = [
|
||||
{"type": "x", "uuid": "a", "n": 1.5},
|
||||
{"type": "y", "n": 2},
|
||||
{"type": "x", "uuid": "a", "n": 999},
|
||||
]
|
||||
await store.append(key, batch)
|
||||
await store.append(key, batch)
|
||||
loaded = await store.load(key)
|
||||
assert loaded == [
|
||||
{"type": "x", "uuid": "a", "n": 1.5},
|
||||
{"type": "y", "n": 2},
|
||||
{"type": "y", "n": 2},
|
||||
]
|
||||
assert isinstance(loaded[0]["n"], float)
|
||||
@@ -0,0 +1,129 @@
|
||||
from beaver_gateway.core.transcript import (
|
||||
CLI_VERSION,
|
||||
build_entries,
|
||||
messages_from_entries,
|
||||
)
|
||||
|
||||
HISTORY = [
|
||||
{"role": "user", "content": "Write бобёр to notes.txt"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "hm", "signature": "sig"},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_1",
|
||||
"name": "Write",
|
||||
"input": {"file_path": "notes.txt"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_1",
|
||||
"content": "ok",
|
||||
"is_error": False,
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [{"type": "text", "text": "Done."}]},
|
||||
{"role": "user", "content": [{"type": "text", "text": "thanks"}]},
|
||||
{"role": "assistant", "content": "np"},
|
||||
]
|
||||
|
||||
|
||||
def test_round_trip() -> None:
|
||||
entries = build_entries(HISTORY, session_id="sid", cwd="/vault", model="claude-x")
|
||||
assert messages_from_entries(entries) == [
|
||||
*HISTORY[:5],
|
||||
{"role": "assistant", "content": [{"type": "text", "text": "np"}]},
|
||||
]
|
||||
|
||||
|
||||
def test_chain_and_shape() -> None:
|
||||
entries = build_entries(HISTORY, session_id="sid", cwd="/vault", model="claude-x")
|
||||
assert [e["type"] for e in entries] == [
|
||||
"user",
|
||||
"assistant",
|
||||
"assistant",
|
||||
"user",
|
||||
"assistant",
|
||||
"user",
|
||||
"assistant",
|
||||
]
|
||||
assert entries[0]["parentUuid"] is None
|
||||
for prev, cur in zip(entries, entries[1:], strict=False):
|
||||
assert cur["parentUuid"] == prev["uuid"]
|
||||
assert len({e["uuid"] for e in entries}) == len(entries)
|
||||
|
||||
user = entries[0]
|
||||
assert user["promptSource"] == "sdk"
|
||||
assert user["entrypoint"] == "sdk-py"
|
||||
assert user["permissionMode"] == "bypassPermissions"
|
||||
assert user["sessionId"] == "sid"
|
||||
assert user["cwd"] == "/vault"
|
||||
assert user["version"] == CLI_VERSION
|
||||
|
||||
thinking, tool_use = entries[1], entries[2]
|
||||
assert thinking["message"]["id"] == tool_use["message"]["id"]
|
||||
assert thinking["requestId"] == tool_use["requestId"]
|
||||
assert tool_use["message"]["stop_reason"] == "tool_use"
|
||||
assert tool_use["message"]["content"] == [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_1",
|
||||
"name": "Write",
|
||||
"input": {"file_path": "notes.txt"},
|
||||
}
|
||||
]
|
||||
|
||||
result = entries[3]
|
||||
assert result["sourceToolAssistantUUID"] == tool_use["uuid"]
|
||||
assert result["promptId"] == user["promptId"]
|
||||
assert result["toolUseResult"] == "ok"
|
||||
assert result["message"]["content"][0]["is_error"] is False
|
||||
|
||||
final = entries[4]
|
||||
assert final["message"]["stop_reason"] == "end_turn"
|
||||
assert final["message"]["usage"]["input_tokens"] == 0
|
||||
|
||||
assert entries[5]["promptId"] != user["promptId"]
|
||||
|
||||
|
||||
def test_tool_results_split_and_merged_back() -> None:
|
||||
history = [
|
||||
{"role": "user", "content": "go"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": "a", "name": "Read", "input": {}},
|
||||
{"type": "tool_use", "id": "b", "name": "Read", "input": {}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "tool_result", "tool_use_id": "a", "content": "1"},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "b",
|
||||
"content": [{"type": "text", "text": "2"}],
|
||||
},
|
||||
],
|
||||
},
|
||||
]
|
||||
entries = build_entries(history, session_id="s", cwd="/", model="m")
|
||||
assert [e["type"] for e in entries] == [
|
||||
"user",
|
||||
"assistant",
|
||||
"assistant",
|
||||
"user",
|
||||
"user",
|
||||
]
|
||||
assert entries[3]["parentUuid"] == entries[1]["uuid"]
|
||||
assert entries[4]["parentUuid"] == entries[2]["uuid"]
|
||||
assert entries[4]["toolUseResult"] == "2"
|
||||
assert messages_from_entries(entries) == history
|
||||
Reference in New Issue
Block a user