feat: reset bot session context while retaining monitoring history

tests/unit_tests/api/service/test_session_context_reset.py

@@ -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'
+    )

tests/unit_tests/provider/test_session_manager.py

@@ -11,7 +11,7 @@ from __future__ import annotations
 import pytest
 import asyncio
 from types import SimpleNamespace
-from unittest.mock import Mock
+from unittest.mock import Mock, AsyncMock
 from importlib import import_module
 
 import langbot_plugin.api.entities.builtin.provider.session as provider_session
@@ -26,6 +26,13 @@ from langbot.pkg.pipeline.pool import (
 )
 
 
+@pytest.fixture(autouse=True)
+def reset_generation(monkeypatch):
+    lookup = AsyncMock(return_value=None)
+    monkeypatch.setattr('langbot.pkg.agent.runner.context_reset.get_reset_generation', lookup)
+    return lookup
+
+
 TEST_CONTEXT = ExecutionContext(
     instance_uuid='instance-test',
     workspace_uuid='workspace-test',
@@ -447,7 +454,7 @@ class TestSessionManagerWorkspaceIsolation:
             sessions.append(await manager.get_session(query))
 
         assert len(manager.session_list) == 2
-        assert sessions[0] not in manager.session_list
+        assert all(session is not sessions[0] for session in manager.session_list)
         assert sessions[1:] == manager.session_list
         assert await manager.get_session(queries[-1]) is manager.session_list[-1]
 
@@ -556,3 +563,18 @@ class TestSessionManagerWorkspaceIsolation:
         content = conversation.messages[0].content
         assert content[1].image_base64 is None
         assert content[2].file_base64 is None
+
+
+@pytest.mark.asyncio
+async def test_reset_generation_replaces_cached_session(reset_generation):
+    app = Mock()
+    app.instance_config.data = {'concurrency': {'session': 5}}
+    manager = get_session_module().SessionManager(app)
+    query = scoped_query()
+    previous = await manager.get_session(query)
+    assert await manager.get_session(query) is previous
+    reset_generation.return_value = 'new-context'
+    fresh = await manager.get_session(query)
+    assert fresh is not previous
+    assert not fresh.conversations
+    assert await manager.get_session(query) is fresh