|
|
@@ -18,6 +18,7 @@ from cyber_agent.application.candidate import (
|
|
|
CandidateLedger,
|
|
|
CandidateOperationRecord,
|
|
|
CandidatePointer,
|
|
|
+ CandidateRepositoryRejected,
|
|
|
CandidateReviewAction,
|
|
|
MAX_CANDIDATES_PER_ROOT,
|
|
|
MAX_CANDIDATE_OPERATIONS_PER_ROOT,
|
|
|
@@ -1304,6 +1305,443 @@ class CandidateServiceTest(unittest.IsolatedAsyncioTestCase):
|
|
|
self.assertEqual("completed", matching[0].status)
|
|
|
self.assertEqual(1, len(completed.candidates))
|
|
|
|
|
|
+ async def test_approved_candidate_publish_failure_keeps_original_grant(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,
|
|
|
+ )
|
|
|
+ request = {
|
|
|
+ "request": {
|
|
|
+ "operation": "create",
|
|
|
+ "content": {"text": "approved durable candidate"},
|
|
|
+ "parent_refs": [],
|
|
|
+ },
|
|
|
+ }
|
|
|
+ assistant = TraceMessage.create(
|
|
|
+ trace_id="child-a",
|
|
|
+ role="assistant",
|
|
|
+ sequence=4,
|
|
|
+ content={
|
|
|
+ "text": "",
|
|
|
+ "tool_calls": [{
|
|
|
+ "id": "approved-recoverable-candidate",
|
|
|
+ "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)
|
|
|
+ batch = runner._create_tool_approval_batch(
|
|
|
+ trace_id="child-a",
|
|
|
+ assistant_msg=assistant,
|
|
|
+ tool_calls=assistant.content["tool_calls"],
|
|
|
+ auto_execute_tools=False,
|
|
|
+ )
|
|
|
+ batch.calls[0].decision = "approved"
|
|
|
+ batch.status = "decided"
|
|
|
+ await self.store.replace_tool_approval_batch("child-a", batch)
|
|
|
+ config = RunConfig(
|
|
|
+ trace_id="child-a",
|
|
|
+ uid="user-1",
|
|
|
+ tools=["manage_candidate"],
|
|
|
+ tool_groups=[],
|
|
|
+ auto_execute_tools=False,
|
|
|
+ approval_batch_id=batch.batch_id,
|
|
|
+ )
|
|
|
+ original_replace = self.store.replace_candidate_ledger
|
|
|
+ writes = 0
|
|
|
+
|
|
|
+ async def fail_first_terminal_publish(root_trace_id, raw_ledger):
|
|
|
+ nonlocal writes
|
|
|
+ writes += 1
|
|
|
+ if writes == 2:
|
|
|
+ raise OSError("approved candidate terminal publish interrupted")
|
|
|
+ await original_replace(root_trace_id, raw_ledger)
|
|
|
+
|
|
|
+ trace = await self.store.get_trace("child-a")
|
|
|
+ original_approval_replace = self.store.replace_tool_approval_batch
|
|
|
+ approval_writes = 0
|
|
|
+
|
|
|
+ async def fail_recoverable_approval_publish(trace_id, raw_batch):
|
|
|
+ nonlocal approval_writes
|
|
|
+ approval_writes += 1
|
|
|
+ if approval_writes == 3:
|
|
|
+ raise OSError("recoverable approval state publish interrupted")
|
|
|
+ await original_approval_replace(trace_id, raw_batch)
|
|
|
+
|
|
|
+ with patch.object(
|
|
|
+ self.store,
|
|
|
+ "replace_candidate_ledger",
|
|
|
+ side_effect=fail_first_terminal_publish,
|
|
|
+ ), patch.object(
|
|
|
+ self.store,
|
|
|
+ "replace_tool_approval_batch",
|
|
|
+ side_effect=fail_recoverable_approval_publish,
|
|
|
+ ):
|
|
|
+ with self.assertRaisesRegex(
|
|
|
+ OSError,
|
|
|
+ "recoverable approval state publish interrupted",
|
|
|
+ ):
|
|
|
+ await runner._resume_approved_tool_batch(
|
|
|
+ trace=trace,
|
|
|
+ history=[assistant.to_llm_dict()],
|
|
|
+ goal_tree=None,
|
|
|
+ config=config,
|
|
|
+ sequence=5,
|
|
|
+ head_seq=4,
|
|
|
+ side_branch_ctx=None,
|
|
|
+ runtime_tool_names={"manage_candidate"},
|
|
|
+ )
|
|
|
+
|
|
|
+ interrupted = await self.store.get_tool_approval_batch("child-a")
|
|
|
+ self.assertEqual("executing", interrupted.status)
|
|
|
+ self.assertEqual("executing", interrupted.calls[0].execution_status)
|
|
|
+ self.assertFalse(
|
|
|
+ await runner._mark_approval_execution_unknown("child-a")
|
|
|
+ )
|
|
|
+ recoverable = await self.store.get_tool_approval_batch("child-a")
|
|
|
+ self.assertEqual("recoverable_candidate_command", recoverable.status)
|
|
|
+ self.assertEqual("approved", recoverable.calls[0].decision)
|
|
|
+ self.assertEqual("not_started", recoverable.calls[0].execution_status)
|
|
|
+ self.assertEqual(
|
|
|
+ "approved-recoverable-candidate",
|
|
|
+ recoverable.calls[0].tool_call_id,
|
|
|
+ )
|
|
|
+ with self.assertRaisesRegex(
|
|
|
+ ValueError,
|
|
|
+ "recoverable_candidate_command_must_resume_first",
|
|
|
+ ):
|
|
|
+ await runner._prepare_trace(
|
|
|
+ [{"role": "user", "content": "change direction"}],
|
|
|
+ RunConfig(trace_id="child-a"),
|
|
|
+ )
|
|
|
+ with self.assertRaisesRegex(
|
|
|
+ ValueError,
|
|
|
+ "recoverable_candidate_command_must_resume_first",
|
|
|
+ ):
|
|
|
+ await runner._prepare_trace(
|
|
|
+ [],
|
|
|
+ RunConfig(trace_id="child-a", after_sequence=3),
|
|
|
+ )
|
|
|
+
|
|
|
+ from cyber_agent.trace import run_api
|
|
|
+
|
|
|
+ async def restore_for_api(**_kwargs):
|
|
|
+ return runner, RunConfig(
|
|
|
+ trace_id="child-a",
|
|
|
+ approval_batch_id=recoverable.batch_id,
|
|
|
+ )
|
|
|
+
|
|
|
+ with patch.object(run_api, "_get_runner", return_value=runner), patch.object(
|
|
|
+ run_api,
|
|
|
+ "_restore_runner_and_config",
|
|
|
+ side_effect=restore_for_api,
|
|
|
+ ):
|
|
|
+ with self.assertRaisesRegex(
|
|
|
+ run_api.HTTPException,
|
|
|
+ "recoverable_candidate_command_must_resume_first",
|
|
|
+ ) as blocked_api:
|
|
|
+ await run_api.run_trace(
|
|
|
+ "child-a",
|
|
|
+ run_api.TraceRunRequest(messages=[{
|
|
|
+ "role": "user",
|
|
|
+ "content": "must wait for candidate recovery",
|
|
|
+ }]),
|
|
|
+ )
|
|
|
+ self.assertEqual(409, blocked_api.exception.status_code)
|
|
|
+
|
|
|
+ resumed_config = config
|
|
|
+ resumed_config.approval_batch_id = recoverable.batch_id
|
|
|
+ _history, next_sequence, head_sequence = (
|
|
|
+ await runner._resume_approved_tool_batch(
|
|
|
+ trace=await self.store.get_trace("child-a"),
|
|
|
+ history=[assistant.to_llm_dict()],
|
|
|
+ goal_tree=None,
|
|
|
+ config=resumed_config,
|
|
|
+ sequence=5,
|
|
|
+ head_seq=4,
|
|
|
+ side_branch_ctx=None,
|
|
|
+ runtime_tool_names={"manage_candidate"},
|
|
|
+ )
|
|
|
+ )
|
|
|
+ completed_batch = await self.store.get_tool_approval_batch("child-a")
|
|
|
+ completed_ledger = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual("completed", completed_batch.status)
|
|
|
+ self.assertEqual(6, next_sequence)
|
|
|
+ self.assertEqual(5, head_sequence)
|
|
|
+ self.assertEqual(1, len(completed_ledger.candidates))
|
|
|
+ self.assertEqual(1, len(self.repository.contents))
|
|
|
+ self.assertEqual(2, self.repository.put_calls)
|
|
|
+
|
|
|
+ # Crash window: the Tool Result is durable, while the last approval
|
|
|
+ # update still shows the call as executing. Recovery must finalize the
|
|
|
+ # existing message instead of replaying the candidate command.
|
|
|
+ completed_batch.status = "executing"
|
|
|
+ completed_batch.calls[0].execution_status = "executing"
|
|
|
+ completed_batch.calls[0].result_message_id = None
|
|
|
+ await self.store.replace_tool_approval_batch("child-a", completed_batch)
|
|
|
+ await self.store.update_trace("child-a", head_sequence=4)
|
|
|
+ self.assertTrue(
|
|
|
+ await runner._recover_candidate_approval_if_durable(
|
|
|
+ "child-a",
|
|
|
+ completed_batch,
|
|
|
+ )
|
|
|
+ )
|
|
|
+ reconciled = await self.store.get_tool_approval_batch("child-a")
|
|
|
+ trace_after_reconcile = await self.store.get_trace("child-a")
|
|
|
+ matching_results = [
|
|
|
+ message
|
|
|
+ for message in await self.store.get_trace_messages("child-a")
|
|
|
+ if message.role == "tool"
|
|
|
+ and message.tool_call_id == "approved-recoverable-candidate"
|
|
|
+ ]
|
|
|
+ self.assertEqual("completed", reconciled.status)
|
|
|
+ self.assertEqual("executed", reconciled.calls[0].execution_status)
|
|
|
+ self.assertEqual(
|
|
|
+ matching_results[0].message_id,
|
|
|
+ reconciled.calls[0].result_message_id,
|
|
|
+ )
|
|
|
+ self.assertEqual(
|
|
|
+ matching_results[0].sequence,
|
|
|
+ trace_after_reconcile.head_sequence,
|
|
|
+ )
|
|
|
+ self.assertEqual(1, len(matching_results))
|
|
|
+ self.assertEqual(2, self.repository.put_calls)
|
|
|
+
|
|
|
+ async def test_failed_outcome_publish_failure_remains_orphan_replayable(self):
|
|
|
+ await self._create_child("child-a")
|
|
|
+ self.repository.bad_artifact = True
|
|
|
+ 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": "invalid artifact candidate"},
|
|
|
+ "parent_refs": [],
|
|
|
+ },
|
|
|
+ }
|
|
|
+ assistant = TraceMessage.create(
|
|
|
+ trace_id="child-a",
|
|
|
+ role="assistant",
|
|
|
+ sequence=4,
|
|
|
+ content={
|
|
|
+ "text": "",
|
|
|
+ "tool_calls": [{
|
|
|
+ "id": "failed-publish-candidate",
|
|
|
+ "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_failed_publish(root_trace_id, raw_ledger):
|
|
|
+ nonlocal writes
|
|
|
+ writes += 1
|
|
|
+ if writes == 2:
|
|
|
+ raise OSError("failed outcome publish interrupted")
|
|
|
+ await original_replace(root_trace_id, raw_ledger)
|
|
|
+
|
|
|
+ with patch.object(
|
|
|
+ self.store,
|
|
|
+ "replace_candidate_ledger",
|
|
|
+ side_effect=fail_first_failed_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="failed-publish-candidate",
|
|
|
+ side_branch_ctx=None,
|
|
|
+ trigger_event=None,
|
|
|
+ ),
|
|
|
+ allowed_tool_names={"manage_candidate"},
|
|
|
+ tool_call_id="failed-publish-candidate",
|
|
|
+ )
|
|
|
+ pending = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual("pending", pending.operations[-1].status)
|
|
|
+ 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,
|
|
|
+ )
|
|
|
+ terminal = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual(6, next_sequence)
|
|
|
+ self.assertEqual("tool", healed[-1].role)
|
|
|
+ self.assertEqual("failed-publish-candidate", healed[-1].tool_call_id)
|
|
|
+ self.assertEqual("failed", terminal.operations[-1].status)
|
|
|
+ self.assertEqual(0, len(terminal.candidates))
|
|
|
+
|
|
|
+ async def test_pending_recovery_survives_transient_repository_error(self):
|
|
|
+ await self._create_child("child-a")
|
|
|
+ command_id = "transient-recovery-candidate"
|
|
|
+ request = {
|
|
|
+ "operation": "create",
|
|
|
+ "content": {"text": "eventually consistent candidate"},
|
|
|
+ "parent_refs": [],
|
|
|
+ "effective_at_sequence": 5,
|
|
|
+ "command_id": command_id,
|
|
|
+ }
|
|
|
+ original_replace = self.store.replace_candidate_ledger
|
|
|
+ writes = 0
|
|
|
+
|
|
|
+ async def fail_first_terminal_publish(root_trace_id, raw_ledger):
|
|
|
+ nonlocal writes
|
|
|
+ writes += 1
|
|
|
+ if writes == 2:
|
|
|
+ raise OSError("candidate terminal publish interrupted")
|
|
|
+ await original_replace(root_trace_id, raw_ledger)
|
|
|
+
|
|
|
+ with patch.object(
|
|
|
+ self.store,
|
|
|
+ "replace_candidate_ledger",
|
|
|
+ side_effect=fail_first_terminal_publish,
|
|
|
+ ):
|
|
|
+ with self.assertRaises(RecoverableToolExecutionError):
|
|
|
+ await self.service.manage("child-a", **request)
|
|
|
+
|
|
|
+ with patch.object(
|
|
|
+ self.repository,
|
|
|
+ "put_version",
|
|
|
+ side_effect=ValueError("repository temporarily unavailable"),
|
|
|
+ ):
|
|
|
+ with self.assertRaises(RecoverableToolExecutionError):
|
|
|
+ await self.service.manage("child-a", **request)
|
|
|
+ still_pending = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual("pending", still_pending.operations[-1].status)
|
|
|
+
|
|
|
+ with patch.object(
|
|
|
+ self.repository,
|
|
|
+ "resolve",
|
|
|
+ side_effect=RuntimeError("artifact resolver temporarily unavailable"),
|
|
|
+ ):
|
|
|
+ with self.assertRaises(RecoverableToolExecutionError):
|
|
|
+ await self.service.manage("child-a", **request)
|
|
|
+ still_pending = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual("pending", still_pending.operations[-1].status)
|
|
|
+
|
|
|
+ recovered = await self.service.manage("child-a", **request)
|
|
|
+ completed = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual(recovered, completed.candidates[-1])
|
|
|
+ self.assertEqual("completed", completed.operations[-1].status)
|
|
|
+ self.assertEqual(1, len(self.repository.contents))
|
|
|
+
|
|
|
+ async def test_first_repository_timeout_after_commit_remains_replayable(self):
|
|
|
+ await self._create_child("child-a")
|
|
|
+ request = {
|
|
|
+ "operation": "create",
|
|
|
+ "content": {"text": "committed before response timeout"},
|
|
|
+ "parent_refs": [],
|
|
|
+ "effective_at_sequence": 5,
|
|
|
+ "command_id": "first-call-timeout-candidate",
|
|
|
+ }
|
|
|
+ original_put = self.repository.put_version
|
|
|
+
|
|
|
+ async def commit_then_timeout(version_request):
|
|
|
+ await original_put(version_request)
|
|
|
+ raise TimeoutError("response was lost after commit")
|
|
|
+
|
|
|
+ with patch.object(
|
|
|
+ self.repository,
|
|
|
+ "put_version",
|
|
|
+ side_effect=commit_then_timeout,
|
|
|
+ ):
|
|
|
+ with self.assertRaises(RecoverableToolExecutionError):
|
|
|
+ await self.service.manage("child-a", **request)
|
|
|
+ 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))
|
|
|
+
|
|
|
+ recovered = await self.service.manage("child-a", **request)
|
|
|
+ completed = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual(recovered, completed.candidates[-1])
|
|
|
+ self.assertEqual("completed", completed.operations[-1].status)
|
|
|
+ self.assertEqual(1, len(self.repository.contents))
|
|
|
+
|
|
|
+ async def test_typed_repository_rejection_is_terminal(self):
|
|
|
+ await self._create_child("child-a")
|
|
|
+ with patch.object(
|
|
|
+ self.repository,
|
|
|
+ "put_version",
|
|
|
+ side_effect=CandidateRepositoryRejected("content policy rejected"),
|
|
|
+ ):
|
|
|
+ with self.assertRaisesRegex(
|
|
|
+ CandidateStateError,
|
|
|
+ "content policy rejected",
|
|
|
+ ):
|
|
|
+ await self.service.manage(
|
|
|
+ "child-a",
|
|
|
+ operation="create",
|
|
|
+ content={"text": "rejected candidate"},
|
|
|
+ parent_refs=[],
|
|
|
+ effective_at_sequence=5,
|
|
|
+ command_id="deterministic-rejection-candidate",
|
|
|
+ )
|
|
|
+ failed = CandidateLedger.model_validate(
|
|
|
+ await self.store.get_candidate_ledger("root-a")
|
|
|
+ )
|
|
|
+ self.assertEqual("failed", failed.operations[-1].status)
|
|
|
+ self.assertEqual(0, len(failed.candidates))
|
|
|
+ self.assertEqual(0, len(self.repository.contents))
|
|
|
+
|
|
|
async def test_descendant_adoption_blocks_ancestor_rewind_without_mapping(self):
|
|
|
await self._create_child("child-a")
|
|
|
await self.store.create_trace(Trace(
|