|
|
@@ -27,6 +27,7 @@ from cyber_agent.application.candidate_service import (
|
|
|
CandidateService,
|
|
|
CandidateStateError,
|
|
|
)
|
|
|
+from cyber_agent.core.runner import AgentRunner, RunConfig
|
|
|
from cyber_agent.core.artifacts import (
|
|
|
ArtifactRef,
|
|
|
ValidationMaterial,
|
|
|
@@ -38,6 +39,8 @@ from cyber_agent.core.validation import (
|
|
|
ValidationCheck,
|
|
|
ValidationResult,
|
|
|
)
|
|
|
+from cyber_agent.tools import get_tool_registry
|
|
|
+from cyber_agent.tools.errors import RecoverableToolExecutionError
|
|
|
from cyber_agent.tools.registry import ToolRegistry
|
|
|
from cyber_agent.trace.models import Message as TraceMessage, Trace
|
|
|
from cyber_agent.trace.store import (
|
|
|
@@ -319,7 +322,10 @@ class CandidateServiceTest(unittest.IsolatedAsyncioTestCase):
|
|
|
"replace_candidate_ledger",
|
|
|
side_effect=fail_completed_publish,
|
|
|
):
|
|
|
- with self.assertRaisesRegex(OSError, "finalize crash"):
|
|
|
+ with self.assertRaisesRegex(
|
|
|
+ RecoverableToolExecutionError,
|
|
|
+ "original command ID",
|
|
|
+ ):
|
|
|
await self.service.manage(
|
|
|
"root-a",
|
|
|
operation="create",
|
|
|
@@ -934,6 +940,111 @@ class CandidateServiceTest(unittest.IsolatedAsyncioTestCase):
|
|
|
self.assertEqual(MAX_CANDIDATE_OPERATIONS_PER_ROOT, len(completed.operations))
|
|
|
self.assertEqual("discarded", completed.current_state(candidate))
|
|
|
|
|
|
+ async def test_review_batch_reserves_action_slot_across_crash_recovery(self):
|
|
|
+ await self._create_child("child-a")
|
|
|
+ candidate = await self.service.manage(
|
|
|
+ "child-a",
|
|
|
+ operation="create",
|
|
|
+ content={"text": "reserved review"},
|
|
|
+ parent_refs=[],
|
|
|
+ effective_at_sequence=4,
|
|
|
+ )
|
|
|
+ await self._record_validation("child-a", candidate)
|
|
|
+ ledger = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ target = MAX_CANDIDATE_OPERATIONS_PER_ROOT - 2
|
|
|
+ fillers = [
|
|
|
+ CandidateOperationRecord(
|
|
|
+ operation_id=f"review-reservation-filler-{index}",
|
|
|
+ operation="review_batch",
|
|
|
+ input_hash="2" * 64,
|
|
|
+ status="completed",
|
|
|
+ )
|
|
|
+ for index in range(target - len(ledger.operations))
|
|
|
+ ]
|
|
|
+ near_limit = CandidateLedger.model_validate({
|
|
|
+ **ledger.model_dump(mode="json"),
|
|
|
+ "operations": [
|
|
|
+ *[item.model_dump(mode="json") for item in ledger.operations],
|
|
|
+ *[item.model_dump(mode="json") for item in fillers],
|
|
|
+ ],
|
|
|
+ })
|
|
|
+ await self.store.replace_candidate_ledger(
|
|
|
+ "root-a",
|
|
|
+ near_limit.model_dump(mode="json"),
|
|
|
+ )
|
|
|
+ action = CandidateReviewAction(
|
|
|
+ action="adopt",
|
|
|
+ candidate_ref=candidate,
|
|
|
+ reason="reserve the complete review batch",
|
|
|
+ )
|
|
|
+ original_replace = self.store.replace_candidate_ledger
|
|
|
+ writes = 0
|
|
|
+
|
|
|
+ async def crash_after_batch_reservation(root_trace_id, raw_ledger):
|
|
|
+ nonlocal writes
|
|
|
+ writes += 1
|
|
|
+ if writes == 2:
|
|
|
+ raise OSError("crash after review reservation")
|
|
|
+ await original_replace(root_trace_id, raw_ledger)
|
|
|
+
|
|
|
+ with patch.object(
|
|
|
+ self.store,
|
|
|
+ "replace_candidate_ledger",
|
|
|
+ side_effect=crash_after_batch_reservation,
|
|
|
+ ):
|
|
|
+ with self.assertRaisesRegex(OSError, "review reservation"):
|
|
|
+ await self.service.apply_review_actions(
|
|
|
+ "root-a",
|
|
|
+ "child-a",
|
|
|
+ report_refs=[candidate],
|
|
|
+ actions=[action],
|
|
|
+ effective_at_sequence=10,
|
|
|
+ command_id="capacity-reserved-review",
|
|
|
+ )
|
|
|
+ reserved = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual(MAX_CANDIDATE_OPERATIONS_PER_ROOT, len(reserved.operations))
|
|
|
+ self.assertEqual(2, sum(
|
|
|
+ item.status == "pending" for item in reserved.operations
|
|
|
+ ))
|
|
|
+ repository_version_calls = self.repository.put_calls
|
|
|
+ with self.assertRaisesRegex(
|
|
|
+ CandidateStateError,
|
|
|
+ "operation limit exceeded before execution",
|
|
|
+ ):
|
|
|
+ await self.service.manage(
|
|
|
+ "child-a",
|
|
|
+ operation="create",
|
|
|
+ content={"text": "must not steal review capacity"},
|
|
|
+ parent_refs=[],
|
|
|
+ effective_at_sequence=11,
|
|
|
+ command_id="sibling-after-review-reservation",
|
|
|
+ )
|
|
|
+ self.assertEqual(repository_version_calls, self.repository.put_calls)
|
|
|
+
|
|
|
+ records = await self.service.apply_review_actions(
|
|
|
+ "root-a",
|
|
|
+ "child-a",
|
|
|
+ report_refs=[candidate],
|
|
|
+ actions=[action],
|
|
|
+ effective_at_sequence=999,
|
|
|
+ command_id="capacity-reserved-review",
|
|
|
+ )
|
|
|
+ self.assertEqual("adopted", records[0].state)
|
|
|
+ self.assertEqual(1, self.repository.adoption_calls)
|
|
|
+ completed = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual(MAX_CANDIDATE_OPERATIONS_PER_ROOT, len(completed.operations))
|
|
|
+ self.assertTrue(all(
|
|
|
+ item.status == "completed"
|
|
|
+ for item in completed.operations
|
|
|
+ if item.command_id == "capacity-reserved-review"
|
|
|
+ ))
|
|
|
+
|
|
|
async def test_candidate_capacity_rejects_before_repository_or_intent_write(self):
|
|
|
await self._create_child("child-a")
|
|
|
candidate = await self.service.manage(
|
|
|
@@ -1026,7 +1137,10 @@ class CandidateServiceTest(unittest.IsolatedAsyncioTestCase):
|
|
|
"replace_candidate_ledger",
|
|
|
side_effect=crash_after_repository,
|
|
|
):
|
|
|
- with self.assertRaisesRegex(OSError, "candidate publish"):
|
|
|
+ with self.assertRaisesRegex(
|
|
|
+ RecoverableToolExecutionError,
|
|
|
+ "original command ID",
|
|
|
+ ):
|
|
|
await self.service.manage(
|
|
|
"child-a",
|
|
|
operation="create",
|
|
|
@@ -1076,6 +1190,120 @@ class CandidateServiceTest(unittest.IsolatedAsyncioTestCase):
|
|
|
self.assertEqual(MAX_CANDIDATES_PER_ROOT, len(completed.candidates))
|
|
|
self.assertEqual(recovered, completed.candidate(recovered))
|
|
|
|
|
|
+ async def test_runner_replays_pending_candidate_with_original_tool_call_id(self):
|
|
|
+ await self._create_child("child-a")
|
|
|
+ registry = get_tool_registry().clone(["manage_candidate"])
|
|
|
+ runner = AgentRunner(
|
|
|
+ trace_store=self.store,
|
|
|
+ tool_registry=registry,
|
|
|
+ llm_call=lambda **_kwargs: None,
|
|
|
+ application_binding=self.binding,
|
|
|
+ candidate_service=self.service,
|
|
|
+ )
|
|
|
+ config = RunConfig(
|
|
|
+ tools=["manage_candidate"],
|
|
|
+ tool_groups=[],
|
|
|
+ auto_execute_tools=True,
|
|
|
+ )
|
|
|
+ request = {
|
|
|
+ "request": {
|
|
|
+ "operation": "create",
|
|
|
+ "content": {"text": "recover with the original command"},
|
|
|
+ "parent_refs": [],
|
|
|
+ },
|
|
|
+ }
|
|
|
+ assistant = TraceMessage.create(
|
|
|
+ trace_id="child-a",
|
|
|
+ role="assistant",
|
|
|
+ sequence=4,
|
|
|
+ content={
|
|
|
+ "text": "",
|
|
|
+ "tool_calls": [{
|
|
|
+ "id": "recoverable-candidate-call",
|
|
|
+ "type": "function",
|
|
|
+ "function": {
|
|
|
+ "name": "manage_candidate",
|
|
|
+ "arguments": json.dumps(request),
|
|
|
+ },
|
|
|
+ }],
|
|
|
+ },
|
|
|
+ )
|
|
|
+ await self.store.add_message(assistant)
|
|
|
+ await self.store.update_trace(
|
|
|
+ "child-a",
|
|
|
+ head_sequence=4,
|
|
|
+ last_sequence=4,
|
|
|
+ )
|
|
|
+ trace = await self.store.get_trace("child-a")
|
|
|
+ original_replace = self.store.replace_candidate_ledger
|
|
|
+ writes = 0
|
|
|
+
|
|
|
+ async def fail_first_publish(root_trace_id, raw_ledger):
|
|
|
+ nonlocal writes
|
|
|
+ writes += 1
|
|
|
+ if writes == 2:
|
|
|
+ raise OSError("candidate publish interrupted")
|
|
|
+ await original_replace(root_trace_id, raw_ledger)
|
|
|
+
|
|
|
+ with patch.object(
|
|
|
+ self.store,
|
|
|
+ "replace_candidate_ledger",
|
|
|
+ side_effect=fail_first_publish,
|
|
|
+ ):
|
|
|
+ with self.assertRaises(RecoverableToolExecutionError):
|
|
|
+ await registry.execute(
|
|
|
+ "manage_candidate",
|
|
|
+ request,
|
|
|
+ context=runner._build_tool_context(
|
|
|
+ config=config,
|
|
|
+ trace=trace,
|
|
|
+ trace_id=trace.trace_id,
|
|
|
+ goal_id=None,
|
|
|
+ goal_tree=None,
|
|
|
+ sequence=5,
|
|
|
+ head_sequence=4,
|
|
|
+ tool_call_id="recoverable-candidate-call",
|
|
|
+ side_branch_ctx=None,
|
|
|
+ trigger_event=None,
|
|
|
+ ),
|
|
|
+ allowed_tool_names={"manage_candidate"},
|
|
|
+ tool_call_id="recoverable-candidate-call",
|
|
|
+ )
|
|
|
+ pending = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual("pending", pending.operations[-1].status)
|
|
|
+ self.assertEqual(1, len(self.repository.contents))
|
|
|
+ self.assertFalse(any(
|
|
|
+ item.role == "tool"
|
|
|
+ for item in await self.store.get_trace_messages("child-a")
|
|
|
+ ))
|
|
|
+
|
|
|
+ healed, next_sequence = await runner._heal_orphaned_tool_calls(
|
|
|
+ [assistant],
|
|
|
+ await self.store.get_trace("child-a"),
|
|
|
+ None,
|
|
|
+ config,
|
|
|
+ 5,
|
|
|
+ )
|
|
|
+ self.assertEqual(6, next_sequence)
|
|
|
+ self.assertEqual("tool", healed[-1].role)
|
|
|
+ self.assertEqual(
|
|
|
+ "recoverable-candidate-call",
|
|
|
+ healed[-1].tool_call_id,
|
|
|
+ )
|
|
|
+ self.assertEqual(1, len(self.repository.contents))
|
|
|
+ completed = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ matching = [
|
|
|
+ item for item in completed.operations
|
|
|
+ if item.command_id == "recoverable-candidate-call"
|
|
|
+ ]
|
|
|
+ self.assertEqual(1, len(matching))
|
|
|
+ self.assertEqual("completed", matching[0].status)
|
|
|
+ self.assertEqual(1, len(completed.candidates))
|
|
|
+
|
|
|
async def test_descendant_adoption_blocks_ancestor_rewind_without_mapping(self):
|
|
|
await self._create_child("child-a")
|
|
|
await self.store.create_trace(Trace(
|