diff --git a/tests/unit_tests/agent/test_chat_handler.py b/tests/unit_tests/agent/test_chat_handler.py index 4d73ba2b9..6eede3f28 100644 --- a/tests/unit_tests/agent/test_chat_handler.py +++ b/tests/unit_tests/agent/test_chat_handler.py @@ -19,7 +19,6 @@ from langbot.pkg.agent.runner.errors import ( RunnerExecutionError, RunnerNotAuthorizedError, ) -from langbot.pkg.agent.runner.config_resolver import RunnerConfigResolver # Define mock classes in dependency order (no forward references needed) @@ -173,36 +172,6 @@ class MockApplication: self.sess_mgr.get_conversation = AsyncMock() -class TestRunnerConfigResolverInChatHandler: - """Tests for RunnerConfigResolver usage in chat handler context.""" - - def test_resolve_runner_id_from_pipeline_config(self): - """Chat handler should use RunnerConfigResolver to resolve runner ID.""" - pipeline_config = { - 'ai': { - 'runner': { - 'id': 'plugin:langbot-team/LocalAgent/default', - }, - }, - } - - runner_id = RunnerConfigResolver.resolve_runner_id(pipeline_config) - assert runner_id == 'plugin:langbot-team/LocalAgent/default' - - def test_old_runner_field_is_not_resolved(self): - """The 4.x path requires ai.runner.id.""" - pipeline_config = { - 'ai': { - 'runner': { - 'runner': 'local-agent', - }, - }, - } - - runner_id = RunnerConfigResolver.resolve_runner_id(pipeline_config) - assert runner_id is None - - class TestErrorHandling: """Tests for orchestrator error handling.""" diff --git a/tests/unit_tests/agent/test_event_log_transcript.py b/tests/unit_tests/agent/test_event_log_transcript.py index eb0ab3b52..28540071c 100644 --- a/tests/unit_tests/agent/test_event_log_transcript.py +++ b/tests/unit_tests/agent/test_event_log_transcript.py @@ -16,7 +16,6 @@ from langbot.pkg.agent.runner.host_models import ( ) from langbot.pkg.agent.runner.event_log_store import EventLogStore from langbot.pkg.agent.runner.transcript_store import TranscriptStore -from langbot.pkg.agent.runner.session_registry import get_session_registry from langbot_plugin.api.entities.builtin.runner.event import ( ActorContext, ) @@ -166,33 +165,6 @@ class TestTranscriptStore: assert transcript_id is not None -class TestHistoryPageAuthorization: - """Test history.page authorization.""" - - @pytest.mark.asyncio - async def test_history_page_validates_conversation_scope(self, mock_db_engine): - """Test history.page only allows access to run's conversation.""" - # This test verifies the authorization logic - # The actual implementation validates conversation_id matches session - session_registry = get_session_registry() - - await session_registry.register( - run_id='run_1', - runner_id='plugin:test/plugin/runner', - query_id=None, - plugin_identity='test/plugin', - resources={'models': [], 'tools': [], 'knowledge_bases': [], 'storage': {'plugin_storage': True}}, - conversation_id='conv_1', - ) - - session = await session_registry.get('run_1') - assert session is not None - assert session['authorization']['conversation_id'] == 'conv_1' - - # Cleanup - await session_registry.unregister('run_1') - - class TestEventLogStoreRealSQLite: """Test EventLogStore with real SQLite database.""" diff --git a/tests/unit_tests/agent/test_handler_auth.py b/tests/unit_tests/agent/test_handler_auth.py index 67d2316da..82be692b7 100644 --- a/tests/unit_tests/agent/test_handler_auth.py +++ b/tests/unit_tests/agent/test_handler_auth.py @@ -1,15 +1,4 @@ -"""Tests for RuntimeConnectionHandler proxy action authorization. - -Tests focus on: -- INVOKE_LLM authorization -- INVOKE_LLM_STREAM authorization -- CALL_TOOL authorization -- RETRIEVE_KNOWLEDGE_BASE authorization - -Authorization paths: -1. Runner calls: has run_id, validates against session_registry -2. Regular plugin calls: no run_id, unscoped plugin action path -""" +"""Tests for shared Runner resource authorization helpers and registry contracts.""" from __future__ import annotations @@ -178,108 +167,6 @@ class MockDisconnectCallback: return True -class TestInvokeLLMAuthorization: - """Tests for INVOKE_LLM authorization.""" - - @pytest.mark.asyncio - async def test_invoke_llm_authorized_with_run_id(self): - """INVOKE_LLM: authorized when model in session.resources.""" - # Setup registry with session - registry = AgentRunSessionRegistry() - resources = make_resources(models=[{'model_id': 'model_001'}]) - - await registry.register( - run_id='run_authorized', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - # Verify authorization logic directly - session = await registry.get('run_authorized') - assert session is not None - assert registry.is_resource_allowed(session, 'model', 'model_001') is True - - # Cleanup - await registry.unregister('run_authorized') - - @pytest.mark.asyncio - async def test_invoke_llm_unauthorized_with_run_id(self): - """INVOKE_LLM: unauthorized when model not in session.resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(models=[{'model_id': 'model_001'}]) - - await registry.register( - run_id='run_unauthorized', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - # Test authorization logic directly - session = await registry.get('run_unauthorized') - assert session is not None - # model_002 is not in resources - assert registry.is_resource_allowed(session, 'model', 'model_002') is False - - await registry.unregister('run_unauthorized') - - @pytest.mark.asyncio - async def test_invoke_llm_session_not_found(self): - """INVOKE_LLM: session not found should return error.""" - registry = AgentRunSessionRegistry() - - # No session registered for this run_id - session = await registry.get('run_nonexistent') - assert session is None - - -class TestInvokeLLMStreamAuthorization: - """Tests for INVOKE_LLM_STREAM authorization.""" - - @pytest.mark.asyncio - async def test_invoke_llm_stream_authorized_with_run_id(self): - """INVOKE_LLM_STREAM: authorized when model in session.resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(models=[{'model_id': 'model_001'}]) - - await registry.register( - run_id='run_stream_authorized', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - session = await registry.get('run_stream_authorized') - assert session is not None - assert registry.is_resource_allowed(session, 'model', 'model_001') is True - - await registry.unregister('run_stream_authorized') - - @pytest.mark.asyncio - async def test_invoke_llm_stream_unauthorized_with_run_id(self): - """INVOKE_LLM_STREAM: unauthorized when model not in session.resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(models=[{'model_id': 'model_001'}]) - - await registry.register( - run_id='run_stream_unauthorized', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - session = await registry.get('run_stream_unauthorized') - assert session is not None - assert registry.is_resource_allowed(session, 'model', 'model_002') is False - - await registry.unregister('run_stream_unauthorized') - - @pytest.mark.asyncio async def test_tool_manager_get_tool_detail_returns_uniform_schema(): """ToolManager.get_tool_detail returns a uniform host-level tool detail shape. @@ -315,436 +202,6 @@ async def test_tool_manager_get_tool_detail_returns_uniform_schema(): assert await mgr.get_tool_detail('workspace-test', 'missing') is None -class TestCallToolAuthorization: - """Tests for CALL_TOOL authorization.""" - - @pytest.mark.asyncio - async def test_call_tool_authorized_with_run_id(self): - """CALL_TOOL: authorized when tool in session.resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(tools=[{'tool_name': 'web_search'}]) - - await registry.register( - run_id='run_tool_authorized', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - session = await registry.get('run_tool_authorized') - assert session is not None - assert registry.is_resource_allowed(session, 'tool', 'web_search') is True - - await registry.unregister('run_tool_authorized') - - @pytest.mark.asyncio - async def test_call_tool_unauthorized_with_run_id(self): - """CALL_TOOL: unauthorized when tool not in session.resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(tools=[{'tool_name': 'web_search'}]) - - await registry.register( - run_id='run_tool_unauthorized', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - session = await registry.get('run_tool_unauthorized') - assert session is not None - assert registry.is_resource_allowed(session, 'tool', 'image_gen') is False - - await registry.unregister('run_tool_unauthorized') - - -class TestRetrieveKnowledgeBaseAuthorization: - """Tests for RETRIEVE_KNOWLEDGE_BASE authorization.""" - - @pytest.mark.asyncio - async def test_retrieve_knowledge_base_authorized_with_run_id(self): - """RETRIEVE_KNOWLEDGE_BASE: authorized when kb in session.resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(knowledge_bases=[{'kb_id': 'kb_001'}]) - - await registry.register( - run_id='run_kb_authorized', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - session = await registry.get('run_kb_authorized') - assert session is not None - assert registry.is_resource_allowed(session, 'knowledge_base', 'kb_001') is True - - await registry.unregister('run_kb_authorized') - - @pytest.mark.asyncio - async def test_retrieve_knowledge_base_unauthorized_with_run_id(self): - """RETRIEVE_KNOWLEDGE_BASE: unauthorized when kb not in session.resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(knowledge_bases=[{'kb_id': 'kb_001'}]) - - await registry.register( - run_id='run_kb_unauthorized', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - session = await registry.get('run_kb_unauthorized') - assert session is not None - assert registry.is_resource_allowed(session, 'knowledge_base', 'kb_999') is False - - await registry.unregister('run_kb_unauthorized') - - @pytest.mark.asyncio - async def test_retrieve_knowledge_base_no_run_id_pipeline_check(self): - """RETRIEVE_KNOWLEDGE_BASE: no run_id checks pipeline config.""" - # When no run_id, the handler checks against pipeline's configured KBs - # This is the unscoped path for regular plugin calls - - from langbot.pkg.agent.runner.config_resolver import RunnerConfigResolver - - # Simulate pipeline config - pipeline_config = { - 'ai': { - 'runner': { - 'id': 'plugin:test/runner/default', - }, - 'runner_config': { - 'plugin:test/runner/default': { - 'knowledge-bases': ['kb_001', 'kb_002'], - }, - }, - }, - } - - runner_id = RunnerConfigResolver.resolve_runner_id(pipeline_config) - assert runner_id == 'plugin:test/runner/default' - - runner_config = RunnerConfigResolver.resolve_runner_config(pipeline_config, runner_id) - allowed_kbs = runner_config.get('knowledge-bases', []) - assert 'kb_001' in allowed_kbs - assert 'kb_999' not in allowed_kbs - - -class TestAuthorizationPathDifferentiation: - """Tests that verify Runner vs regular plugin call differentiation.""" - - @pytest.mark.asyncio - async def test_runner_path_with_run_id(self): - """Runner calls provide run_id and use session_registry.""" - registry = AgentRunSessionRegistry() - - # Runner call has run_id - run_id = 'run_runner_123' - - # Register session with resources - await registry.register( - run_id=run_id, - runner_id='plugin:test/agent/default', - query_id=1, - plugin_identity='test/agent', - resources=make_resources( - models=[{'model_id': 'model_xyz'}], - tools=[{'tool_name': 'agent_tool'}], - knowledge_bases=[{'kb_id': 'kb_agent'}], - ), - ) - - session = await registry.get(run_id) - assert session is not None - - # Authorization checks - assert registry.is_resource_allowed(session, 'model', 'model_xyz') is True - assert registry.is_resource_allowed(session, 'model', 'other_model') is False - assert registry.is_resource_allowed(session, 'tool', 'agent_tool') is True - assert registry.is_resource_allowed(session, 'tool', 'other_tool') is False - assert registry.is_resource_allowed(session, 'knowledge_base', 'kb_agent') is True - assert registry.is_resource_allowed(session, 'knowledge_base', 'kb_other') is False - - await registry.unregister(run_id) - - -class TestRETRIEVEKNOWLEDGEBASEBugFix: - """Tests for the RETRIEVE_KNOWLEDGE_BASE bug fix in handler.py. - - Bug: Previously, the handler directly accessed pipeline_config['ai']['local-agent'] - without first resolving the runner_id, causing issues when non-local-agent runners - were used. - - Fix: Now uses RunnerConfigResolver.resolve_runner_id first, then resolve_runner_config. - """ - - def test_retrieve_kb_fix_local_runner(self): - """Fix should work for local-agent runner.""" - from langbot.pkg.agent.runner.config_resolver import RunnerConfigResolver - - pipeline_config = { - 'ai': { - 'runner': { - 'id': 'plugin:langbot-team/LocalAgent/default', - }, - 'runner_config': { - 'plugin:langbot-team/LocalAgent/default': { - 'knowledge-bases': ['kb_001'], - }, - }, - }, - } - - runner_id = RunnerConfigResolver.resolve_runner_id(pipeline_config) - runner_config = RunnerConfigResolver.resolve_runner_config(pipeline_config, runner_id) - allowed_kbs = runner_config.get('knowledge-bases', []) - - assert 'kb_001' in allowed_kbs - - def test_retrieve_kb_fix_other_runner(self): - """Fix should work for non-local-agent runners.""" - from langbot.pkg.agent.runner.config_resolver import RunnerConfigResolver - - pipeline_config = { - 'ai': { - 'runner': { - 'id': 'plugin:custom/my-agent/default', - }, - 'runner_config': { - 'plugin:custom/my-agent/default': { - 'knowledge-bases': ['kb_custom'], - }, - }, - }, - } - - runner_id = RunnerConfigResolver.resolve_runner_id(pipeline_config) - runner_config = RunnerConfigResolver.resolve_runner_config(pipeline_config, runner_id) - allowed_kbs = runner_config.get('knowledge-bases', []) - - assert 'kb_custom' in allowed_kbs - - def test_retrieve_kb_does_not_read_old_runner_format(self): - """The 4.x authorization path ignores legacy Pipeline runner fields.""" - from langbot.pkg.agent.runner.config_resolver import RunnerConfigResolver - - pipeline_config = { - 'ai': { - 'runner': { - 'runner': 'local-agent', - }, - 'local-agent': { - 'knowledge-bases': ['kb_legacy'], - }, - }, - } - - runner_id = RunnerConfigResolver.resolve_runner_id(pipeline_config) - runner_config = RunnerConfigResolver.resolve_runner_config(pipeline_config, runner_id) - assert runner_id is None - assert runner_config == {} - - -class TestHandlerActionAuthorization: - """Tests for real handler action-level authorization. - - These tests simulate RuntimeConnectionHandler action handlers - to verify actual authorization behavior at the action level. - """ - - @pytest.mark.asyncio - async def test_invoke_llm_handler_authorized_path(self): - """INVOKE_LLM handler: authorized when model in resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(models=[{'model_id': 'model_001'}]) - - await registry.register( - run_id='run_invoke_llm_auth', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - # Simulate handler authorization logic - run_id = 'run_invoke_llm_auth' - llm_model_uuid = 'model_001' - - session_registry = registry - session = await session_registry.get(run_id) - assert session is not None - - # Authorization check (same as handler.py line 352) - is_allowed = session_registry.is_resource_allowed(session, 'model', llm_model_uuid) - assert is_allowed is True - - await registry.unregister(run_id) - - @pytest.mark.asyncio - async def test_invoke_llm_handler_unauthorized_path(self): - """INVOKE_LLM handler: unauthorized when model not in resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(models=[{'model_id': 'model_001'}]) - - await registry.register( - run_id='run_invoke_llm_unauth', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - run_id = 'run_invoke_llm_unauth' - llm_model_uuid = 'model_999' # Not in resources - - session_registry = registry - session = await session_registry.get(run_id) - assert session is not None - - # Authorization check (same as handler.py line 352) - is_allowed = session_registry.is_resource_allowed(session, 'model', llm_model_uuid) - assert is_allowed is False - - # Should return error response (handler.py line 357) - expected_error = f'Model {llm_model_uuid} is not authorized for this agent run' - assert 'not authorized' in expected_error - - await registry.unregister(run_id) - - @pytest.mark.asyncio - async def test_invoke_llm_handler_session_not_found(self): - """INVOKE_LLM handler: session not found returns error.""" - registry = AgentRunSessionRegistry() - - # No session registered - run_id = 'run_nonexistent' - session = await registry.get(run_id) - assert session is None - - # Handler should return error (handler.py line 348) - expected_error = f'Run session {run_id} not found or expired' - assert 'not found' in expected_error - - @pytest.mark.asyncio - async def test_call_tool_handler_authorized_path(self): - """CALL_TOOL handler: authorized when tool in resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(tools=[{'tool_name': 'web_search'}]) - - await registry.register( - run_id='run_call_tool_auth', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - run_id = 'run_call_tool_auth' - tool_name = 'web_search' - - session_registry = registry - session = await session_registry.get(run_id) - assert session is not None - - # Authorization check (handler.py line 475) - is_allowed = session_registry.is_resource_allowed(session, 'tool', tool_name) - assert is_allowed is True - - await registry.unregister(run_id) - - @pytest.mark.asyncio - async def test_call_tool_handler_unauthorized_path(self): - """CALL_TOOL handler: unauthorized when tool not in resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(tools=[{'tool_name': 'web_search'}]) - - await registry.register( - run_id='run_call_tool_unauth', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - run_id = 'run_call_tool_unauth' - tool_name = 'image_gen' # Not in resources - - session_registry = registry - session = await session_registry.get(run_id) - assert session is not None - - # Authorization check - is_allowed = session_registry.is_resource_allowed(session, 'tool', tool_name) - assert is_allowed is False - - # Should return error (handler.py line 480) - expected_error = f'Tool {tool_name} is not authorized for this agent run' - assert 'not authorized' in expected_error - - await registry.unregister(run_id) - - @pytest.mark.asyncio - async def test_retrieve_knowledge_base_handler_authorized_path(self): - """RETRIEVE_KNOWLEDGE_BASE handler: authorized when kb in resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(knowledge_bases=[{'kb_id': 'kb_001'}]) - - await registry.register( - run_id='run_kb_auth', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - run_id = 'run_kb_auth' - kb_id = 'kb_001' - - session_registry = registry - session = await session_registry.get(run_id) - assert session is not None - - # Authorization check (handler.py line 889) - is_allowed = session_registry.is_resource_allowed(session, 'knowledge_base', kb_id) - assert is_allowed is True - - await registry.unregister(run_id) - - @pytest.mark.asyncio - async def test_retrieve_knowledge_base_handler_unauthorized_path(self): - """RETRIEVE_KNOWLEDGE_BASE handler: unauthorized when kb not in resources.""" - registry = AgentRunSessionRegistry() - resources = make_resources(knowledge_bases=[{'kb_id': 'kb_001'}]) - - await registry.register( - run_id='run_kb_unauth', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - run_id = 'run_kb_unauth' - kb_id = 'kb_999' # Not in resources - - session_registry = registry - session = await session_registry.get(run_id) - assert session is not None - - # Authorization check - is_allowed = session_registry.is_resource_allowed(session, 'knowledge_base', kb_id) - assert is_allowed is False - - # Should return error (handler.py line 894) - expected_error = f'Knowledge base {kb_id} is not authorized for this agent run' - assert 'not authorized' in expected_error - - await registry.unregister(run_id) - - class TestSDKRunnerAPIProxyFieldConsistency: """Tests for SDK RunnerAPIProxy field name consistency with Host handler. @@ -766,42 +223,6 @@ class TestSDKRunnerAPIProxyFieldConsistency: assert action != unrestricted_action -class TestNoRunIdBackwardCompatPath: - """Tests for unscoped plugin action path when no run_id is provided. - - Regular plugins (non-Runner) don't have run_id and should - have unrestricted access to certain APIs. - """ - - @pytest.mark.asyncio - async def test_retrieve_knowledge_base_no_run_id_pipeline_check(self): - """RETRIEVE_KNOWLEDGE_BASE: no run_id uses pipeline config check.""" - from langbot.pkg.agent.runner.config_resolver import RunnerConfigResolver - - # When no run_id, handler.py lines 897-914 check pipeline config - pipeline_config = { - 'ai': { - 'runner': { - 'id': 'plugin:test/runner/default', - }, - 'runner_config': { - 'plugin:test/runner/default': { - 'knowledge-bases': ['kb_001', 'kb_002'], - }, - }, - }, - } - - runner_id = RunnerConfigResolver.resolve_runner_id(pipeline_config) - runner_config = RunnerConfigResolver.resolve_runner_config(pipeline_config, runner_id) - allowed_kb_uuids = runner_config.get('knowledge-bases', []) - - # kb_001 should be allowed - assert 'kb_001' in allowed_kb_uuids - # kb_999 should NOT be allowed - assert 'kb_999' not in allowed_kb_uuids - - class TestSessionExpiryAndCleanup: """Tests for session expiry and cleanup scenarios.""" @@ -960,70 +381,6 @@ class TestResourceTypeValidation: assert registry.is_resource_allowed(session, 'unknown_type', 'any_id') is False -class TestBypassPrevention: - """Tests to ensure RunnerAPIProxy cannot bypass authorization.""" - - @pytest.mark.asyncio - async def test_cannot_bypass_via_unrestricted_retrieve_knowledge(self): - """Cannot bypass KB authorization via unrestricted RETRIEVE_KNOWLEDGE action.""" - # RunnerAPIProxy uses RETRIEVE_KNOWLEDGE_BASE (with run_id) - # RETRIEVE_KNOWLEDGE is unrestricted and separate - # Runner should NOT use RETRIEVE_KNOWLEDGE to bypass authorization - - registry = AgentRunSessionRegistry() - resources = make_resources(knowledge_bases=[{'kb_id': 'kb_001'}]) - - await registry.register( - run_id='run_bypass_test', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - session = await registry.get('run_bypass_test') - - # kb_002 is not authorized - assert registry.is_resource_allowed(session, 'knowledge_base', 'kb_002') is False - - # If Runner tried to use RETRIEVE_KNOWLEDGE (unrestricted), - # it would bypass authorization - but RunnerAPIProxy correctly uses - # RETRIE_KNOWLEDGE_BASE which requires authorization - - from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction - - # Verify SDK uses correct action - assert PluginToRuntimeAction.RETRIEVE_KNOWLEDGE_BASE.value == 'retrieve_knowledge_base' - - await registry.unregister('run_bypass_test') - - @pytest.mark.asyncio - async def test_cannot_bypass_via_missing_run_id_in_session(self): - """Cannot bypass by using run_id that doesn't exist in registry.""" - registry = AgentRunSessionRegistry() - resources = make_resources(models=[{'model_id': 'model_001'}]) - - await registry.register( - run_id='run_valid', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - # Try to use a run_id that doesn't exist - fake_run_id = 'run_fake' - session = await registry.get(fake_run_id) - assert session is None - - # Handler should return error for non-existent run_id - # (handler.py line 348, 466, 881) - expected_error = f'Run session {fake_run_id} not found or expired' - assert 'not found' in expected_error - - await registry.unregister('run_valid') - - class TestValidateRunAuthorizationHelper: """Tests for _validate_run_authorization helper function. @@ -1307,114 +664,6 @@ class TestStorageResourcePermissionHelper: assert registry.is_resource_allowed(session, 'storage', 'workspace') is False -class TestRealActionHandlerSimulation: - """Tests that simulate real RuntimeConnectionHandler action registration and execution. - - These tests attempt to verify the actual handler behavior without full integration. - Uses global session registry to match _validate_run_authorization behavior. - """ - - @pytest.mark.asyncio - async def test_action_handler_invoke_llm_flow(self): - """Simulate INVOKE_LLM action handler authorization flow.""" - # Use global session registry - from langbot.pkg.agent.runner.session_registry import get_session_registry - - registry = get_session_registry() - resources = make_resources(models=[{'model_id': 'model_001'}]) - - await registry.register( - run_id='run_invoke_llm_flow_sim', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - # Simulate handler logic - from langbot.pkg.plugin.handler import _validate_run_authorization - - mock_ap = MagicMock() - mock_ap.logger = MagicMock() - - # Step 1: Validate authorization - session, error = await _validate_run_authorization( - 'run_invoke_llm_flow_sim', - 'model', - 'model_001', - mock_ap, - caller_plugin_identity='test/runner', - ) - - # Should pass authorization - assert session is not None - assert error is None - - # Step 2: Handler would invoke LLM (not tested here, would need mock model) - - await registry.unregister('run_invoke_llm_flow_sim') - - @pytest.mark.asyncio - async def test_action_handler_rejects_unauthorized_model(self): - """Simulate INVOKE_LLM handler rejecting unauthorized model.""" - # Use global session registry - from langbot.pkg.agent.runner.session_registry import get_session_registry - - registry = get_session_registry() - resources = make_resources(models=[{'model_id': 'model_001'}]) - - await registry.register( - run_id='run_reject_model_sim', - runner_id='plugin:test/runner/default', - query_id=1, - plugin_identity='test/runner', - resources=resources, - ) - - from langbot.pkg.plugin.handler import _validate_run_authorization - - mock_ap = MagicMock() - mock_ap.logger = MagicMock() - mock_ap.logger.warning = MagicMock() - - # Try to access unauthorized model - session, error = await _validate_run_authorization( - 'run_reject_model_sim', - 'model', - 'model_999', - mock_ap, - caller_plugin_identity='test/runner', - ) - - # Should reject - assert session is None - assert error is not None - assert 'not authorized' in error.message.lower() - assert mock_ap.logger.warning.called - - await registry.unregister('run_reject_model_sim') - - @pytest.mark.asyncio - async def test_action_handler_session_not_found_flow(self): - """Simulate handler behavior when session not found.""" - from langbot.pkg.plugin.handler import _validate_run_authorization - - mock_ap = MagicMock() - mock_ap.logger = MagicMock() - mock_ap.logger.warning = MagicMock() - - # Try to validate with non-existent run_id - session, error = await _validate_run_authorization( - 'run_nonexistent_session_flow', 'model', 'model_001', mock_ap - ) - - # Should return error - assert session is None - assert error is not None - assert 'not found' in error.message.lower() - assert mock_ap.logger.warning.called - - class TestStoragePermissionValidation: """Tests for Host-side storage permission validation via _validate_run_authorization. diff --git a/tests/unit_tests/agent/test_history_event_api_auth.py b/tests/unit_tests/agent/test_history_event_api_auth.py index 732ce4eac..dc18ee2ea 100644 --- a/tests/unit_tests/agent/test_history_event_api_auth.py +++ b/tests/unit_tests/agent/test_history_event_api_auth.py @@ -213,6 +213,57 @@ async def test_event_get_returns_sdk_record_projection(session_registry, db_engi assert result.data['cursor'] == str(result.data['seq']) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ('denial', 'message'), + [('missing_run', 'run_id is required'), ('expired_run', 'not found or expired'), ('capability', 'not authorized')], +) +async def test_event_get_requires_active_run_capability(session_registry, db_engine, denial, message): + await _register_session(session_registry, available_apis={'event_get': denial != 'capability'}) + await EventLogStore(db_engine).append_event( + event_id='evt_denied', + event_type='message.received', + source='platform', + conversation_id='conv_1', + run_id='run_1', + input_summary='must not be returned', + ) + data = {'event_id': 'evt_denied'} + if denial != 'missing_run': + data['run_id'] = 'expired' if denial == 'expired_run' else 'run_1' + result = await _handler(db_engine, session_registry).actions[PluginToRuntimeAction.EVENT_GET.value](data) + assert result.code != 0 + assert message in result.message + assert not result.data + + +@pytest.mark.asyncio +@pytest.mark.parametrize('mismatch', [None, 'conversation_id', 'bot_id', 'workspace_id', 'thread_id']) +async def test_event_get_enforces_run_scope(session_registry, db_engine, mismatch): + scope = {'conversation_id': 'conv_1', 'bot_id': 'bot_1', 'workspace_id': 'workspace-test', 'thread_id': 'thread_1'} + await _register_session(session_registry, **scope, available_apis={'event_get': True}) + event_scope = {**scope, **({mismatch: 'other'} if mismatch else {})} + await EventLogStore(db_engine).append_event( + event_id='evt_other_run', + event_type='message.received', + source='platform', + run_id='other_run', + input_summary='scoped document', + **event_scope, + ) + result = await _handler(db_engine, session_registry).actions[PluginToRuntimeAction.EVENT_GET.value]( + {'run_id': 'run_1', 'event_id': 'evt_other_run'} + ) + if mismatch is None: + assert result.code == 0 + assert result.data['event_id'] == 'evt_other_run' + assert result.data['input_summary'] == 'scoped document' + else: + assert result.code != 0 + assert 'not accessible' in result.message + assert not result.data + + @pytest.mark.asyncio async def test_event_page_returns_sdk_page_projection(session_registry, db_engine): await _register_session(session_registry, available_apis={'event_page': True}) diff --git a/tests/unit_tests/plugin/test_handler_actions.py b/tests/unit_tests/plugin/test_handler_actions.py index bec6c9e19..7f0c10ee5 100644 --- a/tests/unit_tests/plugin/test_handler_actions.py +++ b/tests/unit_tests/plugin/test_handler_actions.py @@ -1867,3 +1867,185 @@ async def test_embedding_api_rejects_model_outside_runner_grants(): assert response.code != 0 assert 'not authorized' in response.message app.model_mgr.get_embedding_model_by_uuid.assert_not_awaited() + + +@pytest.fixture +def knowledge_host(monkeypatch): + from langbot.pkg.agent.runner import session_registry + from langbot.pkg.entity.persistence.rag import KnowledgeBase + from langbot.pkg.rag.knowledge.kbmgr import RAGManager, RuntimeKnowledgeBase + + registry = session_registry.AgentRunSessionRegistry() + monkeypatch.setattr(session_registry, '_global_registry', registry) + app = SimpleNamespace( + logger=Mock(), + persistence_mgr=SimpleNamespace(), + plugin_connector=SimpleNamespace( + require_workspace_context=AsyncMock(side_effect=lambda context: context), + call_rag_retrieve=AsyncMock( + return_value={ + 'results': [ + { + 'id': 'document-a', + 'content': [{'type': 'text', 'text': 'selected document'}], + 'metadata': {}, + 'distance': 0.1, + } + ] + } + ), + ), + ) + runtime = make_handler(app) + app.rag_mgr = RAGManager(app) + for kb_id, workspace_uuid in [('kb-a', 'workspace-a'), ('kb-unselected', 'workspace-a'), ('kb-b', 'workspace-b')]: + context = ExecutionContext(instance_uuid='instance-a', workspace_uuid=workspace_uuid, placement_generation=1) + entity = KnowledgeBase( + uuid=kb_id, + workspace_uuid=workspace_uuid, + name=kb_id, + knowledge_engine_plugin_id='author/engine', + collection_id=kb_id, + retrieval_settings={}, + ) + app.rag_mgr._cache_runtime(RuntimeKnowledgeBase(app, entity, context)) + app.rag_mgr.get_knowledge_base_by_uuid = AsyncMock(wraps=app.rag_mgr.get_knowledge_base_by_uuid) + return SimpleNamespace(app=app, runtime=runtime, registry=registry) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'action', [PluginToRuntimeAction.RETRIEVE_KNOWLEDGE, PluginToRuntimeAction.RETRIEVE_KNOWLEDGE_BASE] +) +@pytest.mark.parametrize('denial', [None, 'unselected', 'operation', 'expired', 'plugin']) +async def test_knowledge_actions_enforce_run_grants_before_resource_access(knowledge_host, action, denial): + host = knowledge_host + await host.registry.register( + run_id='run-a', + runner_id='plugin:test-author/test-plugin/runner', + query_id=None, + plugin_identity='other/plugin' if denial == 'plugin' else 'test-author/test-plugin', + workspace_id='workspace-a', + resources=make_agent_resources( + knowledge_bases=[{'kb_id': 'kb-a', 'operations': ['list'] if denial == 'operation' else ['retrieve']}] + ), + ) + response = await host.runtime.actions[action.value]( + { + 'run_id': 'expired' if denial == 'expired' else 'run-a', + 'kb_id': 'kb-unselected' if denial == 'unselected' else 'kb-a', + 'query_text': 'hello', + } + ) + if denial is None: + assert response.code == 0 + assert [item['id'] for item in response.data['results']] == ['document-a'] + host.app.plugin_connector.call_rag_retrieve.assert_awaited_once() + else: + assert response.code != 0 + assert ( + response.message + == { + 'unselected': 'knowledge_base kb-unselected is not authorized for operation retrieve for this agent run', + 'operation': 'knowledge_base kb-a is not authorized for operation retrieve for this agent run', + 'expired': 'Run session expired not found or expired', + 'plugin': 'Plugin identity mismatch: caller test-author/test-plugin is not authorized for run_id run-a', + }[denial] + ) + assert not response.data + host.app.rag_mgr.get_knowledge_base_by_uuid.assert_not_awaited() + host.app.plugin_connector.call_rag_retrieve.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'action', [PluginToRuntimeAction.RETRIEVE_KNOWLEDGE, PluginToRuntimeAction.RETRIEVE_KNOWLEDGE_BASE] +) +async def test_knowledge_actions_keep_foreign_runs_in_trusted_workspace(knowledge_host, action): + host = knowledge_host + await host.registry.register( + run_id='run-b', + runner_id='plugin:test-author/test-plugin/runner', + query_id=None, + plugin_identity='test-author/test-plugin', + workspace_id='workspace-b', + resources=make_agent_resources(knowledge_bases=[{'kb_id': 'kb-b'}]), + ) + response = await host.runtime.actions[action.value]( + {'run_id': 'run-b', 'kb_id': 'kb-b', 'query_text': 'hello', 'workspace_uuid': 'workspace-b'} + ) + assert response.code != 0 + assert response.message == 'Knowledge base kb-b not found' + assert not response.data + host.app.rag_mgr.get_knowledge_base_by_uuid.assert_awaited_once_with(TEST_EXECUTION_CONTEXT, 'kb-b') + host.app.plugin_connector.call_rag_retrieve.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('selected', [True, False]) +@pytest.mark.parametrize('scope_source', ['preprocessed', 'custom_runner']) +async def test_pipeline_knowledge_action_enforces_query_scope(knowledge_host, selected, scope_source): + from langbot.pkg.agent.runner.descriptor import RunnerDescriptor + + host = knowledge_host + runner_id = 'plugin:custom/runner/default' + query_context = ExecutionContext( + instance_uuid='instance-a', workspace_uuid='workspace-a', placement_generation=1, bot_uuid='bot-a' + ) + query = SimpleNamespace( + _execution_context=query_context, + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=1, + variables={'_knowledge_base_uuids': ['kb-a']} if scope_source == 'preprocessed' else {}, + pipeline_config={ + 'ai': { + 'runner': {'id': runner_id}, + 'runner_config': { + runner_id: {'custom-kb-selector': ['kb-unselected' if scope_source == 'preprocessed' else 'kb-a']} + }, + } + }, + session=SimpleNamespace(launcher_type=SimpleNamespace(value='person'), launcher_id='conversation-a'), + bot_uuid='bot-a', + sender_id='user-a', + ) + host.app.runner_registry = SimpleNamespace( + get=AsyncMock( + return_value=RunnerDescriptor( + usages=['agent'], + id=runner_id, + source='plugin', + label={'en_US': 'Custom Runner'}, + plugin_author='custom', + plugin_name='runner', + runner_name='default', + config_schema=[{'name': 'custom-kb-selector', 'type': 'knowledge-base-multi-selector', 'default': []}], + capabilities={'knowledge_retrieval': True}, + permissions={'knowledge_bases': ['list', 'retrieve']}, + ) + ) + ) + host.app.query_pool = SimpleNamespace(get_query=AsyncMock(return_value=query)) + response = await host.runtime.actions[PluginToRuntimeAction.RETRIEVE_KNOWLEDGE_BASE.value]( + { + 'query_uuid': 'query-a', + 'query_id': 1, + 'kb_id': 'kb-a' if selected else 'kb-unselected', + 'query_text': 'hello', + } + ) + host.app.query_pool.get_query.assert_awaited_once_with('workspace-a', 'query-a') + if scope_source == 'custom_runner': + host.app.runner_registry.get.assert_awaited_once_with(query_context, runner_id, None) + else: + host.app.runner_registry.get.assert_not_awaited() + if selected: + assert response.code == 0 + assert [item['id'] for item in response.data['results']] == ['document-a'] + host.app.plugin_connector.call_rag_retrieve.assert_awaited_once() + else: + assert response.code != 0 + assert 'not configured for this pipeline' in response.message + host.app.rag_mgr.get_knowledge_base_by_uuid.assert_not_awaited() + host.app.plugin_connector.call_rag_retrieve.assert_not_awaited()