Record a stopped investigation when its reader goes away

CHANGELOG.md

@@ -431,6 +431,16 @@ Notable changes to Cloudfall. The format follows
 
 ### Fixed
 
+- `agent investigate` keeps its investigation record when the reader closes
+  stdout mid-run (such as `| head -2`). Since the command streams, such a
+  run left the decisions it proposed without the investigation that made
+  them, and exited 0. The record is now saved ended `stopped`, a new status
+  in `agent-investigation.schema.json`, with the steps, turns and tokens so
+  far, an empty `answer` and no `finding`; nothing more is written to the
+  closed stdout, and the command exits 94 `INVESTIGATION_STOPPED`. A run
+  stopped by a signal or a timeout keeps its own exit code, and one stopped
+  while it waits on the model still leaves no record;
+  `--no-stream` is unchanged (#44)
 - The operator proposal schema refuses `approval.via` on an autonomous
   (policy) approval. Such a receipt used to pass validation and lose the
   field on load; it now fails schema validation like any other invalid

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

@@ -55,7 +55,8 @@
             "answered",
             "unstructured",
             "turn-limit",
-            "model-unavailable"
+            "model-unavailable",
+            "stopped"
           ]
         },
         "turns": {

sdk/src/cloudfall/app.py

@@ -18,6 +18,7 @@
 import errno
 import os
 import shutil
+import threading
 
 # treaty reads a handler's annotations at registration, so a streaming
 # handler's Iterator and a provided MCP tool's Mapping must exist at runtime.
@@ -2186,6 +2187,30 @@ def decisions_approve(args: ApproveDecisionArgs, ctx: Ctx) -> Decided:
         "then run again or raise --max-turns"
     ),
 )
+INVESTIGATION_STOPPED = 94
+"""The exit code of an ``agent investigate`` whose reader closed stdout mid-run."""
+app.exit_code(
+    "INVESTIGATION_STOPPED",
+    INVESTIGATION_STOPPED,
+    description=(
+        "The reader closed stdout mid-run; the investigation so far is recorded "
+        "in investigations/, ended stopped"
+    ),
+    retryable=False,
+    side_effects="partial",
+    suggestion=(
+        "read the newest record in investigations/ and the decisions it names, "
+        "then run again with a reader that stays to the end"
+    ),
+)
+stream_stopped = threading.Event()
+"""Set when an ``agent investigate`` stream was closed mid-run and recorded so.
+
+treaty ends a stream whose reader went away after an event with exit 0, and
+the handler has no say in it, so ``cloudfall.cli`` reads this to exit
+``INVESTIGATION_STOPPED`` instead. A stream closed by a signal or a timeout
+keeps the exit code treaty gives it.
+"""
 
 
 @dataclass(frozen=True, slots=True, kw_only=True)
@@ -2302,6 +2327,7 @@ class InvestigationEvent:
         "PRECONDITION",
         "MODEL_UNAVAILABLE",
         "INVESTIGATION_INCOMPLETE",
+        "INVESTIGATION_STOPPED",
     ],
     timeout=None,
     has_network_io=True,
@@ -2363,6 +2389,17 @@ def agent_investigate(
                 "investigation": failed.investigation_id.value,
             },
         ) from error
+    except GeneratorExit:
+        # The stream was closed before its end, as when the reader closes
+        # stdout: nothing more can be written, but the decisions proposed so
+        # far keep the investigation that made them.
+        stopped = events.close()
+        if stopped is None:
+            message = "the investigation ended before its reader went away"
+            raise RuntimeError(message) from None
+        _save_investigation(store, stopped)
+        stream_stopped.set()
+        raise
     investigation = _save_investigation(store, investigation)
     yield _investigation_event(investigation)
     if investigation.status is not InvestigationStatus.ANSWERED:

sdk/src/cloudfall/cli.py

@@ -5,17 +5,36 @@
 import sys
 from typing import TYPE_CHECKING
 
-from cloudfall.app import app
+from cloudfall.app import INVESTIGATION_STOPPED, app, stream_stopped
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+OUTPUT_CLOSED = (0, 141)
+"""What treaty exits with when the reader closed stdout: after an event, before any."""
+
 
 def main(argv: Sequence[str] | None = None) -> int:
     """Run one command in-process and return its exit code."""
-    return app.run(list(sys.argv[1:] if argv is None else argv))
+    stream_stopped.clear()
+    return _stopped(app.run(list(sys.argv[1:] if argv is None else argv)))
 
 
 def run() -> None:
     """Installed console-script entry point."""
-    app.main()
+    stream_stopped.clear()
+    try:
+        app.main()
+    except SystemExit as exiting:
+        # app.main() exits with treaty's code; only a stopped stream's changes.
+        code = exiting.code
+        if isinstance(code, int) and _stopped(code) != code:
+            raise SystemExit(INVESTIGATION_STOPPED) from None
+        raise
+
+
+def _stopped(code: int) -> int:
+    """Return INVESTIGATION_STOPPED when the reader left an investigation mid-run."""
+    if stream_stopped.is_set() and code in OUTPUT_CLOSED:
+        return INVESTIGATION_STOPPED
+    return code

sdk/src/cloudfall/investigation.py

@@ -100,6 +100,9 @@ class InvestigationStatus(StrEnum):
     MODEL_UNAVAILABLE = "model-unavailable"
     """The endpoint failed mid-run; the record holds the steps before it."""
 
+    STOPPED = "stopped"
+    """Its reader went away mid-run; the record holds the steps before it."""
+
 
 class StepOutcome(StrEnum):
     """What came of one tool call."""
@@ -445,7 +448,9 @@ def investigation_events(  # noqa: PLR0913 - as investigate.
     """Run the investigation, yielding each turn and call as it happens.
 
     The generator's return value is the investigation; a failing endpoint
-    raises ``ModelUnavailableError`` as ``investigate`` does.
+    raises ``ModelUnavailableError`` as ``investigate`` does. Closed before
+    it ends, it returns the investigation so far, ended ``stopped``, from
+    ``close()``.
     """
     started = _utc_now()
     by_name = {tool.name: tool for tool in tools}
@@ -526,6 +531,10 @@ def record(ending: InvestigationStatus, finding: Finding | None) -> Investigatio
         # still empty: the loop ends at the model's last words.
         failed = record(InvestigationStatus.MODEL_UNAVAILABLE, None)
         raise ModelUnavailableError(error, failed) from error
+    except GeneratorExit:
+        # Whoever read the events went away: nothing more is asked or run, and
+        # the steps so far, decisions included, end in a record of their own.
+        return record(InvestigationStatus.STOPPED, None)
     finding = _finding(answer, steps) if answer else None
     if finding is not None:
         status = InvestigationStatus.ANSWERED

sdk/tests/test_investigation.py

@@ -4,6 +4,8 @@
 
 import json
 import socket
+import subprocess
+import sys
 import threading
 from http.server import BaseHTTPRequestHandler, HTTPServer
 from typing import TYPE_CHECKING, Any, ClassVar
@@ -23,6 +25,7 @@
     StepOutcome,
     ToolResult,
     investigate,
+    investigation_events,
 )
 from cloudfall.resources import default_schema_directory
 from cloudfall.validation import SchemaCatalog
@@ -951,3 +954,123 @@ def test_an_over_long_endpoint_id_still_leaves_a_valid_record(
     assert turn.request_id is not None
     assert turn.response_id.startswith("chatcmpl-")
     assert (len(turn.response_id), len(turn.request_id)) == (256, 256)
+
+
+def test_closing_the_events_returns_the_investigation_so_far_stopped() -> None:
+    report, rotate = _tools()
+    script = _Script(_call("shop-logrotate", {}), _say("never asked"))
+    events = investigation_events(
+        "alert", _tool_list(report, rotate), ENDPOINT, script
+    )
+
+    assert next(events).turn == 1  # the first turn
+    assert next(events).turn == 1  # its call, which proposed a decision
+    stopped = events.close()
+
+    assert stopped is not None
+    assert stopped.status is InvestigationStatus.STOPPED
+    assert [step.decision for step in stopped.steps] == [
+        "shop-logrotate-20261008152628"
+    ]
+    assert (stopped.turns, stopped.answer, stopped.finding) == (1, "", None)
+    assert len(script.requests) == 1
+
+
+class _Held(BaseHTTPRequestHandler):
+    """An endpoint whose second answer waits until the test lets it go."""
+
+    replies: ClassVar[list[Mapping[str, object]]] = []
+    answered: ClassVar[list[Mapping[str, object]]] = []
+    released: ClassVar[threading.Event] = threading.Event()
+
+    def do_POST(self) -> None:
+        self.rfile.read(int(self.headers["Content-Length"]))
+        if _Held.answered:
+            assert _Held.released.wait(60), "the test never let the answer go"
+        reply = _Held.replies.pop(0)
+        _Held.answered.append(reply)
+        answer = json.dumps(reply).encode()
+        self.send_response(200)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(answer)))
+        self.end_headers()
+        self.wfile.write(answer)
+
+    def log_message(self, *_args: object) -> None:
+        return
+
+
+def test_a_reader_closing_stdout_mid_run_leaves_the_record_stopped(
+    tmp_path: Path,
+) -> 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")
+    _Held.replies[:] = [
+        _call("operation_facts", {}),
+        _say("ROOT CAUSE: /var/log/shop\nPROPOSED: none\nWHY: read only."),
+    ]
+    _Held.answered.clear()
+    _Held.released.clear()
+    server = HTTPServer(("127.0.0.1", 0), _Held)
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    command = [
+        sys.executable,
+        "-c",
+        "from cloudfall.cli import run; run()",
+        *("agent", "investigate", "--repository", str(repository)),
+        *("--alert", "postgresql-main down on web-1", "--model", "nemotron"),
+        *("--base-url", f"http://127.0.0.1:{server.server_address[1]}/v1"),
+        *("--api-key-from-file", str(key)),
+    ]
+    stderr_path = tmp_path / "stderr.txt"
+    try:
+        with stderr_path.open("w", encoding="utf-8") as stderr:
+            process = subprocess.Popen(  # noqa: S603 - fixed interpreter and arguments.
+                command,
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.PIPE,
+                stderr=stderr,
+                cwd=tmp_path,
+            )
+            try:
+                assert process.stdout is not None
+                # A reader that takes the first turn and its call, then goes
+                # away; only then does the model answer its second turn.
+                first = json.loads(process.stdout.readline())
+                call = json.loads(process.stdout.readline())
+                process.stdout.close()
+                _Held.released.set()
+                exit_code = process.wait(timeout=60)
+            finally:
+                process.kill()
+                process.wait()
+    finally:
+        server.shutdown()
+        server.server_close()
+
+    assert (first["kind"], call["kind"]) == ("turn", "call")
+    assert exit_code == 94
+    # Nothing more is written once stdout is gone: no traceback, and no
+    # "Exception ignored ... BrokenPipeError" at shutdown either. stderr holds
+    # only the stream's JSON warning lines, such as UNTRUSTED_CONTENT.
+    err = stderr_path.read_text(encoding="utf-8")
+    assert "BrokenPipe" not in err
+    assert "Traceback" not in err
+    assert all("code" in json.loads(line) for line in err.splitlines())
+    [recorded] = (repository / "investigations").glob("*.json")
+    spec = json.loads(recorded.read_text(encoding="utf-8"))["spec"]
+    assert spec["status"] == "stopped"
+    assert spec["turns"] == 2
+    assert spec["answer"] == ""
+    assert "finding" not in spec
+    [step] = spec["steps"]
+    assert step["outcome"] == "ran"
+    assert list((repository / "decisions").glob(f"{step['decision']}*"))