wehub-resource-sync
110 lines
3.3 KiB
Python
110 lines
3.3 KiB
Python
"""Regression tests for chat.send middle-message edit branching."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
from opensquilla.engine.types import DoneEvent
|
|
from opensquilla.gateway.agent_tasks import get_agent_task_registry
|
|
from opensquilla.gateway.auth import Principal
|
|
from opensquilla.gateway.config import GatewayConfig
|
|
from opensquilla.gateway.rpc import RpcContext, get_dispatcher
|
|
from opensquilla.session.manager import SessionManager
|
|
from opensquilla.session.storage import SessionStorage
|
|
|
|
_PRINCIPAL = Principal(
|
|
role="operator", scopes=frozenset(["operator.admin"]), is_owner=True, authenticated=True
|
|
)
|
|
|
|
PARENT_KEY = "agent:main:webchat:editbranch01"
|
|
|
|
|
|
class _RecordingTurnRunner:
|
|
def __init__(self) -> None:
|
|
self.run_calls: list[dict[str, Any]] = []
|
|
self._locks: dict[str, asyncio.Lock] = {}
|
|
|
|
def _get_session_lock(self, session_key: str) -> asyncio.Lock:
|
|
return self._locks.setdefault(session_key, asyncio.Lock())
|
|
|
|
async def run(self, message: str, session_key: str, **kwargs: Any):
|
|
self.run_calls.append({"message": message, "session_key": session_key, **kwargs})
|
|
yield DoneEvent()
|
|
|
|
|
|
@pytest.fixture
|
|
def dispatcher():
|
|
return get_dispatcher()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def manager():
|
|
storage = SessionStorage(":memory:")
|
|
await storage.connect()
|
|
mgr = SessionManager(storage, inject_time_prefix=False)
|
|
yield mgr
|
|
await storage.close()
|
|
|
|
|
|
@pytest.fixture
|
|
def ctx(manager):
|
|
runner = _RecordingTurnRunner()
|
|
context = RpcContext(
|
|
conn_id="test-conn",
|
|
principal=_PRINCIPAL,
|
|
config=GatewayConfig(memory={"flush_enabled": False}),
|
|
turn_runner=runner,
|
|
)
|
|
context.session_manager = manager
|
|
return context
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_send_fork_before_message_returns_child_without_future_history(
|
|
dispatcher,
|
|
ctx,
|
|
manager,
|
|
):
|
|
await manager.create(PARENT_KEY, agent_id="main", display_name="Branch edit")
|
|
await manager.append_message(PARENT_KEY, "user", "A marker", token_count=1)
|
|
middle = await manager.append_message(PARENT_KEY, "user", "B marker", token_count=1)
|
|
await manager.append_message(PARENT_KEY, "user", "C marker", token_count=1)
|
|
|
|
res = await dispatcher.dispatch(
|
|
"r-chat-branch-edit",
|
|
"chat.send",
|
|
{
|
|
"sessionKey": PARENT_KEY,
|
|
"message": "B edited",
|
|
"forkBeforeMessageId": middle.message_id,
|
|
},
|
|
ctx,
|
|
)
|
|
|
|
task_key = res.payload.get("sessionKey", PARENT_KEY) if res.ok else PARENT_KEY
|
|
task = get_agent_task_registry().get(task_key) or get_agent_task_registry().get(PARENT_KEY)
|
|
if task is not None:
|
|
await task
|
|
|
|
assert res.ok is True
|
|
child_key = res.payload["sessionKey"]
|
|
assert child_key != PARENT_KEY
|
|
assert res.payload["key"] == child_key
|
|
|
|
parent_entries = await manager.get_transcript(PARENT_KEY)
|
|
assert [entry.content for entry in parent_entries] == [
|
|
"A marker",
|
|
"B marker",
|
|
"C marker",
|
|
]
|
|
|
|
child_entries = await manager.get_transcript(child_key)
|
|
assert [entry.content for entry in child_entries] == ["A marker", "B edited"]
|
|
|
|
assert ctx.turn_runner.run_calls[0]["session_key"] == child_key
|
|
assert ctx.turn_runner.run_calls[0]["message"] == "B edited"
|