330 lines
13 KiB
Python
330 lines
13 KiB
Python
import argparse
|
|
import asyncio
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import stat
|
|
import subprocess
|
|
import sys
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from tempfile import TemporaryDirectory
|
|
from unittest import mock
|
|
|
|
|
|
SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "dispatch.py"
|
|
SPEC = importlib.util.spec_from_file_location("agent_task_dispatch_test", SCRIPT)
|
|
dispatch = importlib.util.module_from_spec(SPEC)
|
|
assert SPEC.loader is not None
|
|
sys.modules[SPEC.name] = dispatch
|
|
SPEC.loader.exec_module(dispatch)
|
|
|
|
|
|
def catalog_value(command: str = "/bin/true") -> dict:
|
|
targets = {
|
|
"primary": {
|
|
"agent": "runner-primary",
|
|
"model": "model-primary",
|
|
"execution_class": "local_model",
|
|
"selfcheck_required": True,
|
|
"runtime": {
|
|
"command": [command, "--workspace", "{workspace}", "--model", "{model}", "{prompt}"],
|
|
"resume_command": [command, "--resume", "{resume_session}", "{prompt}"],
|
|
"environment": {"TARGET_ID": "{target_id}"},
|
|
"output_format": "jsonl",
|
|
"native_session_monitor": True,
|
|
"session_path": "sessions/{session_id}.jsonl",
|
|
},
|
|
},
|
|
"alternate": {
|
|
"agent": "runner-alternate",
|
|
"model": "model-alternate",
|
|
"execution_class": "cloud_model",
|
|
"runtime": {"command": [command, "{prompt}"]},
|
|
},
|
|
}
|
|
routes = {"worker": {}, "review": {}}
|
|
for stage in routes:
|
|
for lane in ("local", "cloud"):
|
|
for grade in range(1, 11):
|
|
routes[stage][f"{lane}-G{grade:02d}"] = {
|
|
"candidates": ["primary", "alternate"],
|
|
"rule_id": f"{stage}-{lane}-{grade:02d}",
|
|
"reason_codes": ["injected-route"],
|
|
}
|
|
return {"schema_version": "1.0", "targets": targets, "routes": routes}
|
|
|
|
|
|
def write_catalog(root: Path, value: dict | None = None) -> Path:
|
|
path = root / "execution-catalog.json"
|
|
path.write_text(json.dumps(value or catalog_value()), encoding="utf-8")
|
|
return path
|
|
|
|
|
|
def write_plan(root: Path, *, task_name: str = "group/01_task") -> Path:
|
|
directory = root / "agent-task" / task_name
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
path = directory / "PLAN-cloud-G05.md"
|
|
path.write_text(
|
|
f"<!-- task={task_name} plan=0 tag=API -->\n\n"
|
|
"# Plan\n\n## Modified Files Summary\n\n"
|
|
"| File | Action |\n|---|---|\n| `src/item.txt` | modify |\n",
|
|
encoding="utf-8",
|
|
)
|
|
return path
|
|
|
|
|
|
def task_from_plan(root: Path, plan: Path) -> dispatch.Task:
|
|
directory = plan.parent
|
|
return dispatch.Task(
|
|
name="group/01_task",
|
|
directory=directory,
|
|
plan=plan,
|
|
review=None,
|
|
user_review=None,
|
|
recovery=False,
|
|
index=1,
|
|
write_set={"src/item.txt"},
|
|
write_set_known=True,
|
|
plan_hash=dispatch.sha256_file(plan),
|
|
)
|
|
|
|
|
|
class RuntimeCatalogDispatcherTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.previous_catalog = dispatch.EXECUTION_CATALOG_PATH
|
|
|
|
def tearDown(self):
|
|
dispatch.EXECUTION_CATALOG_PATH = self.previous_catalog
|
|
|
|
def test_agent_spec_is_loaded_from_persisted_catalog_evidence(self):
|
|
with TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
catalog = write_catalog(root)
|
|
plan = write_plan(root)
|
|
dispatch.EXECUTION_CATALOG_PATH = catalog
|
|
selector = dispatch._selector_module()
|
|
decision = selector.select_execution_target(plan, catalog_path=catalog)
|
|
spec = dispatch.agent_spec_from_decision(decision)
|
|
self.assertEqual(spec.target_id, "primary")
|
|
self.assertEqual(spec.cli, "runner-primary")
|
|
self.assertEqual(spec.model, "model-primary")
|
|
self.assertTrue(spec.native_resume)
|
|
self.assertEqual(spec.runtime["command"][0], "/bin/true")
|
|
|
|
def test_agent_spec_rejects_catalog_change_after_selection(self):
|
|
with TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
catalog = write_catalog(root)
|
|
plan = write_plan(root)
|
|
selector = dispatch._selector_module()
|
|
decision = selector.select_execution_target(plan, catalog_path=catalog)
|
|
changed = catalog_value()
|
|
changed["targets"]["primary"]["model"] = "changed"
|
|
catalog.write_text(json.dumps(changed), encoding="utf-8")
|
|
with self.assertRaisesRegex(dispatch.ExecutionDecisionError, "변경"):
|
|
dispatch.agent_spec_from_decision(decision)
|
|
|
|
def test_command_is_expanded_only_from_runtime_template(self):
|
|
spec = dispatch.AgentSpec(
|
|
"opaque-agent",
|
|
"opaque-model",
|
|
"opaque-agent/opaque-model",
|
|
target_id="opaque-id",
|
|
runtime={
|
|
"command": ["runner", "{workspace}", "{model}", "{session_id}", "{attempt_dir}", "{prompt}"],
|
|
"resume_command": ["runner", "resume", "{resume_session}", "{prompt}"],
|
|
},
|
|
)
|
|
command = dispatch.build_command(
|
|
spec,
|
|
"do work",
|
|
Path("/workspace"),
|
|
"session-1",
|
|
Path("/attempt"),
|
|
)
|
|
resumed = dispatch.build_command(
|
|
spec,
|
|
"continue",
|
|
Path("/workspace"),
|
|
"session-1",
|
|
Path("/attempt"),
|
|
native_resume_session=Path("/attempt/session.jsonl"),
|
|
)
|
|
self.assertEqual(command, ["runner", "/workspace", "opaque-model", "session-1", "/attempt", "do work"])
|
|
self.assertEqual(resumed, ["runner", "resume", "/attempt/session.jsonl", "continue"])
|
|
|
|
def test_preflight_checks_executable_and_optional_probe(self):
|
|
with TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
value = catalog_value("/bin/true")
|
|
value["targets"]["primary"]["runtime"]["preflight_command"] = ["/bin/true", "--check"]
|
|
catalog = write_catalog(root, value)
|
|
dispatch.preflight_execution_catalog(catalog)
|
|
|
|
def test_preflight_rejects_missing_command(self):
|
|
with TemporaryDirectory() as tmp:
|
|
catalog = write_catalog(Path(tmp), catalog_value("definitely-missing-command"))
|
|
with self.assertRaisesRegex(dispatch.ExecutionDecisionError, "command not found"):
|
|
dispatch.preflight_execution_catalog(catalog)
|
|
|
|
def test_persisted_decision_failover_uses_next_runtime_target(self):
|
|
with TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
catalog = write_catalog(root)
|
|
plan = write_plan(root)
|
|
task = task_from_plan(root, plan)
|
|
dispatch.EXECUTION_CATALOG_PATH = catalog
|
|
with mock.patch.dict(os.environ, {"XDG_STATE_HOME": str(root / "state")}):
|
|
store = dispatch.StateStore(root)
|
|
try:
|
|
initial, first_spec = dispatch.persisted_execution_decision(store, task, stage="worker")
|
|
failed, second_spec = dispatch.persisted_execution_decision(
|
|
store,
|
|
task,
|
|
stage="worker",
|
|
transition="failover",
|
|
failure_class="provider-quota",
|
|
)
|
|
finally:
|
|
store.close()
|
|
self.assertEqual(initial["selected"]["target_id"], "primary")
|
|
self.assertEqual(first_spec.model, "model-primary")
|
|
self.assertEqual(failed["selected"]["target_id"], "alternate")
|
|
self.assertEqual(second_spec.model, "model-alternate")
|
|
self.assertNotIn("quota", failed)
|
|
|
|
def test_retry_blocked_marks_failover_without_quota_state(self):
|
|
with TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
catalog = write_catalog(root)
|
|
plan = write_plan(root)
|
|
task = task_from_plan(root, plan)
|
|
dispatch.EXECUTION_CATALOG_PATH = catalog
|
|
with mock.patch.dict(os.environ, {"XDG_STATE_HOME": str(root / "state")}):
|
|
store = dispatch.StateStore(root)
|
|
try:
|
|
decision, _ = dispatch.persisted_execution_decision(store, task, stage="worker")
|
|
state = store.task_state(task)
|
|
state.update(
|
|
blocked="runtime failure",
|
|
blocker_evidence={
|
|
"role": "worker",
|
|
"failure_class": "provider-quota",
|
|
"locator": "/tmp/locator.json",
|
|
"selected": decision["selected"],
|
|
"work_unit_id": decision["work_unit_id"],
|
|
},
|
|
)
|
|
store.save()
|
|
store.mark_retry_failover("group")
|
|
state = store.task_state(task)
|
|
finally:
|
|
store.close()
|
|
self.assertTrue(state["retry_failover_pending"])
|
|
self.assertNotIn("quota_snapshot", state)
|
|
self.assertNotIn("retry_quota_refresh_pending", state)
|
|
|
|
def test_runtime_error_classifier_keeps_provider_quota(self):
|
|
failure, evidence = dispatch.classify_failure_with_evidence(
|
|
"HTTP 429 resource exhausted: quota reached"
|
|
)
|
|
self.assertEqual(failure, "provider-quota")
|
|
self.assertIsNotNone(evidence)
|
|
|
|
def test_generic_json_terminal_diagnostic_has_no_agent_branch(self):
|
|
diagnostic = dispatch.terminal_diagnostic(
|
|
"opaque-agent",
|
|
"stdout",
|
|
json.dumps({"type": "turn.failed", "error": {"code": 429}}),
|
|
)
|
|
self.assertIn("429", diagnostic or "")
|
|
self.assertIsNone(
|
|
dispatch.terminal_diagnostic(
|
|
"opaque-agent",
|
|
"stdout",
|
|
json.dumps({"type": "message", "text": "quota design notes"}),
|
|
)
|
|
)
|
|
|
|
def test_catalog_source_is_in_runtime_audit_evidence(self):
|
|
with TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
catalog = write_catalog(root)
|
|
plan = write_plan(root)
|
|
selector = dispatch._selector_module()
|
|
decision = selector.select_execution_target(plan, catalog_path=catalog)
|
|
evidence = dispatch.selector_runtime_evidence(decision)
|
|
self.assertEqual(evidence["catalog"]["source"], str(catalog.resolve()))
|
|
self.assertNotIn("quota", evidence)
|
|
|
|
|
|
class GenericDispatcherContractTests(unittest.TestCase):
|
|
def test_parallel_limit_contract(self):
|
|
self.assertEqual(dispatch.validated_max_parallel(0), 0)
|
|
self.assertEqual(dispatch.validated_max_parallel(3), 3)
|
|
with self.assertRaises(ValueError):
|
|
dispatch.validated_max_parallel(-1)
|
|
|
|
def test_modified_files_summary_is_canonicalized(self):
|
|
with TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
plan = write_plan(root)
|
|
write_set, diagnostics = dispatch.inspect_write_set(plan, root)
|
|
self.assertEqual(diagnostics, [])
|
|
self.assertEqual(write_set, {str((root / "src/item.txt").resolve())})
|
|
|
|
def test_outside_workspace_claim_is_rejected(self):
|
|
with TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
plan = write_plan(root)
|
|
plan.write_text(
|
|
plan.read_text(encoding="utf-8").replace("`src/item.txt`", "`../outside.txt`"),
|
|
encoding="utf-8",
|
|
)
|
|
_, diagnostics = dispatch.inspect_write_set(plan, root)
|
|
self.assertTrue(any("outside" in item.lower() or "workspace" in item.lower() for item in diagnostics))
|
|
|
|
def test_selector_evidence_uses_agent_model_fields(self):
|
|
decision = {
|
|
"work_unit_id": "group/01::plan-0::tag-API",
|
|
"selected": {"target_id": "a", "agent": "runner", "model": "model"},
|
|
"candidates": [
|
|
{"candidate_rank": 1, "target_id": "a", "agent": "runner", "model": "model"}
|
|
],
|
|
"decision": {"rule_id": "rule", "policy_priority": 1, "reason_codes": []},
|
|
"transition": {"trigger": "initial"},
|
|
}
|
|
lines = dispatch.selector_evidence_lines(decision)
|
|
self.assertIn("candidates=#1:runner/model", lines)
|
|
self.assertFalse(any("quota" in line for line in lines))
|
|
|
|
def test_validate_plan_mode_does_not_require_catalog(self):
|
|
with TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
plan = write_plan(root)
|
|
completed = subprocess.run(
|
|
[sys.executable, str(SCRIPT), "--workspace", str(root), "--validate-plan", str(plan)],
|
|
capture_output=True,
|
|
text=True,
|
|
env={key: value for key, value in os.environ.items() if key != "AGENT_TASK_EXECUTION_CATALOG"},
|
|
check=False,
|
|
)
|
|
self.assertEqual(completed.returncode, 0, completed.stderr)
|
|
|
|
def test_dry_run_requires_catalog(self):
|
|
with TemporaryDirectory() as tmp:
|
|
completed = subprocess.run(
|
|
[sys.executable, str(SCRIPT), "--workspace", tmp, "--dry-run"],
|
|
capture_output=True,
|
|
text=True,
|
|
env={key: value for key, value in os.environ.items() if key != "AGENT_TASK_EXECUTION_CATALOG"},
|
|
check=False,
|
|
)
|
|
self.assertEqual(completed.returncode, 2)
|
|
self.assertIn("missing_execution_catalog", completed.stderr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|