Record the investigation when the model endpoint fails partway

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

config/schemas/v1/agent-investigation.schema.json

@@ -54,7 +54,8 @@
           "enum": [
             "answered",
             "unstructured",
-            "turn-limit"
+            "turn-limit",
+            "model-unavailable"
           ]
         },
         "turns": {

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

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: