feat: reset bot session context while retaining monitoring history

2026-10-21 07:55:18 +08:00 · 2026-10-10 22:35:27 +08:00
parent b8a7649ac2
commit 1c70218383
20 changed files with 408 additions and 5 deletions
@@ -0,0 +1,148 @@
+import contextlib
+import datetime
+import json
+from types import SimpleNamespace
+
+import pytest
+import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from langbot.pkg.api.http.service.monitoring import MonitoringService
+from langbot.pkg.agent.runner.context_reset import get_reset_generation
+from langbot.pkg.agent.runner.transcript_store import TranscriptStore
+from langbot.pkg.entity.persistence.base import Base
+from langbot.pkg.entity.persistence.monitoring import MonitoringSession
+from langbot.pkg.entity.persistence.runner_state import RunnerState
+from langbot.pkg.entity.persistence.agent_run import AgentRun
+from langbot.pkg.entity.persistence.transcript import Transcript
+
+
+class Persistence:
+    def __init__(self, engine):
+        self.engine = engine
+
+    def get_db_engine(self):
+        return self.engine
+
+    @contextlib.asynccontextmanager
+    async def tenant_uow(self, workspace):
+        async with self.engine.begin() as conn:
+            self.conn = conn
+            yield
+
+    async def execute_async(self, statement):
+        return await self.conn.execute(statement)
+
+
+@pytest.fixture
+async def context_service(tmp_path):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "reset.db"}')
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+        await conn.execute(
+            sa.insert(MonitoringSession).values(
+                workspace_uuid='w',
+                bot_id='b',
+                session_id='person_u',
+                bot_name='Bot',
+                pipeline_id='p',
+                pipeline_name='Pipeline',
+                start_time=datetime.datetime.utcnow(),
+                last_activity=datetime.datetime.utcnow(),
+            )
+        )
+        for workspace, bot, conversation, scope in [
+            ('w', 'b', 'person_u', 'conversation'),
+            ('w', 'b2', 'person_u', 'conversation'),
+            ('w2', 'b', 'person_u', 'conversation'),
+            ('w', 'b', 'person_other', 'conversation'),
+            ('w', 'b', 'person_u', 'actor'),
+        ]:
+            await conn.execute(
+                sa.insert(RunnerState).values(
+                    workspace_id=workspace,
+                    bot_id=bot,
+                    conversation_id=conversation,
+                    scope=scope,
+                    scope_key=f'{workspace}:{bot}:{conversation}:{scope}',
+                    state_key='checkpoint',
+                    runner_id='runner',
+                    binding_identity='binding',
+                    value_json=json.dumps('old'),
+                )
+            )
+    yield MonitoringService(SimpleNamespace(persistence_mgr=Persistence(engine))), engine
+    await engine.dispose()
+
+
+@pytest.mark.asyncio
+async def test_reset_preserves_logs_and_isolates_context(context_service):
+    service, engine = context_service
+    store = TranscriptStore(engine)
+    for bot in ['b', 'b2']:
+        await store.append_transcript(None, 'event', 'person_u', 'user', bot_id=bot, workspace_id='w', content='old')
+    await service.reset_session_context('w', 'b', 'person_u')
+    assert await get_reset_generation(engine, 'w', 'b', 'person_u')
+    async with engine.connect() as conn:
+        assert (await conn.execute(sa.select(sa.func.count()).select_from(Transcript))).scalar() == 2
+        assert (await conn.execute(sa.select(sa.func.count()).select_from(MonitoringSession))).scalar() == 1
+        assert (
+            await conn.execute(
+                sa.select(sa.func.count()).select_from(RunnerState).where(RunnerState.runner_id == 'runner')
+            )
+        ).scalar() == 4
+    page = await store.page_transcript('person_u', bot_id='b', workspace_id='w')
+    assert not page[0]
+    assert (await store.page_transcript('person_u', bot_id='b2', workspace_id='w'))[0]
+    await store.append_transcript(None, 'new', 'person_u', 'user', bot_id='b', workspace_id='w', content='fresh')
+    assert len((await store.page_transcript('person_u', bot_id='b', workspace_id='w'))[0]) == 1
+    await service.reset_session_context('w', 'b', 'person_u')
+    assert not (await store.page_transcript('person_u', bot_id='b', workspace_id='w'))[0]
+
+
+@pytest.mark.asyncio
+async def test_reset_refuses_wrong_scope_and_running_tasks(context_service):
+    service, engine = context_service
+    for workspace, bot in [('other', 'b'), ('w', 'other')]:
+        with pytest.raises(LookupError):
+            await service.reset_session_context(workspace, bot, 'person_u')
+    async with engine.begin() as conn:
+        await conn.execute(
+            sa.insert(AgentRun).values(
+                run_id='run',
+                event_id='evt',
+                agent_id='agent',
+                binding_id='binding',
+                runner_id='runner',
+                workspace_id='w',
+                bot_id='b',
+                conversation_id='person_u',
+                status='running',
+            )
+        )
+    with pytest.raises(RuntimeError):
+        await service.reset_session_context('w', 'b', 'person_u')
+    assert await get_reset_generation(engine, 'w', 'b', 'person_u') is None
+
+
+@pytest.mark.asyncio
+async def test_reset_starts_new_external_runner_conversation(context_service):
+    from unit_tests.agent.test_state_store import FakeBinding, FakeEventEnvelope, make_descriptor
+    from langbot.pkg.agent.runner.persistent_state_store import PersistentStateStore
+
+    service, engine = context_service
+    store = PersistentStateStore(engine)
+    event = FakeEventEnvelope(workspace_id='w', bot_id='b', conversation_id='person_u')
+    binding = FakeBinding()
+    descriptor = make_descriptor()
+    await store.apply_update_from_event(
+        event, binding, descriptor, 'conversation', 'external.conversation_id', 'remote-old', None
+    )
+    assert (await store.build_snapshot_from_event(event, binding, descriptor))['conversation'][
+        'external.conversation_id'
+    ] == 'remote-old'
+    await service.reset_session_context('w', 'b', 'person_u')
+    snapshot = await store.build_snapshot_from_event(event, binding, descriptor)
+    assert snapshot['conversation']['external.conversation_id'] == await get_reset_generation(
+        engine, 'w', 'b', 'person_u'
+    )