From 3ff7f00e949da4bfe850672c333d505ccbd4b14e Mon Sep 17 00:00:00 2001 From: Romamo Date: Fri, 9 Oct 2026 11:08:47 +0300 Subject: [PATCH 1/2] Keep each model turn's raw trace in the stream and the record A turn now carries the model's reasoning (Nemotron's reasoning_content), the tool calls as written, the response id and x-request-id, its time and its prompt, completion and reasoning tokens. The transport returns a ModelReply so the request id survives. The record's spec.trace is optional, so records written before stay valid. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01APTdZZw8Jw9PQF4TD55wMv --- CHANGELOG.md | 7 + .../v1/agent-investigation.schema.json | 38 ++++++ sdk/src/cloudfall/app.py | 17 ++- sdk/src/cloudfall/investigation.py | 123 ++++++++++++++++-- sdk/tests/test_investigation.py | 41 +++++- 5 files changed, 209 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index df046f8..c6fb8be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,13 @@ Notable changes to Cloudfall. The format follows ### Added +- Every turn of `agent investigate` keeps its raw trace, in the stream's + `turn` event (`trace`) and in the investigation record (`spec.trace`): the + model's reasoning when it returns one (Nemotron's `reasoning_content`), + each tool call exactly as written (name and raw arguments), the + endpoint's response id and `x-request-id`, the milliseconds the turn took, + and its prompt, completion and reasoning tokens. The record can now say + what the model thought when it proposed a change - `agent investigate` streams: one bare JSON line with `_seq` per model turn (its words and how many tools it called), one per tool call (the step and what the tool handed back, `effect: created` when it wrote a decision), diff --git a/config/schemas/v1/agent-investigation.schema.json b/config/schemas/v1/agent-investigation.schema.json index 133da8d..f691a7b 100644 --- a/config/schemas/v1/agent-investigation.schema.json +++ b/config/schemas/v1/agent-investigation.schema.json @@ -117,6 +117,44 @@ "answer": { "type": "string" }, + "trace": { + "type": "array", + "items": { + "type": "object", + "properties": { + "turn": {"type": "integer", "minimum": 1}, + "text": {"type": "string"}, + "reasoning": {"type": "string"}, + "toolCalls": { + "type": "array", + "items": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "arguments": {"type": "string"} + }, + "required": ["name", "arguments"], + "additionalProperties": false + } + }, + "responseId": {"type": "string", "maxLength": 256}, + "requestId": {"type": "string", "maxLength": 256}, + "elapsedMs": {"type": "integer", "minimum": 0}, + "tokens": { + "type": "object", + "properties": { + "prompt": {"type": "integer", "minimum": 0}, + "completion": {"type": "integer", "minimum": 0}, + "reasoning": {"type": "integer", "minimum": 0} + }, + "required": ["prompt", "completion", "reasoning"], + "additionalProperties": false + } + }, + "required": ["turn", "text", "reasoning", "toolCalls", "elapsedMs", "tokens"], + "additionalProperties": false + } + }, "replay": { "type": "string", "minLength": 1, diff --git a/sdk/src/cloudfall/app.py b/sdk/src/cloudfall/app.py index 70d2b9b..1c5f070 100644 --- a/sdk/src/cloudfall/app.py +++ b/sdk/src/cloudfall/app.py @@ -140,6 +140,7 @@ InvestigationStatus, InvestigationStore, ModelEndpoint, + ModelReply, ModelUnavailableError, Progress, StepOutcome, @@ -2231,9 +2232,10 @@ def endpoint(self) -> ModelEndpoint: class InvestigationEvent: """``agent investigate``: one event of the stream. - ``turn`` is the model's words and how many tools it called; ``call`` is one - tool call and what the tool handed back; the last event, ``investigation``, - is the record as it was saved. + ``turn`` is one model turn as the endpoint answered it (``trace``: its + reasoning, the tool calls as written, the response and request ids, the + time and the tokens); ``call`` is one tool call and what the tool handed + back; the last event, ``investigation``, is the record as it was saved. """ effect: str @@ -2241,6 +2243,7 @@ class InvestigationEvent: turn: int | None = None text: str | None = None calls: int | None = None + trace: dict[str, object] = Out(default_factory=dict, ordered=True) step: dict[str, object] = Out(default_factory=dict, ordered=True) result: dict[str, object] = Out(default_factory=dict, ordered=True) status: str | None = None @@ -2345,6 +2348,7 @@ def _progress_event(progress: Progress) -> InvestigationEvent: turn=progress.turn, text=progress.text, calls=progress.calls, + trace=progress.as_dict(), ) wrote = ( progress.step.decision is not None @@ -2444,7 +2448,7 @@ def _tool_result( def _chat_transport(ctx: Ctx, api_key: str) -> ChatTransport: """Send one chat request through ``ctx.http``.""" - def send(url: str, body: Mapping[str, object]) -> Mapping[str, object]: + def send(url: str, body: Mapping[str, object]) -> ModelReply: try: response = ctx.http.post( url, @@ -2473,7 +2477,10 @@ def send(url: str, body: Mapping[str, object]) -> Mapping[str, object]: if not isinstance(answer, Mapping): message = "the model endpoint answered JSON that is not an object" raise InvestigationError(ERROR_MODEL_ANSWER_MALFORMED, message) - return cast("Mapping[str, object]", answer) + return ModelReply( + body=cast("Mapping[str, object]", answer), + request_id=response.headers.get("x-request-id"), + ) return send diff --git a/sdk/src/cloudfall/investigation.py b/sdk/src/cloudfall/investigation.py index 6c19f25..ebdfa59 100644 --- a/sdk/src/cloudfall/investigation.py +++ b/sdk/src/cloudfall/investigation.py @@ -16,6 +16,7 @@ import json import re +import time from collections.abc import Callable, Generator, Mapping, Sequence from dataclasses import dataclass, replace from datetime import UTC, datetime @@ -164,11 +165,18 @@ class AgentTool: call: Callable[[Mapping[str, object]], ToolResult] -ChatTransport = Callable[[str, Mapping[str, object]], Mapping[str, object]] -"""Send one chat request body to a URL and return the answer's JSON. +@dataclass(frozen=True, slots=True) +class ModelReply: + """One answer of the endpoint: its JSON, and the request id it was served under.""" -A failed request raises ``InvestigationError``. -""" + body: Mapping[str, object] + request_id: str | None = None + + +ChatTransport = Callable[[str, Mapping[str, object]], ModelReply] +"""Send one chat request body to a URL and return the endpoint's reply. + +A failed request raises ``InvestigationError``.""" @dataclass(frozen=True, slots=True) @@ -247,6 +255,8 @@ class Investigation: finding: Finding | None replay: str | None = None """The recording played back instead of the hosts, when there was one.""" + trace: tuple[Turn, ...] = () + """Every turn as the endpoint answered it, reasoning included.""" def as_document(self) -> dict[str, object]: """Serialize the investigation as a schema-valid record.""" @@ -268,6 +278,8 @@ def as_document(self) -> dict[str, object]: spec["finding"] = self.finding.as_dict() if self.replay is not None: spec["replay"] = self.replay + if self.trace: + spec["trace"] = [turn.as_dict() for turn in self.trace] return { "apiVersion": "cloudfall/v1", "kind": "AgentInvestigation", @@ -322,13 +334,62 @@ def save(self, investigation: Investigation) -> Investigation: raise InvestigationError(ERROR_INVESTIGATION_EXISTS, message) +@dataclass(frozen=True, slots=True) +class ToolCallText: + """One tool call exactly as the model wrote it: the name and the raw arguments.""" + + name: str + arguments: str + + def as_dict(self) -> dict[str, object]: + """Serialize the call for the record.""" + return {"name": self.name, "arguments": self.arguments} + + @dataclass(frozen=True, slots=True) class Turn: - """The model answered one turn: its words, and how many tools it called.""" + """One turn of the model, as the endpoint answered it. + + Its words, its reasoning when the model returns it (Nemotron does, as + ``reasoning_content``), the tool calls as written, the ids the endpoint + served it under, how long it took and the tokens it cost. + """ turn: int text: str - calls: int + reasoning: str + tool_calls: tuple[ToolCallText, ...] + response_id: str | None + request_id: str | None + elapsed_ms: int + prompt_tokens: int + completion_tokens: int + reasoning_tokens: int + + @property + def calls(self) -> int: + """How many tools the model called this turn.""" + return len(self.tool_calls) + + def as_dict(self) -> dict[str, object]: + """Serialize the turn for the record's trace.""" + document: dict[str, object] = { + "turn": self.turn, + "text": self.text, + "reasoning": self.reasoning, + "toolCalls": [call.as_dict() for call in self.tool_calls], + "elapsedMs": self.elapsed_ms, + "tokens": { + "prompt": self.prompt_tokens, + "completion": self.completion_tokens, + "reasoning": self.reasoning_tokens, + }, + } + if self.response_id is not None: + document["responseId"] = self.response_id + if self.request_id is not None: + document["requestId"] = self.request_id + return document @dataclass(frozen=True, slots=True) @@ -392,6 +453,7 @@ def investigation_events( # noqa: PLR0913 - as investigate. steps: list[Step] = [] seen: dict[str, ToolResult] = {} usage = Usage() + trace: list[Turn] = [] answer = "" status = InvestigationStatus.TURN_LIMIT turns = 0 @@ -412,11 +474,13 @@ def record(ending: InvestigationStatus, finding: Finding | None) -> Investigatio answer=answer, finding=finding, replay=replay, + trace=tuple(trace), ) try: while turns < max_turns: turns += 1 + asked = time.monotonic() reply = transport( endpoint.chat_url, { @@ -426,14 +490,17 @@ def record(ending: InvestigationStatus, finding: Finding | None) -> Investigatio "temperature": 0.2, }, ) - usage = usage.plus(reply) - message = _first_message(reply) + elapsed_ms = int((time.monotonic() - asked) * 1000) + usage = usage.plus(reply.body) + message = _first_message(reply.body) messages.append(message) calls = cast( "Sequence[Mapping[str, object]]", message.get("tool_calls") or () ) text = str(message.get("content") or "").strip() - yield Turn(turn=turns, text=text, calls=len(calls)) + turn = _turn(turns, reply, message, calls, elapsed_ms) + trace.append(turn) + yield turn if not calls: answer = text status = InvestigationStatus.UNSTRUCTURED @@ -482,6 +549,44 @@ def _tool_definition(tool: AgentTool) -> dict[str, object]: } +def _turn( + number: int, + reply: ModelReply, + message: Mapping[str, object], + calls: Sequence[Mapping[str, object]], + elapsed_ms: int, +) -> Turn: + """Read one turn's raw trace out of the endpoint's answer.""" + usage = reply.body.get("usage") + usage = usage if isinstance(usage, Mapping) else {} + details = usage.get("completion_tokens_details") + details = details if isinstance(details, Mapping) else {} + reasoning = message.get("reasoning_content") or message.get("reasoning") or "" + response_id = reply.body.get("id") + return Turn( + turn=number, + text=str(message.get("content") or "").strip(), + reasoning=str(reasoning).strip(), + tool_calls=tuple(_call_text(call) for call in calls), + response_id=response_id if isinstance(response_id, str) else None, + request_id=reply.request_id, + elapsed_ms=elapsed_ms, + prompt_tokens=_count(usage, "prompt_tokens"), + completion_tokens=_count(usage, "completion_tokens"), + reasoning_tokens=_count(details, "reasoning_tokens"), + ) + + +def _call_text(call: Mapping[str, object]) -> ToolCallText: + function = call.get("function") + function = function if isinstance(function, Mapping) else {} + arguments = function.get("arguments", "") + return ToolCallText( + name=str(function.get("name", "")), + arguments=arguments if isinstance(arguments, str) else json.dumps(arguments), + ) + + def _first_message(reply: Mapping[str, object]) -> dict[str, object]: choices = reply.get("choices") if not isinstance(choices, list) or not choices: diff --git a/sdk/tests/test_investigation.py b/sdk/tests/test_investigation.py index 1f48fbc..233bb43 100644 --- a/sdk/tests/test_investigation.py +++ b/sdk/tests/test_investigation.py @@ -17,6 +17,7 @@ InvestigationStatus, InvestigationStore, ModelEndpoint, + ModelReply, ModelUnavailableError, StepOutcome, ToolResult, @@ -74,10 +75,10 @@ def __init__(self, *replies: Mapping[str, object]) -> None: self.replies = list(replies) self.requests: list[Mapping[str, object]] = [] - def __call__(self, url: str, body: Mapping[str, object]) -> Mapping[str, object]: + def __call__(self, url: str, body: Mapping[str, object]) -> ModelReply: assert url == "https://models.example/v1/chat/completions" self.requests.append(json.loads(json.dumps(body))) - return self.replies.pop(0) + return ModelReply(body=self.replies.pop(0), request_id="req-test") class _Tool: @@ -270,7 +271,7 @@ def test_a_malformed_model_answer_is_refused() -> None: 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]: + def __call__(self, url: str, body: Mapping[str, object]) -> ModelReply: if not self.replies: message = "the model endpoint answered 500: overloaded" raise InvestigationError(ERROR_MODEL_REFUSED, message) @@ -828,3 +829,37 @@ def test_an_incomplete_investigation_ends_with_the_record_then_93( # The error line holds no data: the record is the event before it. assert "data.investigation" not in error["suggestion"] assert "kind investigation" in error["suggestion"] + + +def test_each_turn_keeps_the_raw_trace() -> None: + report, rotate = _tools() + thinking = _call("disk-report", {}) + thinking["id"] = "chatcmpl-1" + first = thinking["choices"] + assert isinstance(first, list) + first[0]["message"]["reasoning_content"] = "The disk may be full; read it first." + thinking["usage"] = { + "prompt_tokens": 100, + "completion_tokens": 40, + "completion_tokens_details": {"reasoning_tokens": 30}, + } + script = _Script( + thinking, _say("ROOT CAUSE: disk\nPROPOSED: none\nWHY: read only.") + ) + + done = investigate("alert", _tool_list(report, rotate), ENDPOINT, script) + + turn = done.trace[0] + assert turn.reasoning == "The disk may be full; read it first." + assert turn.tool_calls[0].name == "disk-report" + assert turn.tool_calls[0].arguments == "{}" + assert (turn.response_id, turn.request_id) == ("chatcmpl-1", "req-test") + assert (turn.prompt_tokens, turn.completion_tokens, turn.reasoning_tokens) == ( + 100, + 40, + 30, + ) + recorded = done.as_document()["spec"] + assert isinstance(recorded, dict) + assert recorded["trace"][0]["reasoning"] == turn.reasoning + assert recorded["trace"][1]["toolCalls"] == [] From 22b4f1db5c39c80774546294f8ad717da46bba0b Mon Sep 17 00:00:00 2001 From: "shipmill-romamo[bot]" <338557544+shipmill-romamo[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 11:58:17 +0300 Subject: [PATCH 2/2] Bound the endpoint's response and request ids the trace keeps The schema caps responseId and requestId at 256 characters, but the turn copied both from the endpoint as sent. An over-long id made the record fail schema validation on save with a raw jsonschema error, after the run's decisions were already written, so the investigation was lost. Keep at most 256 characters, as the record already does for tool names. The CLI stream test now also checks the turn event's trace over the real transport, with no x-request-id header and no API key in the stream. Co-Authored-By: Claude Opus 5.5 --- sdk/src/cloudfall/investigation.py | 11 ++++++++-- sdk/tests/test_investigation.py | 34 ++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/sdk/src/cloudfall/investigation.py b/sdk/src/cloudfall/investigation.py index ebdfa59..5edb975 100644 --- a/sdk/src/cloudfall/investigation.py +++ b/sdk/src/cloudfall/investigation.py @@ -40,6 +40,8 @@ """Most characters of host output one tool result hands the model.""" TOOL_NAME_LIMIT = 128 """Most characters of a tool name the record keeps, as its schema allows.""" +REPLY_ID_LIMIT = 256 +"""Most characters of an endpoint's response or request id the trace keeps.""" SAVE_ATTEMPTS = 100 """Most investigations recorded under one second's id, with their suffixes.""" ENDPOINT_LIMIT = 512 @@ -563,13 +565,18 @@ def _turn( details = details if isinstance(details, Mapping) else {} reasoning = message.get("reasoning_content") or message.get("reasoning") or "" response_id = reply.body.get("id") + # The ids are the endpoint's text; the record keeps what its schema takes. return Turn( turn=number, text=str(message.get("content") or "").strip(), reasoning=str(reasoning).strip(), tool_calls=tuple(_call_text(call) for call in calls), - response_id=response_id if isinstance(response_id, str) else None, - request_id=reply.request_id, + response_id=( + response_id[:REPLY_ID_LIMIT] if isinstance(response_id, str) else None + ), + request_id=( + None if reply.request_id is None else reply.request_id[:REPLY_ID_LIMIT] + ), elapsed_ms=elapsed_ms, prompt_tokens=_count(usage, "prompt_tokens"), completion_tokens=_count(usage, "completion_tokens"), diff --git a/sdk/tests/test_investigation.py b/sdk/tests/test_investigation.py index 233bb43..e0b4786 100644 --- a/sdk/tests/test_investigation.py +++ b/sdk/tests/test_investigation.py @@ -803,6 +803,11 @@ def test_the_cli_streams_each_turn_and_call_then_the_record( assert events[1]["effect"] == "created" assert "/var/log/shop holds 36G" in events[1]["result"]["output"] assert events[2]["text"].startswith("ROOT CAUSE") + # A turn carries its raw trace; this endpoint sends no x-request-id. + trace = events[0]["trace"] + assert trace["toolCalls"][0]["name"] == "operation_facts" + assert "requestId" not in trace + assert "sk-test" not in json.dumps(events) # The decision the call wrote, and the saved record. assert stream["end"]["_summary"] is True assert stream["end"]["effects"] == {"created": 2, "noop": 2} @@ -863,3 +868,32 @@ def test_each_turn_keeps_the_raw_trace() -> None: assert isinstance(recorded, dict) assert recorded["trace"][0]["reasoning"] == turn.reasoning assert recorded["trace"][1]["toolCalls"] == [] + + +class _LongIds(_Script): + """An endpoint that serves its answers under over-long ids.""" + + def __call__(self, url: str, body: Mapping[str, object]) -> ModelReply: + reply = super().__call__(url, body) + return ModelReply(body=reply.body, request_id="r" * 300) + + +def test_an_over_long_endpoint_id_still_leaves_a_valid_record( + tmp_path: Path, +) -> None: + report, rotate = _tools() + answer = _say("ROOT CAUSE: disk\nPROPOSED: none\nWHY: read only.") + answer["id"] = "chatcmpl-" + "x" * 300 + script = _LongIds(answer) + done = investigate("alert", _tool_list(report, rotate), ENDPOINT, script) + + saved = InvestigationStore( + directory=tmp_path / "investigations", + catalog=SchemaCatalog(default_schema_directory()), + ).save(done) + + turn = saved.trace[0] + assert turn.response_id is not None + assert turn.request_id is not None + assert turn.response_id.startswith("chatcmpl-") + assert (len(turn.response_id), len(turn.request_id)) == (256, 256)