mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-22 08:25:21 +08:00
test(auth): exercise knowledge and event action boundaries (#2667)
This commit is contained in:
huanghuoguoguo
@@ -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."""
|
||||
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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})
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user