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, [])