fix: enforce agent run API permissions

tests/unit_tests/agent/test_history_event_api_auth.py

@@ -0,0 +1,146 @@
+"""Tests for AgentRunner history/event pull API authorization."""
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from langbot.pkg.agent.runner.session_registry import AgentRunSessionRegistry
+from langbot.pkg.plugin.handler import RuntimeConnectionHandler
+from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction
+
+from .conftest import make_resources
+
+
+class FakeConnection:
+    pass
+
+
+class FakeApplication:
+    def __init__(self, db_engine):
+        self.logger = MagicMock()
+        self.persistence_mgr = MagicMock()
+        self.persistence_mgr.get_db_engine = MagicMock(return_value=db_engine)
+
+
+@pytest.fixture
+def session_registry(monkeypatch):
+    registry = AgentRunSessionRegistry()
+    monkeypatch.setattr(
+        'langbot.pkg.agent.runner.session_registry._global_registry',
+        registry,
+    )
+    return registry
+
+
+@pytest.fixture
+async def db_engine():
+    engine = create_async_engine('sqlite+aiosqlite:///:memory:')
+    yield engine
+    await engine.dispose()
+
+
+def _handler(db_engine, session_registry):
+    async def fake_disconnect():
+        return True
+
+    fake_app = FakeApplication(db_engine)
+    return RuntimeConnectionHandler(FakeConnection(), fake_disconnect, fake_app)
+
+
+async def _register_session(
+    session_registry,
+    *,
+    run_id='run_1',
+    conversation_id='conv_1',
+    permissions=None,
+):
+    await session_registry.register(
+        run_id=run_id,
+        runner_id='plugin:test/runner/default',
+        query_id=None,
+        plugin_identity='test/runner',
+        resources=make_resources(),
+        conversation_id=conversation_id,
+        permissions=permissions or {},
+    )
+
+
+@pytest.mark.asyncio
+async def test_history_page_requires_manifest_permission(session_registry, db_engine):
+    await _register_session(session_registry, permissions={'history': []})
+    handler = _handler(db_engine, session_registry)
+    history_page = handler.actions[PluginToRuntimeAction.HISTORY_PAGE.value]
+
+    result = await history_page({
+        'run_id': 'run_1',
+        'caller_plugin_identity': 'test/runner',
+    })
+
+    assert result.code != 0
+    assert 'not authorized' in result.message.lower()
+
+
+@pytest.mark.asyncio
+async def test_history_page_rejects_cross_conversation(session_registry, db_engine):
+    await _register_session(session_registry, permissions={'history': ['page']})
+    handler = _handler(db_engine, session_registry)
+    history_page = handler.actions[PluginToRuntimeAction.HISTORY_PAGE.value]
+
+    result = await history_page({
+        'run_id': 'run_1',
+        'conversation_id': 'conv_other',
+        'caller_plugin_identity': 'test/runner',
+    })
+
+    assert result.code != 0
+    assert 'not accessible' in result.message.lower()
+
+
+@pytest.mark.asyncio
+async def test_history_search_rejects_filter_conversation_override(session_registry, db_engine):
+    await _register_session(session_registry, permissions={'history': ['search']})
+    handler = _handler(db_engine, session_registry)
+    history_search = handler.actions[PluginToRuntimeAction.HISTORY_SEARCH.value]
+
+    result = await history_search({
+        'run_id': 'run_1',
+        'query': 'hello',
+        'filters': {'conversation_id': 'conv_other'},
+        'caller_plugin_identity': 'test/runner',
+    })
+
+    assert result.code != 0
+    assert 'not accessible' in result.message.lower()
+
+
+@pytest.mark.asyncio
+async def test_event_page_requires_manifest_permission(session_registry, db_engine):
+    await _register_session(session_registry, permissions={'events': []})
+    handler = _handler(db_engine, session_registry)
+    event_page = handler.actions[PluginToRuntimeAction.EVENT_PAGE.value]
+
+    result = await event_page({
+        'run_id': 'run_1',
+        'caller_plugin_identity': 'test/runner',
+    })
+
+    assert result.code != 0
+    assert 'not authorized' in result.message.lower()
+
+
+@pytest.mark.asyncio
+async def test_event_page_rejects_cross_conversation(session_registry, db_engine):
+    await _register_session(session_registry, permissions={'events': ['page']})
+    handler = _handler(db_engine, session_registry)
+    event_page = handler.actions[PluginToRuntimeAction.EVENT_PAGE.value]
+
+    result = await event_page({
+        'run_id': 'run_1',
+        'conversation_id': 'conv_other',
+        'caller_plugin_identity': 'test/runner',
+    })
+
+    assert result.code != 0
+    assert 'not accessible' in result.message.lower()

tests/unit_tests/agent/test_resource_builder.py

@@ -18,6 +18,7 @@ def make_descriptor(
     *,
     permissions: dict | None = None,
     config_schema: list[dict] | None = None,
+    capabilities: dict | None = None,
 ) -> AgentRunnerDescriptor:
     return AgentRunnerDescriptor(
         id=RUNNER_ID,
@@ -26,6 +27,7 @@ def make_descriptor(
         plugin_author='test',
         plugin_name='runner',
         runner_name='default',
+        capabilities=capabilities or {},
         permissions=permissions or {'models': ['invoke', 'stream']},
         config_schema=config_schema or [],
     )
@@ -99,6 +101,7 @@ async def test_build_models_authorizes_config_declared_llm_and_rerank_models(app
     app.model_mgr.get_model_by_uuid = AsyncMock(side_effect=get_model_by_uuid)
     app.model_mgr.get_rerank_model_by_uuid = AsyncMock(side_effect=get_rerank_model_by_uuid)
     descriptor = make_descriptor(
+        permissions={'models': ['invoke', 'stream', 'rerank']},
         config_schema=[
             {'name': 'model', 'type': 'model-fallback-selector'},
             {'name': 'aux-model', 'type': 'llm-model-selector'},
@@ -145,6 +148,33 @@ async def test_build_models_still_honors_manifest_permissions(app):
     app.model_mgr.get_rerank_model_by_uuid.assert_not_awaited()
 
 
+@pytest.mark.asyncio
+async def test_build_models_authorizes_rerank_only_runner(app):
+    """A rerank-only runner should receive config-selected rerank models."""
+    app.model_mgr.get_model_by_uuid = AsyncMock(return_value=make_model())
+    app.model_mgr.get_rerank_model_by_uuid = AsyncMock(
+        return_value=make_model(model_type='rerank', provider='rerank-provider')
+    )
+    descriptor = make_descriptor(
+        permissions={'models': ['rerank']},
+        config_schema=[
+            {'name': 'model', 'type': 'llm-model-selector'},
+            {'name': 'rerank-model', 'type': 'rerank-model-selector'},
+        ],
+    )
+    query = make_query({
+        'model': 'llm',
+        'rerank-model': 'rerank',
+    })
+
+    resources = await build_resources(app, query, descriptor)
+
+    assert resources['models'] == [
+        {'model_id': 'rerank', 'model_type': 'rerank', 'provider': 'rerank-provider'},
+    ]
+    app.model_mgr.get_model_by_uuid.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_build_models_deduplicates_query_and_config_models(app):
     """A model selected by both preproc and runner config should appear once."""
@@ -197,3 +227,37 @@ async def test_build_tools_authorizes_query_declared_tools(app):
             'description': None,
         },
     ]
+
+
+@pytest.mark.asyncio
+async def test_build_knowledge_bases_unions_config_and_policy_grants(app):
+    descriptor = make_descriptor(
+        capabilities={'knowledge_retrieval': True},
+        permissions={
+            'models': [],
+            'knowledge_bases': ['retrieve'],
+        },
+        config_schema=[
+            {'name': 'knowledge-bases', 'type': 'knowledge-base-multi-selector'},
+        ],
+    )
+    query = make_query(
+        {'knowledge-bases': ['kb_config']},
+        variables={'_knowledge_base_uuids': ['kb_policy']},
+    )
+
+    async def get_kb(kb_uuid):
+        return SimpleNamespace(
+            uuid=kb_uuid,
+            get_name=lambda: f'name-{kb_uuid}',
+            knowledge_base_entity=SimpleNamespace(kb_type='default'),
+        )
+
+    app.rag_mgr.get_knowledge_base_by_uuid = AsyncMock(side_effect=get_kb)
+
+    resources = await build_resources(app, query, descriptor)
+
+    assert resources['knowledge_bases'] == [
+        {'kb_id': 'kb_config', 'kb_name': 'name-kb_config', 'kb_type': 'default'},
+        {'kb_id': 'kb_policy', 'kb_name': 'name-kb_policy', 'kb_type': 'default'},
+    ]

tests/unit_tests/agent/test_state_store.py

@@ -213,6 +213,30 @@ class TestPersistentStateStore:
         snapshot = await persistent_store.build_snapshot_from_event(event, binding, descriptor)
         assert snapshot['conversation']['test_key'] == {'nested': 'value'}
 
+    @pytest.mark.asyncio
+    async def test_state_api_methods_normalize_public_key_aliases(self, persistent_store):
+        scope_key = 'conversation:runner:binding:conv_001'
+
+        success, error = await persistent_store.state_set(
+            scope_key=scope_key,
+            state_key='conversation_id',
+            value='conv_001',
+            runner_id='plugin:test/my-runner/default',
+            binding_identity='binding_001',
+            scope='conversation',
+        )
+
+        assert success is True
+        assert error is None
+        assert await persistent_store.state_get(scope_key, 'external.conversation_id') == 'conv_001'
+        assert await persistent_store.state_get(scope_key, 'conversation_id') == 'conv_001'
+
+        keys, _ = await persistent_store.state_list(scope_key, prefix='conversation_id')
+        assert keys == ['external.conversation_id']
+
+        assert await persistent_store.state_delete(scope_key, 'conversation_id') is True
+        assert await persistent_store.state_get(scope_key, 'external.conversation_id') is None
+
     @pytest.mark.asyncio
     async def test_binding_isolation(self, persistent_store):
         descriptor = make_descriptor()