mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-14 09:46:03 +00:00
fix(agent-runner): harden run lifecycle and protocol stores
This commit is contained in:
@@ -106,6 +106,28 @@ class TestQueryToEventEnvelope:
|
||||
"message_id": "source-message-1",
|
||||
}
|
||||
|
||||
def test_query_to_event_keeps_large_payloads_out_of_event_data(self, mock_query):
|
||||
"""Large or nested platform payloads should not be duplicated into event.data."""
|
||||
source_event = Mock()
|
||||
source_event.type = "platform.message.created"
|
||||
source_event.time = 1700000000
|
||||
source_event.sender = None
|
||||
source_event.model_dump = Mock(return_value={
|
||||
"type": "platform.message.created",
|
||||
"message_id": "source-message-1",
|
||||
"message_chain": [{"type": "Image", "base64": "data:image/png;base64," + ("a" * 1024)}],
|
||||
"raw_text": "x" * 1024,
|
||||
"source_platform_object": {"large": "payload"},
|
||||
})
|
||||
mock_query.message_event = source_event
|
||||
|
||||
event = QueryEntryAdapter.query_to_event(mock_query)
|
||||
|
||||
assert event.data == {
|
||||
"type": "platform.message.created",
|
||||
"message_id": "source-message-1",
|
||||
}
|
||||
|
||||
def test_query_to_event_handles_missing_message_chain(self, mock_query):
|
||||
"""Test delivery context building when Query has no message_chain."""
|
||||
delattr(mock_query, "message_chain")
|
||||
@@ -137,6 +159,29 @@ class TestQueryConfigToAgentConfig:
|
||||
|
||||
assert agent_config.runner_id == "plugin:author/plugin/runner"
|
||||
|
||||
def test_config_to_agent_config_uses_legacy_runner_config_migration(self, mock_query):
|
||||
"""Temporary query adapter must share the normal runner config resolver."""
|
||||
mock_query.pipeline_config = {
|
||||
"ai": {
|
||||
"runner": {"runner": "local-agent"},
|
||||
"local-agent": {
|
||||
"model": "model-primary",
|
||||
"knowledge-base": "kb-001",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
agent_config = QueryEntryAdapter.config_to_agent_config(
|
||||
mock_query,
|
||||
"plugin:langbot/local-agent/default",
|
||||
)
|
||||
|
||||
assert agent_config.runner_config["model"] == {
|
||||
"primary": "model-primary",
|
||||
"fallbacks": [],
|
||||
}
|
||||
assert agent_config.runner_config["knowledge-bases"] == ["kb-001"]
|
||||
|
||||
def test_resolver_projects_agent_scope(self, mock_query):
|
||||
"""Test binding scope projection through the resolver."""
|
||||
event = QueryEntryAdapter.query_to_event(mock_query)
|
||||
|
||||
@@ -2024,6 +2024,75 @@ class TestCallerPluginIdentityValidation:
|
||||
|
||||
await registry.unregister('run_no_caller_identity')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_missing_plugin_identity_denied(self):
|
||||
"""Malformed legacy sessions without plugin_identity fail closed."""
|
||||
from langbot.pkg.agent.runner.session_registry import get_session_registry
|
||||
|
||||
registry = get_session_registry()
|
||||
resources = make_resources(models=[{'model_id': 'model_001'}])
|
||||
session = make_session(
|
||||
run_id='run_missing_session_identity',
|
||||
runner_id='plugin:test/runner/default',
|
||||
plugin_identity='',
|
||||
resources=resources,
|
||||
)
|
||||
async with registry._lock:
|
||||
registry._sessions['run_missing_session_identity'] = session
|
||||
|
||||
from langbot.pkg.plugin.handler import _validate_run_authorization
|
||||
|
||||
mock_ap = MagicMock()
|
||||
mock_ap.logger = MagicMock()
|
||||
|
||||
session, error = await _validate_run_authorization(
|
||||
'run_missing_session_identity',
|
||||
'model',
|
||||
'model_001',
|
||||
mock_ap,
|
||||
caller_plugin_identity='test/runner',
|
||||
)
|
||||
|
||||
assert session is None
|
||||
assert error is not None
|
||||
assert 'no plugin_identity' in error.message
|
||||
|
||||
await registry.unregister('run_missing_session_identity')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pull_api_session_missing_plugin_identity_denied(self):
|
||||
"""Pull API validation also fails closed for missing session identity."""
|
||||
from langbot.pkg.agent.runner.session_registry import get_session_registry
|
||||
|
||||
registry = get_session_registry()
|
||||
session = make_session(
|
||||
run_id='run_missing_pull_identity',
|
||||
runner_id='plugin:test/runner/default',
|
||||
plugin_identity='',
|
||||
available_apis={'history_page': True},
|
||||
)
|
||||
async with registry._lock:
|
||||
registry._sessions['run_missing_pull_identity'] = session
|
||||
|
||||
from langbot.pkg.plugin.handler import _validate_agent_run_session
|
||||
|
||||
mock_ap = MagicMock()
|
||||
mock_ap.logger = MagicMock()
|
||||
|
||||
session, error = await _validate_agent_run_session(
|
||||
'run_missing_pull_identity',
|
||||
'test/runner',
|
||||
mock_ap,
|
||||
'HISTORY_PAGE',
|
||||
'history_page',
|
||||
)
|
||||
|
||||
assert session is None
|
||||
assert error is not None
|
||||
assert 'no plugin_identity' in error.message
|
||||
|
||||
await registry.unregister('run_missing_pull_identity')
|
||||
|
||||
|
||||
class TestBackwardCompatStorageNoRunId:
|
||||
"""Tests for unscoped storage actions without run_id.
|
||||
|
||||
@@ -275,6 +275,37 @@ def test_context_builder_includes_consumable_base64_attachments():
|
||||
assert input_data.attachments[1].name == "hello.txt"
|
||||
|
||||
|
||||
def test_context_builder_deduplicates_message_chain_attachments():
|
||||
query = make_query()
|
||||
query.user_message = None
|
||||
query.message_chain = platform_message.MessageChain(
|
||||
[platform_message.Image(base64="data:image/jpeg;base64,aGVsbG8=")]
|
||||
)
|
||||
|
||||
input_data = QueryEntryAdapter._build_input(query)
|
||||
|
||||
assert [content.type for content in input_data.contents] == ["image_base64"]
|
||||
assert len(input_data.attachments) == 1
|
||||
assert input_data.attachments[0].artifact_type == "image"
|
||||
assert input_data.attachments[0].content == "data:image/jpeg;base64,aGVsbG8="
|
||||
|
||||
|
||||
def test_context_builder_preserves_same_source_duplicate_attachments():
|
||||
query = make_query()
|
||||
query.user_message = provider_message.Message(
|
||||
role="user",
|
||||
content=[
|
||||
provider_message.ContentElement.from_image_base64("data:image/png;base64,aGVsbG8="),
|
||||
provider_message.ContentElement.from_image_base64("data:image/png;base64,aGVsbG8="),
|
||||
],
|
||||
)
|
||||
query.message_chain = platform_message.MessageChain([])
|
||||
|
||||
input_data = QueryEntryAdapter._build_input(query)
|
||||
|
||||
assert [attachment.artifact_type for attachment in input_data.attachments] == ["image", "image"]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
async def clean_agent_state():
|
||||
"""Reset all singleton stores and create a test database engine."""
|
||||
@@ -546,6 +577,29 @@ async def test_orchestrator_unregisters_session_after_runner_failure(clean_agent
|
||||
assert await get_session_registry().get(context["run_id"]) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_orchestrator_unregisters_session_after_event_log_failure(clean_agent_state):
|
||||
"""Journal failures before runner invocation must not leave steerable sessions."""
|
||||
db_engine = clean_agent_state
|
||||
descriptor = make_descriptor()
|
||||
plugin_connector = FakePluginConnector(
|
||||
results=[
|
||||
{
|
||||
"type": "message.completed",
|
||||
"data": {"message": {"role": "assistant", "content": "unused"}},
|
||||
}
|
||||
]
|
||||
)
|
||||
orchestrator = AgentRunOrchestrator(FakeApplication(plugin_connector, db_engine), FakeRegistry(descriptor))
|
||||
orchestrator.journal.write_event_log = AsyncMock(side_effect=RuntimeError("journal unavailable"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="journal unavailable"):
|
||||
[message async for message in orchestrator.run_from_query(make_query())]
|
||||
|
||||
assert plugin_connector.contexts == []
|
||||
assert await get_session_registry().list_active_runs() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_orchestrator_enforces_total_runner_deadline(clean_agent_state):
|
||||
"""Test that orchestrator enforces total runner timeout."""
|
||||
|
||||
@@ -49,6 +49,20 @@ class TestSessionRegistryBasic:
|
||||
assert 'permissions' not in result
|
||||
assert '_authorized_ids' not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_requires_plugin_identity(self):
|
||||
"""Agent run sessions must always have an owning plugin identity."""
|
||||
registry = AgentRunSessionRegistry()
|
||||
|
||||
with pytest.raises(ValueError, match='plugin_identity is required'):
|
||||
await registry.register(
|
||||
run_id='run_missing_identity',
|
||||
runner_id='plugin:test/my-runner/default',
|
||||
query_id=1,
|
||||
plugin_identity='',
|
||||
resources=make_resources(),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_freezes_authorization_snapshot(self):
|
||||
"""Register should freeze authorization data for the run."""
|
||||
|
||||
Reference in New Issue
Block a user