test_coordinator_integration.py 34 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866
  1. import asyncio
  2. from collections import deque
  3. import pytest
  4. from agent.orchestration.config import OrchestrationConfig
  5. from agent.orchestration.coordinator import TaskConflict, TaskCoordinator
  6. from agent.orchestration.models import (
  7. AgentRole,
  8. ArtifactRef,
  9. AttemptSubmission,
  10. CriterionResult,
  11. DecisionAction,
  12. OperationKind,
  13. OperationStatus,
  14. TaskStatus,
  15. ValidationVerdict,
  16. )
  17. from agent.orchestration.protocols import ValidatorRunResult, WorkerRunResult
  18. from agent.orchestration.store import FileSystemArtifactStore, FileSystemTaskStore, TraceEventSink
  19. from agent.trace.goal_models import GoalTree
  20. from agent.trace.models import Trace
  21. from agent.trace.store import FileSystemTraceStore
  22. from agent.core.runner import AgentRunner
  23. from agent.orchestration.wiring import wire_orchestration
  24. ROOT_TASK_SPEC = {
  25. "objective": "mission",
  26. "acceptance_criteria": [
  27. {"criterion_id": "mission-done", "description": "mission is complete", "hard": True}
  28. ],
  29. }
  30. class FakeExecutor:
  31. def __init__(self, verdicts, submit_worker=True, submit_validator=True, delay=0):
  32. self.verdicts = deque(verdicts)
  33. self.submit_worker = submit_worker
  34. self.submit_validator = submit_validator
  35. self.delay = delay
  36. self.coordinator = None
  37. self.active = 0
  38. self.max_active = 0
  39. self.worker_calls = 0
  40. self.validator_calls = 0
  41. async def run_worker(self, context):
  42. self.worker_calls += 1
  43. self.active += 1
  44. self.max_active = max(self.max_active, self.active)
  45. if self.delay:
  46. await asyncio.sleep(self.delay)
  47. try:
  48. if self.submit_worker:
  49. await self.coordinator.submit_attempt(
  50. {
  51. **context,
  52. "role": AgentRole.WORKER.value,
  53. "trace_id": context["worker_trace_id"],
  54. "tool_call_id": f"submit-{context['attempt_id']}",
  55. },
  56. AttemptSubmission(
  57. summary="done",
  58. artifact_refs=[ArtifactRef(uri=f"memory://{context['attempt_id']}", version="1")],
  59. ),
  60. )
  61. return WorkerRunResult(context["worker_trace_id"], "completed")
  62. finally:
  63. self.active -= 1
  64. async def run_validator(self, context):
  65. self.validator_calls += 1
  66. verdict = self.verdicts.popleft() if self.verdicts else ValidationVerdict.PASSED
  67. if self.submit_validator:
  68. criteria = [
  69. CriterionResult(
  70. criterion_id=item["criterion_id"], verdict=verdict,
  71. reason="checked",
  72. )
  73. for item in context["task_spec"]["acceptance_criteria"]
  74. ]
  75. await self.coordinator.submit_validation(
  76. {
  77. **context,
  78. "role": AgentRole.VALIDATOR.value,
  79. "trace_id": context["validator_trace_id"],
  80. "tool_call_id": f"validate-{context['validation_id']}",
  81. },
  82. verdict,
  83. criteria,
  84. "independent report",
  85. [], [], [], "accept" if verdict == ValidationVerdict.PASSED else "replan",
  86. )
  87. return ValidatorRunResult(context["validator_trace_id"], "completed")
  88. async def make_coordinator(
  89. tmp_path,
  90. executor,
  91. artifact_store=None,
  92. trace_store=None,
  93. task_store=None,
  94. ):
  95. trace_store = trace_store or FileSystemTraceStore(str(tmp_path))
  96. await trace_store.create_trace(Trace(trace_id="root", mode="agent", task="mission", agent_role="planner"))
  97. await trace_store.update_goal_tree("root", GoalTree(mission="mission"))
  98. task_store = task_store or FileSystemTaskStore(str(tmp_path))
  99. coordinator = TaskCoordinator(
  100. task_store,
  101. artifact_store or FileSystemArtifactStore(str(tmp_path)),
  102. trace_store,
  103. OrchestrationConfig(max_parallel_tasks=4),
  104. TraceEventSink(str(tmp_path)),
  105. executor,
  106. )
  107. executor.coordinator = coordinator
  108. await coordinator.ensure_ledger("root", ROOT_TASK_SPEC)
  109. return coordinator, task_store, trace_store
  110. async def create_task(coordinator, objective="task", parent_task_id=None):
  111. result = await coordinator.create_tasks(
  112. "root",
  113. [{
  114. "objective": objective,
  115. "acceptance_criteria": [{"criterion_id": "c1", "description": "must pass", "hard": True}],
  116. }],
  117. parent_task_id=parent_task_id,
  118. )
  119. return result["tasks"][0]["task_id"]
  120. @pytest.mark.asyncio
  121. async def test_passed_validation_requires_planner_accept(tmp_path):
  122. executor = FakeExecutor([ValidationVerdict.PASSED])
  123. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  124. task_id = await create_task(coordinator)
  125. cycle = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  126. assert cycle.task_status == TaskStatus.AWAITING_DECISION
  127. ledger = await store.load("root")
  128. assert ledger.tasks[task_id].status == TaskStatus.AWAITING_DECISION
  129. await coordinator.decide_task(
  130. "root", task_id, cycle.validation_id, DecisionAction.ACCEPT,
  131. {"reason": "all hard criteria passed"}, "decision-1",
  132. )
  133. assert (await store.load("root")).tasks[task_id].status == TaskStatus.COMPLETED
  134. @pytest.mark.asyncio
  135. async def test_failed_or_inconclusive_cannot_be_accepted(tmp_path):
  136. executor = FakeExecutor([ValidationVerdict.FAILED])
  137. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  138. task_id = await create_task(coordinator)
  139. cycle = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  140. with pytest.raises(TaskConflict, match="passed"):
  141. await coordinator.decide_task(
  142. "root", task_id, cycle.validation_id, DecisionAction.ACCEPT,
  143. {"reason": "override"}, "bad-accept",
  144. )
  145. assert (await store.load("root")).tasks[task_id].status == TaskStatus.AWAITING_DECISION
  146. @pytest.mark.asyncio
  147. async def test_repair_is_limited_to_once_per_spec_version(tmp_path):
  148. executor = FakeExecutor([ValidationVerdict.FAILED, ValidationVerdict.FAILED])
  149. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  150. task_id = await create_task(coordinator)
  151. first = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  152. await coordinator.decide_task(
  153. "root", task_id, first.validation_id, DecisionAction.REPAIR,
  154. {"reason": "small local correction"}, "repair-1",
  155. )
  156. second = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  157. ledger = await store.load("root")
  158. attempts = [ledger.attempts[x] for x in ledger.tasks[task_id].attempt_ids]
  159. assert attempts[0].worker_trace_id == attempts[1].worker_trace_id
  160. with pytest.raises(TaskConflict, match="limit"):
  161. await coordinator.decide_task(
  162. "root", task_id, second.validation_id, DecisionAction.REPAIR,
  163. {"reason": "try again"}, "repair-2",
  164. )
  165. @pytest.mark.asyncio
  166. async def test_worker_without_submit_attempt_needs_replan(tmp_path):
  167. executor = FakeExecutor([], submit_worker=False)
  168. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  169. task_id = await create_task(coordinator)
  170. result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  171. ledger = await store.load("root")
  172. assert result.task_status == TaskStatus.NEEDS_REPLAN
  173. assert ledger.tasks[task_id].status == TaskStatus.NEEDS_REPLAN
  174. assert ledger.attempts[result.attempt_id].status.value == "failed"
  175. @pytest.mark.asyncio
  176. async def test_validator_without_submit_is_error_then_revalidates_with_new_trace(tmp_path):
  177. executor = FakeExecutor([ValidationVerdict.PASSED], submit_validator=False)
  178. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  179. task_id = await create_task(coordinator)
  180. first = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  181. ledger = await store.load("root")
  182. first_validation = ledger.validations[first.validation_id]
  183. assert first.task_status == TaskStatus.NEEDS_REPLAN
  184. assert first_validation.status.value == "error"
  185. assert first_validation.verdict is None
  186. executor.submit_validator = True
  187. executor.verdicts.append(ValidationVerdict.PASSED)
  188. second = await coordinator.revalidate_attempt(
  189. "root", task_id, first.attempt_id, "revalidate-1"
  190. )
  191. assert second.task_status == TaskStatus.AWAITING_DECISION
  192. assert second.validation_id != first.validation_id
  193. assert second.validation.validator_trace_id != first_validation.validator_trace_id
  194. @pytest.mark.asyncio
  195. async def test_retry_uses_new_trace_and_revise_invalidates_old_validation(tmp_path):
  196. executor = FakeExecutor([
  197. ValidationVerdict.FAILED,
  198. ValidationVerdict.INCONCLUSIVE,
  199. ValidationVerdict.PASSED,
  200. ])
  201. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  202. task_id = await create_task(coordinator)
  203. first = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  204. await coordinator.decide_task(
  205. "root", task_id, first.validation_id, DecisionAction.RETRY,
  206. {"reason": "use a fresh worker"}, "retry-1",
  207. )
  208. second = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  209. ledger = await store.load("root")
  210. first_attempt = ledger.attempts[first.attempt_id]
  211. second_attempt = ledger.attempts[second.attempt_id]
  212. assert first_attempt.worker_trace_id != second_attempt.worker_trace_id
  213. await coordinator.decide_task(
  214. "root", task_id, second.validation_id, DecisionAction.REVISE,
  215. {
  216. "reason": "clarify criterion",
  217. "objective": "revised task",
  218. "acceptance_criteria": [{"criterion_id": "c1", "description": "revised", "hard": True}],
  219. },
  220. "revise-1",
  221. )
  222. ledger = await store.load("root")
  223. assert ledger.tasks[task_id].current_spec_version == 2
  224. third = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  225. with pytest.raises(TaskConflict, match="obsolete"):
  226. await coordinator.decide_task(
  227. "root", task_id, first.validation_id, DecisionAction.ACCEPT,
  228. {"reason": "stale"}, "stale-accept",
  229. )
  230. await coordinator.decide_task(
  231. "root", task_id, third.validation_id, DecisionAction.ACCEPT,
  232. {"reason": "current validation passed"}, "current-accept",
  233. )
  234. @pytest.mark.asyncio
  235. async def test_block_unblock_cancel_and_supersede(tmp_path):
  236. executor = FakeExecutor([])
  237. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  238. task_id = await create_task(coordinator, "blocked")
  239. await coordinator.decide_task(
  240. "root", task_id, None, DecisionAction.BLOCK,
  241. {"reason": "external dependency"}, "block-1",
  242. )
  243. assert (await store.load("root")).tasks[task_id].blocked_reason == "external dependency"
  244. await coordinator.decide_task(
  245. "root", task_id, None, DecisionAction.UNBLOCK, {}, "unblock-1"
  246. )
  247. assert (await store.load("root")).tasks[task_id].status == TaskStatus.NEEDS_REPLAN
  248. await coordinator.decide_task(
  249. "root", task_id, None, DecisionAction.CANCEL,
  250. {"reason": "no longer needed"}, "cancel-1",
  251. )
  252. assert (await store.load("root")).tasks[task_id].status == TaskStatus.CANCELLED
  253. old_id = await create_task(coordinator, "old")
  254. result = await coordinator.decide_task(
  255. "root", old_id, None, DecisionAction.SUPERSEDE,
  256. {
  257. "reason": "replace spec",
  258. "replacement": {
  259. "objective": "replacement",
  260. "acceptance_criteria": [
  261. {"criterion_id": "replacement-c1", "description": "replacement passes"}
  262. ],
  263. "context_refs": ["replacement-context"],
  264. },
  265. },
  266. "supersede-1",
  267. )
  268. replacement_id = result["payload"]["replacement_task_id"]
  269. ledger = await store.load("root")
  270. assert ledger.tasks[old_id].status == TaskStatus.SUPERSEDED
  271. assert ledger.tasks[replacement_id].status == TaskStatus.PENDING
  272. replacement_spec = ledger.tasks[replacement_id].current_spec
  273. assert replacement_spec.objective == "replacement"
  274. assert replacement_spec.acceptance_criteria[0].criterion_id == "replacement-c1"
  275. assert replacement_spec.context_refs == ("replacement-context",)
  276. @pytest.mark.asyncio
  277. async def test_insert_after_keeps_stable_ids_and_updates_display_order(tmp_path):
  278. executor = FakeExecutor([])
  279. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  280. first = await create_task(coordinator, "first")
  281. third = await create_task(coordinator, "third")
  282. inserted = await coordinator.create_tasks(
  283. "root",
  284. [{
  285. "objective": "second",
  286. "acceptance_criteria": [
  287. {"criterion_id": "second-c1", "description": "second passes"}
  288. ],
  289. "context_refs": ["second-context"],
  290. }],
  291. placement={"after_task_id": first, "focus": True},
  292. idempotency_key="insert",
  293. )
  294. second = inserted["tasks"][0]["task_id"]
  295. ledger = await store.load("root")
  296. assert ledger.tasks[first].display_path == "0.1"
  297. assert ledger.tasks[second].display_path == "0.2"
  298. assert ledger.tasks[third].display_path == "0.3"
  299. assert ledger.tasks[second].current_spec.context_refs == ("second-context",)
  300. assert ledger.focused_task_id == second
  301. assert len({first, second, third}) == 3
  302. repeated = await coordinator.create_tasks(
  303. "root",
  304. [{
  305. "objective": "second",
  306. "acceptance_criteria": [
  307. {"criterion_id": "second-c1", "description": "second passes"}
  308. ],
  309. "context_refs": ["second-context"],
  310. }],
  311. placement={"after_task_id": first, "focus": True},
  312. idempotency_key="insert",
  313. )
  314. assert repeated["tasks"][0]["task_id"] == second
  315. @pytest.mark.asyncio
  316. async def test_split_children_do_not_auto_complete_parent(tmp_path):
  317. executor = FakeExecutor([ValidationVerdict.FAILED, ValidationVerdict.PASSED, ValidationVerdict.PASSED])
  318. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  319. parent_id = await create_task(coordinator, "parent")
  320. first = (await coordinator.dispatch_tasks("root", [parent_id]))[0]
  321. decision = await coordinator.decide_task(
  322. "root", parent_id, first.validation_id, DecisionAction.SPLIT,
  323. {
  324. "reason": "split work",
  325. "tasks": [
  326. {
  327. "objective": "child one",
  328. "acceptance_criteria": [{"criterion_id": "c1", "description": "pass"}],
  329. "context_refs": ["child-one-context"],
  330. },
  331. {
  332. "objective": "child two",
  333. "acceptance_criteria": [{"criterion_id": "c2", "description": "pass"}],
  334. "context_refs": ["child-two-context"],
  335. },
  336. ],
  337. },
  338. "split-1",
  339. )
  340. child_ids = decision["payload"]["child_task_ids"]
  341. ledger = await store.load("root")
  342. assert ledger.tasks[parent_id].child_task_ids == child_ids
  343. assert [ledger.tasks[x].display_path for x in child_ids] == ["0.1.1", "0.1.2"]
  344. assert [ledger.tasks[x].current_spec.context_refs for x in child_ids] == [
  345. ("child-one-context",),
  346. ("child-two-context",),
  347. ]
  348. assert [
  349. ledger.tasks[x].current_spec.acceptance_criteria[0].criterion_id
  350. for x in child_ids
  351. ] == ["c1", "c2"]
  352. cycles = await coordinator.dispatch_tasks("root", child_ids)
  353. for child_id, cycle in zip(child_ids, cycles):
  354. await coordinator.decide_task(
  355. "root", child_id, cycle.validation_id, DecisionAction.ACCEPT,
  356. {"reason": "passed"}, f"accept-{child_id}",
  357. )
  358. ledger = await store.load("root")
  359. assert all(ledger.tasks[x].status == TaskStatus.COMPLETED for x in child_ids)
  360. assert ledger.tasks[parent_id].status == TaskStatus.NEEDS_REPLAN
  361. @pytest.mark.asyncio
  362. async def test_parallel_tasks_are_isolated_and_bounded(tmp_path):
  363. executor = FakeExecutor([ValidationVerdict.PASSED] * 4, delay=0.02)
  364. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  365. task_ids = [await create_task(coordinator, f"task-{i}") for i in range(4)]
  366. results = await coordinator.dispatch_tasks("root", task_ids, idempotency_key="batch")
  367. assert [x.task_id for x in results] == task_ids
  368. assert executor.max_active <= 4
  369. assert all(x.task_status == TaskStatus.AWAITING_DECISION for x in results)
  370. ledger = await store.load("root")
  371. assert len({ledger.attempts[ledger.tasks[x].attempt_ids[-1]].worker_trace_id for x in task_ids}) == 4
  372. repeated = await coordinator.dispatch_tasks("root", task_ids, idempotency_key="batch")
  373. assert [x.attempt_id for x in repeated] == [x.attempt_id for x in results]
  374. @pytest.mark.asyncio
  375. async def test_real_runner_local_executor_creates_independent_terminal_traces(tmp_path):
  376. import json
  377. async def fake_llm(messages, tools, **kwargs):
  378. names = {item["function"]["name"] for item in tools or []}
  379. if "submit_attempt" in names:
  380. arguments = {
  381. "summary": "implemented",
  382. "artifact_refs": [{"uri": "memory://result", "version": "1"}],
  383. "evidence_refs": [],
  384. }
  385. tool_name = "submit_attempt"
  386. elif "submit_validation" in names:
  387. arguments = {
  388. "verdict": "passed",
  389. "criterion_results": [{"criterion_id": "c1", "verdict": "passed", "reason": "verified"}],
  390. "summary": "independent pass",
  391. "evidence_refs": [],
  392. "unverified_claims": [],
  393. "risks": [],
  394. "recommendation": "accept",
  395. }
  396. tool_name = "submit_validation"
  397. else:
  398. raise AssertionError(f"unexpected tools: {names}")
  399. return {
  400. "content": "",
  401. "tool_calls": [{
  402. "id": f"call-{tool_name}",
  403. "type": "function",
  404. "function": {"name": tool_name, "arguments": json.dumps(arguments)},
  405. }],
  406. "finish_reason": "tool_calls",
  407. }
  408. trace_store = FileSystemTraceStore(str(tmp_path))
  409. await trace_store.create_trace(Trace(trace_id="root", mode="agent", task="mission", agent_role="planner"))
  410. await trace_store.update_goal_tree("root", GoalTree(mission="mission"))
  411. runner = AgentRunner(trace_store=trace_store, llm_call=fake_llm)
  412. coordinator = wire_orchestration(
  413. runner,
  414. FileSystemTaskStore(str(tmp_path)),
  415. FileSystemArtifactStore(str(tmp_path)),
  416. )
  417. await coordinator.ensure_ledger("root", ROOT_TASK_SPEC)
  418. task_id = await create_task(coordinator)
  419. cycle = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  420. ledger = await coordinator.task_store.load("root")
  421. worker = await trace_store.get_trace(ledger.attempts[cycle.attempt_id].worker_trace_id)
  422. validator = await trace_store.get_trace(cycle.validation.validator_trace_id)
  423. assert cycle.task_status == TaskStatus.AWAITING_DECISION
  424. assert worker.agent_role == "worker" and worker.result_summary
  425. assert validator.agent_role == "validator" and validator.result_summary
  426. assert worker.trace_id != validator.trace_id
  427. class OneShotGetFailureArtifactStore(FileSystemArtifactStore):
  428. def __init__(self, base_path):
  429. super().__init__(base_path)
  430. self.get_calls = 0
  431. def for_root(self, root_trace_id):
  432. self.root_trace_id = root_trace_id
  433. return self
  434. async def get(self, root_trace_id, snapshot_id=None):
  435. self.get_calls += 1
  436. if self.get_calls == 2:
  437. raise RuntimeError("injected artifact read failure")
  438. return await super().get(root_trace_id, snapshot_id)
  439. class OneShotGoalProjectionFailureStore(FileSystemTraceStore):
  440. def __init__(self, base_path):
  441. super().__init__(base_path)
  442. self.fail_next_goal_update = False
  443. async def update_goal(self, trace_id, goal_id, cascade_completion=True, **updates):
  444. if self.fail_next_goal_update:
  445. self.fail_next_goal_update = False
  446. raise RuntimeError("injected goal projection failure")
  447. return await super().update_goal(
  448. trace_id,
  449. goal_id,
  450. cascade_completion=cascade_completion,
  451. **updates,
  452. )
  453. class OneShotGoalTreeProjectionFailureStore(FileSystemTraceStore):
  454. def __init__(self, base_path):
  455. super().__init__(base_path)
  456. self.fail_next_goal_tree_update = False
  457. async def update_goal_tree(self, trace_id, goal_tree):
  458. if self.fail_next_goal_tree_update:
  459. self.fail_next_goal_tree_update = False
  460. raise RuntimeError("injected goal tree projection failure")
  461. return await super().update_goal_tree(trace_id, goal_tree)
  462. class FailingDispatchCompletionTaskStore(FileSystemTaskStore):
  463. def __init__(self, base_path):
  464. super().__init__(base_path)
  465. self.fail_next_dispatch_completion = True
  466. async def commit(self, ledger, expected_revision, idempotency_key=None, event=None):
  467. if (
  468. self.fail_next_dispatch_completion
  469. and idempotency_key == "root:batch"
  470. ):
  471. self.fail_next_dispatch_completion = False
  472. raise RuntimeError("injected dispatch completion persistence failure")
  473. return await super().commit(
  474. ledger,
  475. expected_revision,
  476. idempotency_key=idempotency_key,
  477. event=event,
  478. )
  479. class OneShotLoadFailureTaskStore(FileSystemTaskStore):
  480. def __init__(self, base_path):
  481. super().__init__(base_path)
  482. self.fail_next_load = False
  483. self._active_operation_loads_before_failure = 1
  484. async def load(self, root_trace_id):
  485. ledger = await super().load(root_trace_id)
  486. if self.fail_next_load and any(
  487. operation.status.value == "running"
  488. for operation in ledger.operations.values()
  489. ):
  490. if self._active_operation_loads_before_failure:
  491. self._active_operation_loads_before_failure -= 1
  492. else:
  493. self.fail_next_load = False
  494. raise RuntimeError("injected reservation load failure")
  495. return ledger
  496. class FailingEventSink:
  497. async def emit(self, root_trace_id, event_type, payload):
  498. raise RuntimeError("injected event sink failure")
  499. @pytest.mark.asyncio
  500. async def test_parallel_unexpected_branch_error_does_not_cancel_siblings(tmp_path):
  501. executor = FakeExecutor([ValidationVerdict.PASSED] * 4, delay=0.01)
  502. artifact_store = OneShotGetFailureArtifactStore(str(tmp_path))
  503. coordinator, store, _ = await make_coordinator(
  504. tmp_path, executor, artifact_store=artifact_store
  505. )
  506. task_ids = [await create_task(coordinator, f"isolated-{index}") for index in range(4)]
  507. results = await coordinator.dispatch_tasks("root", task_ids)
  508. assert [result.task_id for result in results] == task_ids
  509. failures = [result for result in results if result.error]
  510. assert len(failures) == 1
  511. assert "artifact read failure" in failures[0].error
  512. assert failures[0].task_status == TaskStatus.NEEDS_REPLAN
  513. successes = [result for result in results if not result.error]
  514. assert len(successes) == 3
  515. assert all(result.task_status == TaskStatus.AWAITING_DECISION for result in successes)
  516. ledger = await store.load("root")
  517. assert all(
  518. ledger.tasks[task_id].status not in {TaskStatus.RUNNING, TaskStatus.VALIDATING}
  519. for task_id in task_ids
  520. )
  521. @pytest.mark.asyncio
  522. async def test_batch_reservation_conflict_does_not_strand_other_tasks(tmp_path):
  523. executor = FakeExecutor([ValidationVerdict.PASSED, ValidationVerdict.PASSED])
  524. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  525. task_ids = [await create_task(coordinator, f"reserve-{index}") for index in range(3)]
  526. existing = await coordinator._create_attempt(
  527. "root", task_ids[1], "worker", "existing-reservation"
  528. )
  529. results = await coordinator.dispatch_tasks("root", task_ids)
  530. assert [result.task_id for result in results] == task_ids
  531. assert results[0].task_status == TaskStatus.AWAITING_DECISION
  532. assert results[1].attempt_id is None
  533. assert "Dispatch conflict" in results[1].error
  534. assert results[2].task_status == TaskStatus.AWAITING_DECISION
  535. ledger = await store.load("root")
  536. assert ledger.tasks[task_ids[0]].status == TaskStatus.AWAITING_DECISION
  537. assert ledger.tasks[task_ids[2]].status == TaskStatus.AWAITING_DECISION
  538. assert ledger.attempts[existing["attempt_id"]].status.value == "running"
  539. @pytest.mark.asyncio
  540. async def test_goal_projection_failure_does_not_abort_committed_attempt(tmp_path):
  541. executor = FakeExecutor([ValidationVerdict.PASSED])
  542. trace_store = OneShotGoalProjectionFailureStore(str(tmp_path))
  543. coordinator, store, _ = await make_coordinator(
  544. tmp_path,
  545. executor,
  546. trace_store=trace_store,
  547. )
  548. task_id = await create_task(coordinator, "reservation-projection-failure")
  549. trace_store.fail_next_goal_update = True
  550. result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
  551. assert result.task_id == task_id
  552. assert result.task_status == TaskStatus.AWAITING_DECISION
  553. assert result.attempt_id is not None
  554. assert result.error is None
  555. ledger = await store.load("root")
  556. assert ledger.tasks[task_id].status == TaskStatus.AWAITING_DECISION
  557. assert ledger.attempts[result.attempt_id].status.value == "submitted"
  558. @pytest.mark.asyncio
  559. async def test_task_creation_survives_goal_tree_projection_failure_and_reconciles(tmp_path):
  560. executor = FakeExecutor([])
  561. trace_store = OneShotGoalTreeProjectionFailureStore(str(tmp_path))
  562. coordinator, store, _ = await make_coordinator(
  563. tmp_path,
  564. executor,
  565. trace_store=trace_store,
  566. )
  567. trace_store.fail_next_goal_tree_update = True
  568. task_id = await create_task(coordinator, "projection-outage")
  569. ledger = await store.load("root")
  570. assert ledger.tasks[task_id].status == TaskStatus.PENDING
  571. assert ledger.tasks[task_id].goal_id is None
  572. result = await coordinator.reconcile_goal_tree("root")
  573. ledger = await store.load("root")
  574. assert result["reconciled_tasks"] == 2
  575. assert ledger.tasks[task_id].goal_id is not None
  576. @pytest.mark.asyncio
  577. async def test_dispatch_returns_results_when_batch_result_persistence_fails(tmp_path):
  578. executor = FakeExecutor([ValidationVerdict.PASSED, ValidationVerdict.PASSED])
  579. task_store = FailingDispatchCompletionTaskStore(str(tmp_path))
  580. coordinator, store, _ = await make_coordinator(
  581. tmp_path,
  582. executor,
  583. task_store=task_store,
  584. )
  585. task_ids = [await create_task(coordinator, f"batch-{index}") for index in range(2)]
  586. first = await coordinator.dispatch_tasks(
  587. "root",
  588. task_ids,
  589. idempotency_key="batch",
  590. )
  591. assert [result.task_status for result in first] == [
  592. TaskStatus.AWAITING_DECISION,
  593. TaskStatus.AWAITING_DECISION,
  594. ]
  595. assert executor.worker_calls == 2
  596. assert executor.validator_calls == 2
  597. repeated = await coordinator.dispatch_tasks(
  598. "root",
  599. task_ids,
  600. idempotency_key="batch",
  601. )
  602. assert [result.attempt_id for result in repeated] == [
  603. result.attempt_id for result in first
  604. ]
  605. assert executor.worker_calls == 2
  606. assert executor.validator_calls == 2
  607. ledger = await store.load("root")
  608. assert all(len(ledger.tasks[task_id].attempt_ids) == 1 for task_id in task_ids)
  609. assert all(len(ledger.tasks[task_id].validation_ids) == 1 for task_id in task_ids)
  610. @pytest.mark.asyncio
  611. async def test_dispatch_idempotency_key_is_bound_to_tasks_and_presets(tmp_path):
  612. executor = FakeExecutor([ValidationVerdict.PASSED])
  613. coordinator, _, _ = await make_coordinator(tmp_path, executor)
  614. first_task = await create_task(coordinator, "first-binding")
  615. second_task = await create_task(coordinator, "second-binding")
  616. await coordinator.dispatch_tasks(
  617. "root",
  618. [first_task],
  619. worker_presets=["worker"],
  620. idempotency_key="bound-batch",
  621. )
  622. with pytest.raises(TaskConflict, match="different task_ids"):
  623. await coordinator.dispatch_tasks(
  624. "root",
  625. [second_task],
  626. worker_presets=["worker"],
  627. idempotency_key="bound-batch",
  628. )
  629. with pytest.raises(TaskConflict, match="different worker_presets"):
  630. await coordinator.dispatch_tasks(
  631. "root",
  632. [first_task],
  633. worker_presets=["alternate-worker"],
  634. idempotency_key="bound-batch",
  635. )
  636. assert executor.worker_calls == 1
  637. assert executor.validator_calls == 1
  638. @pytest.mark.asyncio
  639. async def test_concurrent_dispatch_replay_does_not_duplicate_agent_execution(tmp_path):
  640. executor = FakeExecutor([ValidationVerdict.PASSED], delay=0.02)
  641. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  642. task_id = await create_task(coordinator, "concurrent-idempotency")
  643. first, replay = await asyncio.gather(
  644. coordinator.dispatch_tasks("root", [task_id], idempotency_key="concurrent"),
  645. coordinator.dispatch_tasks("root", [task_id], idempotency_key="concurrent"),
  646. )
  647. assert executor.worker_calls == 1
  648. assert executor.validator_calls == 1
  649. assert any(
  650. result[0].task_status == TaskStatus.AWAITING_DECISION
  651. for result in (first, replay)
  652. )
  653. assert first[0].attempt_id == replay[0].attempt_id
  654. assert all(result[0].task_status == TaskStatus.AWAITING_DECISION for result in (first, replay))
  655. final = await coordinator.dispatch_tasks(
  656. "root",
  657. [task_id],
  658. idempotency_key="concurrent",
  659. )
  660. ledger = await store.load("root")
  661. assert final[0].task_status == TaskStatus.AWAITING_DECISION
  662. assert len(ledger.tasks[task_id].attempt_ids) == 1
  663. assert len(ledger.tasks[task_id].validation_ids) == 1
  664. @pytest.mark.asyncio
  665. async def test_reservation_load_failure_does_not_corrupt_existing_attempt(tmp_path):
  666. executor = FakeExecutor([ValidationVerdict.PASSED])
  667. task_store = OneShotLoadFailureTaskStore(str(tmp_path))
  668. coordinator, store, _ = await make_coordinator(
  669. tmp_path,
  670. executor,
  671. task_store=task_store,
  672. )
  673. running_task = await create_task(coordinator, "already-running")
  674. existing = await coordinator._create_attempt(
  675. "root", running_task, "worker", "existing-running-attempt"
  676. )
  677. sibling_task = await create_task(coordinator, "unrelated-sibling")
  678. task_store.fail_next_load = True
  679. results = await coordinator.dispatch_tasks(
  680. "root",
  681. [running_task, sibling_task],
  682. )
  683. assert [result.task_id for result in results] == [running_task, sibling_task]
  684. assert results[0].attempt_id is None
  685. assert results[0].task_status == TaskStatus.RUNNING
  686. assert "reservation load failure" in results[0].error
  687. assert results[1].task_status == TaskStatus.AWAITING_DECISION
  688. ledger = await store.load("root")
  689. assert ledger.tasks[running_task].status == TaskStatus.RUNNING
  690. assert ledger.attempts[existing["attempt_id"]].status.value == "running"
  691. @pytest.mark.asyncio
  692. async def test_event_sink_failure_does_not_rollback_committed_ledger(tmp_path):
  693. executor = FakeExecutor([])
  694. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  695. coordinator.event_sink = FailingEventSink()
  696. task_id = await create_task(coordinator, "event-outage")
  697. ledger = await store.load("root")
  698. assert ledger.tasks[task_id].status == TaskStatus.PENDING
  699. @pytest.mark.asyncio
  700. async def test_background_operation_is_durable_and_idempotently_bound(tmp_path):
  701. executor = FakeExecutor([ValidationVerdict.PASSED])
  702. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  703. task_id = await create_task(coordinator, "durable-operation")
  704. operation = await coordinator.start_operation(
  705. "root",
  706. OperationKind.DISPATCH,
  707. task_ids=[task_id],
  708. idempotency_key="operation-key",
  709. )
  710. completed = await coordinator.await_operation("root", operation.operation_id)
  711. replay = await coordinator.start_operation(
  712. "root",
  713. OperationKind.DISPATCH,
  714. task_ids=[task_id],
  715. idempotency_key="operation-key",
  716. )
  717. ledger = await store.load("root")
  718. assert completed.status == OperationStatus.COMPLETED
  719. assert replay.operation_id == completed.operation_id
  720. assert ledger.operations[completed.operation_id].attempt_ids
  721. assert ledger.operations[completed.operation_id].validation_ids
  722. assert executor.worker_calls == executor.validator_calls == 1
  723. other = await create_task(coordinator, "different-operation")
  724. with pytest.raises(TaskConflict, match="different task_ids"):
  725. await coordinator.start_operation(
  726. "root",
  727. OperationKind.DISPATCH,
  728. task_ids=[other],
  729. idempotency_key="operation-key",
  730. )
  731. @pytest.mark.asyncio
  732. async def test_submitted_attempt_advances_to_validation_without_worker_replay(tmp_path):
  733. executor = FakeExecutor([ValidationVerdict.PASSED])
  734. coordinator, store, _ = await make_coordinator(tmp_path, executor)
  735. task_id = await create_task(coordinator, "resume-after-submit")
  736. reserved = await coordinator._create_attempt(
  737. "root", task_id, "worker", "resume-reservation"
  738. )
  739. await coordinator.submit_attempt(
  740. {
  741. "role": AgentRole.WORKER.value,
  742. "root_trace_id": "root",
  743. "task_id": task_id,
  744. "attempt_id": reserved["attempt_id"],
  745. "trace_id": reserved["worker_trace_id"],
  746. "spec_version": 1,
  747. "tool_call_id": "resume-submission",
  748. },
  749. AttemptSubmission(
  750. summary="already executed",
  751. artifact_refs=[ArtifactRef(uri="memory://submitted", version="1")],
  752. ),
  753. )
  754. result = await coordinator.advance_cycle("root", task_id, reserved["attempt_id"])
  755. assert result.task_status == TaskStatus.AWAITING_DECISION
  756. assert executor.worker_calls == 0
  757. assert executor.validator_calls == 1
  758. assert len((await store.load("root")).tasks[task_id].attempt_ids) == 1