|
|
@@ -0,0 +1,282 @@
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import json
|
|
|
+from math import nan
|
|
|
+from typing import Any, Mapping
|
|
|
+
|
|
|
+import pytest
|
|
|
+
|
|
|
+from agent import (
|
|
|
+ AgentRole,
|
|
|
+ FileSystemArtifactStore,
|
|
|
+ FileSystemTaskStore,
|
|
|
+ RunConfig,
|
|
|
+ RoleRunConfigOverrides,
|
|
|
+ RoleRunConfigResolver,
|
|
|
+ wire_orchestration,
|
|
|
+)
|
|
|
+from agent.core.prompts.orchestration import PLANNER_ROLE_CONTRACT
|
|
|
+from agent.orchestration.executor import LocalAgentExecutor
|
|
|
+from agent.orchestration.models import CompletionPolicy, FailureCode
|
|
|
+
|
|
|
+
|
|
|
+class CapturingRunner:
|
|
|
+ trace_store = None
|
|
|
+
|
|
|
+ def __init__(self) -> None:
|
|
|
+ self.configs = []
|
|
|
+ self.messages = []
|
|
|
+
|
|
|
+ async def run_result(self, *, messages, config):
|
|
|
+ self.messages.append(messages)
|
|
|
+ self.configs.append(config)
|
|
|
+ return {"status": "completed", "summary": "done"}
|
|
|
+
|
|
|
+
|
|
|
+class RecordingResolver:
|
|
|
+ def __init__(self, overrides: RoleRunConfigOverrides) -> None:
|
|
|
+ self.overrides = overrides
|
|
|
+ self.calls = []
|
|
|
+
|
|
|
+ async def resolve(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ role: AgentRole,
|
|
|
+ preset: str,
|
|
|
+ context: Mapping[str, Any],
|
|
|
+ ) -> RoleRunConfigOverrides:
|
|
|
+ with pytest.raises(TypeError):
|
|
|
+ context["injected"] = True # type: ignore[index]
|
|
|
+ self.calls.append((role, preset, dict(context)))
|
|
|
+ return self.overrides
|
|
|
+
|
|
|
+
|
|
|
+def worker_context(**updates):
|
|
|
+ context = {
|
|
|
+ "worker_preset": "custom_worker",
|
|
|
+ "worker_trace_id": "worker-trace",
|
|
|
+ "root_trace_id": "root",
|
|
|
+ "task_id": "task",
|
|
|
+ "spec_version": 2,
|
|
|
+ "attempt_id": "attempt",
|
|
|
+ "task_spec": {"objective": "test"},
|
|
|
+ "continue_trace_id": None,
|
|
|
+ }
|
|
|
+ context.update(updates)
|
|
|
+ return context
|
|
|
+
|
|
|
+
|
|
|
+def validator_context(**updates):
|
|
|
+ context = {
|
|
|
+ "validator_preset": "custom_validator",
|
|
|
+ "validator_trace_id": "validator-trace",
|
|
|
+ "root_trace_id": "root",
|
|
|
+ "task_id": "task",
|
|
|
+ "spec_version": 2,
|
|
|
+ "attempt_id": "attempt",
|
|
|
+ "snapshot_id": "snapshot",
|
|
|
+ "validation_id": "validation",
|
|
|
+ "task_spec": {"objective": "test"},
|
|
|
+ "artifact_snapshot": {"artifact_refs": []},
|
|
|
+ }
|
|
|
+ context.update(updates)
|
|
|
+ return context
|
|
|
+
|
|
|
+
|
|
|
+def test_role_run_config_types_are_publicly_exported():
|
|
|
+ assert RoleRunConfigResolver.__name__ == "RoleRunConfigResolver"
|
|
|
+ assert RoleRunConfigOverrides().__dict__ == {
|
|
|
+ "model": None,
|
|
|
+ "temperature": None,
|
|
|
+ "max_iterations": None,
|
|
|
+ }
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_local_executor_without_resolver_preserves_run_config_defaults():
|
|
|
+ runner = CapturingRunner()
|
|
|
+
|
|
|
+ result = await LocalAgentExecutor(runner).run_worker(worker_context())
|
|
|
+
|
|
|
+ assert result.status == "completed"
|
|
|
+ config = runner.configs[0]
|
|
|
+ assert config.model == "gpt-4o"
|
|
|
+ assert config.temperature == 0.3
|
|
|
+ assert config.max_iterations == 200
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_worker_resolver_overrides_only_model_run_fields():
|
|
|
+ runner = CapturingRunner()
|
|
|
+ resolver = RecordingResolver(
|
|
|
+ RoleRunConfigOverrides(
|
|
|
+ model="worker-model",
|
|
|
+ temperature=0.15,
|
|
|
+ max_iterations=17,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+ result = await LocalAgentExecutor(runner, resolver).run_worker(
|
|
|
+ worker_context(untrusted="resolver-can-see-but-runner-must-not-receive")
|
|
|
+ )
|
|
|
+
|
|
|
+ assert result.status == "completed"
|
|
|
+ assert resolver.calls[0][0:2] == (AgentRole.WORKER, "custom_worker")
|
|
|
+ assert resolver.calls[0][2]["untrusted"] == (
|
|
|
+ "resolver-can-see-but-runner-must-not-receive"
|
|
|
+ )
|
|
|
+ config = runner.configs[0]
|
|
|
+ assert (config.model, config.temperature, config.max_iterations) == (
|
|
|
+ "worker-model",
|
|
|
+ 0.15,
|
|
|
+ 17,
|
|
|
+ )
|
|
|
+ assert config.agent_type == "custom_worker"
|
|
|
+ assert config.completion_policy == CompletionPolicy.EXPLICIT_VALIDATION
|
|
|
+ assert config.new_trace_id == "worker-trace"
|
|
|
+ assert config.parent_trace_id == "root"
|
|
|
+ assert config.tools is None and config.tool_groups is None
|
|
|
+ assert config.parallel_tool_execution is False
|
|
|
+ assert config.enable_memory is False
|
|
|
+ assert config.enable_research_flow is False
|
|
|
+ assert "untrusted" not in config.context
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_resolver_context_is_recursively_read_only_and_cannot_mutate_prompt():
|
|
|
+ class NestedMutationResolver:
|
|
|
+ async def resolve(self, *, role, preset, context):
|
|
|
+ del role, preset
|
|
|
+ with pytest.raises(TypeError):
|
|
|
+ context["task_spec"]["objective"] = "mutated"
|
|
|
+ with pytest.raises(AttributeError):
|
|
|
+ context["task_spec"]["refs"].append("mutated")
|
|
|
+ return RoleRunConfigOverrides()
|
|
|
+
|
|
|
+ runner = CapturingRunner()
|
|
|
+ result = await LocalAgentExecutor(runner, NestedMutationResolver()).run_worker(
|
|
|
+ worker_context(task_spec={"objective": "original", "refs": ["safe"]})
|
|
|
+ )
|
|
|
+
|
|
|
+ assert result.status == "completed"
|
|
|
+ prompt = json.loads(runner.messages[0][0]["content"])
|
|
|
+ assert prompt["task_spec"] == {"objective": "original", "refs": ["safe"]}
|
|
|
+ assert runner.configs[0].context["task_id"] == "task"
|
|
|
+
|
|
|
+
|
|
|
+def test_trusted_resolver_values_override_preset_model_policy():
|
|
|
+ from agent import AgentPreset, ToolRegistry
|
|
|
+ from agent.orchestration.policy import DefaultToolPolicy
|
|
|
+
|
|
|
+ config = RunConfig(
|
|
|
+ temperature=0.7,
|
|
|
+ max_iterations=99,
|
|
|
+ completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
|
|
|
+ _role_run_config_override_fields=frozenset({"temperature", "max_iterations"}),
|
|
|
+ )
|
|
|
+ preset = AgentPreset(
|
|
|
+ role=AgentRole.VALIDATOR,
|
|
|
+ allowed_tools=[],
|
|
|
+ temperature=0.0,
|
|
|
+ max_iterations=30,
|
|
|
+ )
|
|
|
+
|
|
|
+ resolved = DefaultToolPolicy().resolve(config, preset, ToolRegistry())
|
|
|
+
|
|
|
+ assert resolved.temperature == 0.7
|
|
|
+ assert resolved.max_iterations == 99
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_validator_resolver_receives_validator_role_and_preset():
|
|
|
+ runner = CapturingRunner()
|
|
|
+ resolver = RecordingResolver(
|
|
|
+ RoleRunConfigOverrides(model="validator-model", temperature=0.0)
|
|
|
+ )
|
|
|
+
|
|
|
+ result = await LocalAgentExecutor(runner, resolver).run_validator(
|
|
|
+ validator_context()
|
|
|
+ )
|
|
|
+
|
|
|
+ assert result.status == "completed"
|
|
|
+ assert resolver.calls[0][0:2] == (AgentRole.VALIDATOR, "custom_validator")
|
|
|
+ config = runner.configs[0]
|
|
|
+ assert config.model == "validator-model"
|
|
|
+ assert config.temperature == 0.0
|
|
|
+ assert config.max_iterations == 200
|
|
|
+ assert config.context["validation_id"] == "validation"
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.parametrize(
|
|
|
+ "kwargs",
|
|
|
+ [
|
|
|
+ {"model": " "},
|
|
|
+ {"model": 1},
|
|
|
+ {"temperature": nan},
|
|
|
+ {"temperature": -0.1},
|
|
|
+ {"temperature": True},
|
|
|
+ {"max_iterations": 0},
|
|
|
+ {"max_iterations": True},
|
|
|
+ ],
|
|
|
+)
|
|
|
+def test_role_run_config_overrides_reject_invalid_values(kwargs):
|
|
|
+ with pytest.raises(ValueError):
|
|
|
+ RoleRunConfigOverrides(**kwargs)
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_invalid_resolver_result_is_normalized_as_executor_failure():
|
|
|
+ class InvalidResolver:
|
|
|
+ async def resolve(self, **kwargs):
|
|
|
+ return {"model": "unsafe"}
|
|
|
+
|
|
|
+ runner = CapturingRunner()
|
|
|
+ result = await LocalAgentExecutor(runner, InvalidResolver()).run_worker(
|
|
|
+ worker_context()
|
|
|
+ )
|
|
|
+
|
|
|
+ assert result.status == "failed"
|
|
|
+ assert "must return RoleRunConfigOverrides" in result.error
|
|
|
+ assert result.execution_stats.failure_code == FailureCode.EXECUTOR_ERROR
|
|
|
+ assert runner.configs == []
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_resolver_exception_is_normalized_as_executor_failure():
|
|
|
+ class BrokenResolver:
|
|
|
+ async def resolve(self, **kwargs):
|
|
|
+ raise RuntimeError("model manifest unavailable")
|
|
|
+
|
|
|
+ runner = CapturingRunner()
|
|
|
+ result = await LocalAgentExecutor(runner, BrokenResolver()).run_validator(
|
|
|
+ validator_context()
|
|
|
+ )
|
|
|
+
|
|
|
+ assert result.status == "failed"
|
|
|
+ assert result.error == "model manifest unavailable"
|
|
|
+ assert result.execution_stats.failure_code == FailureCode.EXECUTOR_ERROR
|
|
|
+ assert runner.configs == []
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_wire_orchestration_injects_role_run_config_resolver(tmp_path):
|
|
|
+ runner = CapturingRunner()
|
|
|
+ runner.task_coordinator = None
|
|
|
+ resolver = RecordingResolver(RoleRunConfigOverrides(model="wired-model"))
|
|
|
+
|
|
|
+ coordinator = wire_orchestration(
|
|
|
+ runner,
|
|
|
+ FileSystemTaskStore(str(tmp_path)),
|
|
|
+ FileSystemArtifactStore(str(tmp_path)),
|
|
|
+ role_run_config_resolver=resolver,
|
|
|
+ )
|
|
|
+ result = await coordinator.executor.run_worker(worker_context())
|
|
|
+
|
|
|
+ assert runner.task_coordinator is coordinator
|
|
|
+ assert result.status == "completed"
|
|
|
+ assert runner.configs[0].model == "wired-model"
|
|
|
+
|
|
|
+
|
|
|
+def test_planner_contract_allows_host_specific_dispatch_tool_name():
|
|
|
+ assert "dispatch_tasks" not in PLANNER_ROLE_CONTRACT
|
|
|
+ assert "task-dispatch tool exposed by your preset" in PLANNER_ROLE_CONTRACT
|
SamLee