fix(agent-runner): harden run lifecycle and protocol stores

This commit is contained in:
huanghuoguoguo
2026-06-13 21:22:13 +08:00
parent 735a0011b0
commit 1153433693
16 changed files with 450 additions and 103 deletions

View File

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

View File

@@ -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.

View File

@@ -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."""

View File

@@ -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."""