From 1ebab9e4bc1572973917207b6c445a44ef74de5a Mon Sep 17 00:00:00 2001 From: "shipmill-romamo[bot]" <338557544+shipmill-romamo[bot]@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:47:09 +0300 Subject: [PATCH] Record the investigation when the model endpoint fails partway agent investigate now saves its AgentInvestigation, ended model-unavailable with the steps, turns and tokens so far, before exiting 92 MODEL_UNAVAILABLE, and names it in error.context.investigation. Decisions proposed in earlier turns are no longer left without the investigation that produced them. The chat transport raises InvestigationError instead of exiting, so investigate() can keep the run up to the failure. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 8 ++ .../v1/agent-investigation.schema.json | 3 +- sdk/src/cloudfall/app.py | 58 ++++---- sdk/src/cloudfall/investigation.py | 122 +++++++++++------ sdk/tests/test_investigation.py | 129 +++++++++++++++++- 5 files changed, 248 insertions(+), 72 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f72020f..b204d36 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,14 @@ Notable changes to Cloudfall. The format follows behind a PostgreSQL alert was traced to unrotated application logs in five turns, and the two proposals it recorded, once approved, brought the database back +- `agent investigate` records the investigation when the model endpoint + fails partway (refused, unreachable, not JSON, or a malformed answer), + so the decisions earlier turns proposed are never left without one: the + record ends `model-unavailable` with the steps, turns and tokens so far, + an empty `answer` and no `finding`, the command still exits 92 + `MODEL_UNAVAILABLE`, and `error.context.investigation` names the record. + A failure on the first request records an investigation with no steps; + a refusal before any request (exit 2) still records nothing (#29) - An operation tool of `mcp serve --repository` describes a read operation as running and returning what the hosts reported, instead of saying it runs check mode only diff --git a/config/schemas/v1/agent-investigation.schema.json b/config/schemas/v1/agent-investigation.schema.json index 32d9fee..fec9092 100644 --- a/config/schemas/v1/agent-investigation.schema.json +++ b/config/schemas/v1/agent-investigation.schema.json @@ -54,7 +54,8 @@ "enum": [ "answered", "unstructured", - "turn-limit" + "turn-limit", + "model-unavailable" ] }, "turns": { diff --git a/sdk/src/cloudfall/app.py b/sdk/src/cloudfall/app.py index 7c73b2a..37abcab 100644 --- a/sdk/src/cloudfall/app.py +++ b/sdk/src/cloudfall/app.py @@ -128,13 +128,18 @@ from cloudfall.inventory import PlatformInventory from cloudfall.investigation import ( DEFAULT_MAX_TURNS, + ERROR_MODEL_ANSWER_MALFORMED, + ERROR_MODEL_REFUSED, + ERROR_MODEL_UNREACHABLE, INVESTIGATION_DIRECTORY, AgentTool, ChatTransport, + Investigation, InvestigationError, InvestigationStatus, InvestigationStore, ModelEndpoint, + ModelUnavailableError, StepOutcome, ToolResult, host_text, @@ -2119,7 +2124,8 @@ def decisions_approve(args: ApproveDecisionArgs, ctx: Ctx) -> Decided: side_effects="partial", suggestion=( "check the endpoint, the model name and the API key, then run again; " - "decisions already recorded stay" + "decisions already recorded stay, and the investigation that recorded " + "them is named in error.context.investigation" ), ) app.exit_code( @@ -2250,6 +2256,10 @@ def agent_investigate(args: InvestigateArgs, ctx: Ctx) -> Investigated: """Work the alert through the catalog and keep the record.""" catalog = args.catalog() tools = [_agent_tool(args, operation, ctx) for operation in catalog.operations] + store = InvestigationStore( + directory=args.root / args.investigations, + catalog=SchemaCatalog(args.schemas), + ) try: investigation = investigate( args.alert, @@ -2258,18 +2268,18 @@ def agent_investigate(args: InvestigateArgs, ctx: Ctx) -> Investigated: _chat_transport(ctx, args.api_key), max_turns=args.max_turns, ) - except InvestigationError as error: + except ModelUnavailableError as error: + # Earlier turns may have recorded decisions: the investigation that + # produced them is recorded too, then the model's failure is reported. + failed = _save_investigation(store, error.investigation) raise Exit.MODEL_UNAVAILABLE( - error.detail, context={"code": error.code} + error.detail, + context={ + "code": error.code, + "investigation": failed.investigation_id.value, + }, ) from error - try: - investigation = InvestigationStore( - directory=args.root / args.investigations, - catalog=SchemaCatalog(args.schemas), - ).save(investigation) - except InvestigationError as error: - # Recording failed, not the model: never MODEL_UNAVAILABLE. - raise Exit.PRECONDITION(error.detail, context={"code": error.code}) from error + investigation = _save_investigation(store, investigation) investigated = Investigated( effect="created", status=investigation.status.value, @@ -2281,6 +2291,16 @@ def agent_investigate(args: InvestigateArgs, ctx: Ctx) -> Investigated: return investigated +def _save_investigation( + store: InvestigationStore, investigation: Investigation +) -> Investigation: + """Record the investigation; a failure to record it is never the model's.""" + try: + return store.save(investigation) + except InvestigationError as error: + raise Exit.PRECONDITION(error.detail, context={"code": error.code}) from error + + def _agent_tool(args: InvestigateArgs, operation: Operation, ctx: Ctx) -> AgentTool: """One declared operation as a tool of the investigating model.""" @@ -2351,29 +2371,21 @@ def send(url: str, body: Mapping[str, object]) -> Mapping[str, object]: # ctx.http's own exit says the command had no side effects, but # decisions of earlier turns may already be recorded. message = f"the model endpoint failed: {error.message}" - raise Exit.MODEL_UNAVAILABLE( - message, context={"code": "agent_model_unreachable"} - ) from error + raise InvestigationError(ERROR_MODEL_UNREACHABLE, message) from error if response.status != HTTPStatus.OK: message = ( f"the model endpoint answered {response.status}: " f"{response.body[:300].decode('utf-8', errors='replace')}" ) - raise Exit.MODEL_UNAVAILABLE( - message, context={"code": "agent_model_refused"} - ) + raise InvestigationError(ERROR_MODEL_REFUSED, message) try: answer = response.json() except ValueError as error: message = "the model endpoint answered something that is not JSON" - raise Exit.MODEL_UNAVAILABLE( - message, context={"code": "agent_model_answer_malformed"} - ) from error + raise InvestigationError(ERROR_MODEL_ANSWER_MALFORMED, message) from error if not isinstance(answer, Mapping): message = "the model endpoint answered JSON that is not an object" - raise Exit.MODEL_UNAVAILABLE( - message, context={"code": "agent_model_answer_malformed"} - ) + raise InvestigationError(ERROR_MODEL_ANSWER_MALFORMED, message) return cast("Mapping[str, object]", answer) return send diff --git a/sdk/src/cloudfall/investigation.py b/sdk/src/cloudfall/investigation.py index 3f95e62..1d4eab7 100644 --- a/sdk/src/cloudfall/investigation.py +++ b/sdk/src/cloudfall/investigation.py @@ -46,6 +46,8 @@ ERROR_ENDPOINT_INVALID = "agent_endpoint_invalid" ERROR_MODEL_ANSWER_MALFORMED = "agent_model_answer_malformed" +ERROR_MODEL_REFUSED = "agent_model_refused" +ERROR_MODEL_UNREACHABLE = "agent_model_unreachable" ERROR_INVESTIGATION_EXISTS = "agent_investigation_exists" SYSTEM_PROMPT = """\ @@ -90,6 +92,9 @@ class InvestigationStatus(StrEnum): TURN_LIMIT = "turn-limit" """The model was still calling tools when its turns ran out.""" + MODEL_UNAVAILABLE = "model-unavailable" + """The endpoint failed mid-run; the record holds the steps before it.""" + class StepOutcome(StrEnum): """What came of one tool call.""" @@ -160,7 +165,10 @@ class AgentTool: ChatTransport = Callable[[str, Mapping[str, object]], Mapping[str, object]] -"""Send one chat request body to a URL and return the answer's JSON.""" +"""Send one chat request body to a URL and return the answer's JSON. + +A failed request raises ``InvestigationError``. +""" @dataclass(frozen=True, slots=True) @@ -264,6 +272,19 @@ def as_document(self) -> dict[str, object]: } +class ModelUnavailableError(InvestigationError): + """The model endpoint failed after the run started. + + It carries the investigation up to the failure, ended ``model-unavailable``, + so the steps that already recorded decisions are kept too. + """ + + def __init__(self, cause: InvestigationError, investigation: Investigation) -> None: + """Keep the cause's code and message beside the investigation so far.""" + super().__init__(cause.code, cause.detail) + self.investigation = investigation + + @dataclass(frozen=True, slots=True) class InvestigationStore: """Schema-validated investigation records in one directory.""" @@ -305,7 +326,11 @@ def investigate( *, max_turns: int = DEFAULT_MAX_TURNS, ) -> Investigation: - """Let the model work the alert through the tools until it answers.""" + """Let the model work the alert through the tools until it answers. + + When the endpoint fails or answers malformed once the run started, this + raises ``ModelUnavailableError`` carrying the investigation so far. + """ started = _utc_now() by_name = {tool.name: tool for tool in tools} messages: list[Mapping[str, object]] = [ @@ -319,53 +344,64 @@ def investigate( answer = "" status = InvestigationStatus.TURN_LIMIT turns = 0 - while turns < max_turns: - turns += 1 - reply = transport( - endpoint.chat_url, - { - "model": endpoint.model, - "messages": messages, - "tools": request_tools, - "temperature": 0.2, - }, + + def record(ending: InvestigationStatus, finding: Finding | None) -> Investigation: + return Investigation( + investigation_id=ResourceId.from_boundary( + f"investigation-{started.strftime('%Y%m%d%H%M%S')}" + ), + alert=alert, + endpoint=endpoint, + started_at=_timestamp(started), + finished_at=_timestamp(_utc_now()), + status=ending, + turns=turns, + usage=usage, + steps=tuple(steps), + answer=answer, + finding=finding, ) - usage = usage.plus(reply) - message = _first_message(reply) - messages.append(message) - calls = message.get("tool_calls") - if not calls: - answer = str(message.get("content") or "").strip() - status = InvestigationStatus.UNSTRUCTURED - break - for call in cast("Sequence[Mapping[str, object]]", calls): - step, result = _run_call(call, by_name, seen) - steps.append(step) - messages.append( + + try: + while turns < max_turns: + turns += 1 + reply = transport( + endpoint.chat_url, { - "role": "tool", - "tool_call_id": str(call.get("id", "")), - "content": json.dumps(result.content, sort_keys=True), - } + "model": endpoint.model, + "messages": messages, + "tools": request_tools, + "temperature": 0.2, + }, ) + usage = usage.plus(reply) + message = _first_message(reply) + messages.append(message) + calls = message.get("tool_calls") + if not calls: + answer = str(message.get("content") or "").strip() + status = InvestigationStatus.UNSTRUCTURED + break + for call in cast("Sequence[Mapping[str, object]]", calls): + step, result = _run_call(call, by_name, seen) + steps.append(step) + messages.append( + { + "role": "tool", + "tool_call_id": str(call.get("id", "")), + "content": json.dumps(result.content, sort_keys=True), + } + ) + except InvestigationError as error: + # The record names no error text: an endpoint's answer can echo the + # request, and the record is kept beside the decisions. The answer is + # still empty: the loop ends at the model's last words. + failed = record(InvestigationStatus.MODEL_UNAVAILABLE, None) + raise ModelUnavailableError(error, failed) from error finding = _finding(answer, steps) if answer else None if finding is not None: status = InvestigationStatus.ANSWERED - return Investigation( - investigation_id=ResourceId.from_boundary( - f"investigation-{started.strftime('%Y%m%d%H%M%S')}" - ), - alert=alert, - endpoint=endpoint, - started_at=_timestamp(started), - finished_at=_timestamp(_utc_now()), - status=status, - turns=turns, - usage=usage, - steps=tuple(steps), - answer=answer, - finding=finding, - ) + return record(status, finding) def host_text(text: str) -> str: diff --git a/sdk/tests/test_investigation.py b/sdk/tests/test_investigation.py index 69e6f0e..477a8de 100644 --- a/sdk/tests/test_investigation.py +++ b/sdk/tests/test_investigation.py @@ -11,11 +11,13 @@ import pytest from cloudfall.cli import main from cloudfall.investigation import ( + ERROR_MODEL_REFUSED, AgentTool, InvestigationError, InvestigationStatus, InvestigationStore, ModelEndpoint, + ModelUnavailableError, StepOutcome, ToolResult, investigate, @@ -265,6 +267,54 @@ def test_a_malformed_model_answer_is_refused() -> None: assert error.value.code == "agent_model_answer_malformed" +class _Failing(_Script): + """A model that answers from a list, then its endpoint refuses.""" + + def __call__(self, url: str, body: Mapping[str, object]) -> Mapping[str, object]: + if not self.replies: + message = "the model endpoint answered 500: overloaded" + raise InvestigationError(ERROR_MODEL_REFUSED, message) + return super().__call__(url, body) + + +@pytest.mark.parametrize( + "script", + [ + _Failing(_call("shop-logrotate", {})), + _Script( + _call("shop-logrotate", {}), + {"choices": [{"message": {"tool_calls": [{"id": "c2"}]}}]}, + ), + ], + ids=["refused", "malformed-tool-call"], +) +def test_a_model_failing_mid_run_keeps_the_investigation_so_far( + tmp_path: Path, script: _Script +) -> None: + report, rotate = _tools() + + with pytest.raises(ModelUnavailableError) as error: + investigate("alert", _tool_list(report, rotate), ENDPOINT, script) + + failed = error.value.investigation + assert failed.status is InvestigationStatus.MODEL_UNAVAILABLE + assert [step.decision for step in failed.steps] == [ + "shop-logrotate-20261008152628" + ] + assert failed.turns == 2 + assert (failed.usage.prompt, failed.usage.completion) == (100, 10) + assert failed.answer == "" + assert failed.finding is None + saved = InvestigationStore( + directory=tmp_path / "investigations", + catalog=SchemaCatalog(default_schema_directory()), + ).save(failed) + spec = saved.as_document()["spec"] + assert isinstance(spec, dict) + assert spec["status"] == "model-unavailable" + assert "finding" not in spec + + @pytest.mark.parametrize( "url", ["ftp://models.example/v1", "https://user:secret@models.example/v1"] ) @@ -276,9 +326,12 @@ def test_an_endpoint_must_be_plain_http(url: str) -> None: class _Model(BaseHTTPRequestHandler): - """An OpenAI-compatible endpoint answering from a class-level script.""" + """An OpenAI-compatible endpoint answering from a class-level script. - replies: ClassVar[list[dict[str, object]]] = [] + A reply is JSON sent with 200, or a status and the raw body to send. + """ + + replies: ClassVar[list[object]] = [] seen: ClassVar[list[dict[str, Any]]] = [] def do_POST(self) -> None: @@ -286,8 +339,11 @@ def do_POST(self) -> None: _Model.seen.append( {"auth": self.headers["Authorization"], "body": json.loads(body)} ) - answer = json.dumps(_Model.replies.pop(0)).encode() - self.send_response(200) + reply = _Model.replies.pop(0) + status, answer = ( + reply if isinstance(reply, tuple) else (200, json.dumps(reply).encode()) + ) + self.send_response(status) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(answer))) self.end_headers() @@ -412,6 +468,63 @@ def test_an_answer_that_is_not_json_exits_model_unavailable( assert payload["error"]["context"]["code"] == "agent_model_answer_malformed" +@pytest.mark.parametrize( + ("failure", "code"), + [ + # The refusal echoes the request's key, as a debugging proxy might. + ((500, b"upstream failed for Bearer sk-test"), "agent_model_refused"), + ((200, b"gateway"), "agent_model_answer_malformed"), + ((200, b"[]"), "agent_model_answer_malformed"), + ({"choices": []}, "agent_model_answer_malformed"), + ( + {"choices": [{"message": {"tool_calls": [{"id": "c2"}]}}]}, + "agent_model_answer_malformed", + ), + ], + ids=["refused", "not-json", "not-an-object", "no-choices", "bad-tool-call"], +) +def test_a_model_failing_mid_run_records_the_investigation_and_exits_92( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + model_url: str, + failure: object, + code: str, +) -> None: + repository = _repository(tmp_path) + (repository / "playbooks" / "facts.yml").write_text( + "---\n- hosts: localhost\n connection: local\n gather_facts: false\n" + " tasks:\n - name: Report\n ansible.builtin.debug:\n" + " msg: /var/log/shop holds 36G\n", + encoding="utf-8", + ) + key = tmp_path / "model.key" + key.write_text("sk-test\n", encoding="utf-8") + _Model.seen.clear() + _Model.replies[:] = [_call("operation_facts", {}), failure] + + exit_code = _investigate_against(repository, key, model_url) + + out = capsys.readouterr() + assert "sk-test" not in out.out + out.err + payload = json.loads(out.out) + assert exit_code == 92, payload + context = payload["error"]["context"] + assert context["code"] == code + recorded = repository / "investigations" / f"{context['investigation']}.json" + text = recorded.read_text(encoding="utf-8") + assert "sk-test" not in text + spec = json.loads(text)["spec"] + assert spec["status"] == "model-unavailable" + assert spec["turns"] == 2 + assert spec["tokens"] == {"prompt": 100, "completion": 10} + assert spec["answer"] == "" + assert "finding" not in spec + [step] = spec["steps"] + assert step["outcome"] == "ran" + decisions = list((repository / "decisions").glob(f"{step['decision']}*")) + assert decisions, "the decision of the first turn is recorded" + + @pytest.mark.parametrize( "url", [ @@ -505,4 +618,10 @@ def test_an_unreachable_endpoint_exits_model_unavailable( payload = json.loads(capsys.readouterr().out) assert exit_code == 92, payload - assert payload["error"]["context"]["code"] == "agent_model_unreachable" + context = payload["error"]["context"] + assert context["code"] == "agent_model_unreachable" + # The first request was sent: the run started, so it is recorded, empty. + recorded = repository / "investigations" / f"{context['investigation']}.json" + spec = json.loads(recorded.read_text(encoding="utf-8"))["spec"] + assert spec["status"] == "model-unavailable" + assert (spec["turns"], spec["steps"]) == (1, [])