-
Notifications
You must be signed in to change notification settings - Fork 589
[feat] Last-message-only turns: runner rebuilds history from records #5486
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: fix/transcript-hygiene
Are you sure you want to change the base?
Changes from all commits
af87b05
23ab026
2ad666a
602838b
5a2f595
882eebc
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -12,6 +12,7 @@ | |
|
|
||
| from oss.src.core.sessions.records.dtos import SessionRecordEvent | ||
| from oss.src.dbs.redis.shared.engine import get_streams_engine | ||
| from oss.src.utils.env import env | ||
| from oss.src.utils.logging import get_module_logger | ||
|
|
||
| log = get_module_logger(__name__) | ||
|
|
@@ -21,13 +22,61 @@ | |
| # Truncate attributes at ingest to avoid storing unbounded record bodies. | ||
| MAX_ATTRIBUTES_BYTES = 64 * 1024 # 64 KB per record | ||
|
|
||
| _TRUNCATION_MARKER = "…[truncated]" | ||
|
|
||
|
|
||
| def _orjson_default(obj): | ||
| if AsyncpgUUID is not None and isinstance(obj, AsyncpgUUID): | ||
| return str(obj) | ||
| raise TypeError(f"Type is not JSON serializable: {type(obj)}") | ||
|
|
||
|
|
||
| def _truncate_attributes(attributes, budget: int, original_bytes: int): | ||
| """Shrink an oversized record body to fit `budget` while PRESERVING structure: small fields | ||
| (``type``/``id``/``name``…) stay intact and only the largest string values are trimmed, each | ||
| marked, so server-side history reconstruction still gets the event shape + partial content. | ||
| Falls back to a minimal discriminator-only shape when non-string bloat can't be trimmed.""" | ||
| if original_bytes <= budget: | ||
| return attributes | ||
| if not isinstance(attributes, dict): | ||
| return {"_truncated": True, "_original_bytes": original_bytes} | ||
|
|
||
| result = dict(attributes) | ||
| trimmed: list[str] = [] | ||
| # Trim the largest string field repeatedly until the serialized body fits (or none remain). | ||
| for _ in range(len(attributes)): | ||
| size = len(dumps(result, default=_orjson_default)) | ||
| if size <= budget: | ||
| break | ||
| str_fields = [ | ||
| (k, v) | ||
| for k, v in result.items() | ||
| if isinstance(v, str) and not k.startswith("_") | ||
| ] | ||
| if not str_fields: | ||
| break | ||
| key, value = max(str_fields, key=lambda kv: len(kv[1])) | ||
| overflow = size - budget | ||
| keep = max( | ||
| 0, len(value) - overflow - 128 | ||
| ) # margin for the marker + json overhead | ||
| result[key] = value[:keep] + _TRUNCATION_MARKER | ||
| if key not in trimmed: | ||
| trimmed.append(key) | ||
|
|
||
| if len(dumps(result, default=_orjson_default)) > budget: | ||
| # Non-string bloat remains → keep only the discriminator fields reconstruction needs. | ||
| return { | ||
| "type": attributes.get("type"), | ||
| "id": attributes.get("id"), | ||
| "_truncated": True, | ||
| "_original_bytes": original_bytes, | ||
| } | ||
|
|
||
| result["_truncated"] = {"fields": trimmed, "original_bytes": original_bytes} | ||
| return result | ||
|
|
||
|
|
||
| def _get_redis(): | ||
| engine = get_streams_engine() | ||
| return engine.get_redis() if engine else None | ||
|
|
@@ -72,8 +121,19 @@ async def publish_record( | |
| session_id=str(record_event.session_id), | ||
| original_bytes=len(raw_attributes), | ||
| ) | ||
| # Smart truncation keeps the event shape + partial content so records stay | ||
| # reconstructable; legacy path drops the whole body. Flag-gated (additive). | ||
| new_attributes = ( | ||
| _truncate_attributes( | ||
| record_event.attributes, | ||
| MAX_ATTRIBUTES_BYTES, | ||
| len(raw_attributes), | ||
| ) | ||
| if env.agenta.sessions.records.smart_truncation | ||
| else {"_truncated": True} | ||
| ) | ||
| truncated_event = record_event.model_copy( | ||
| update={"attributes": {"_truncated": True}} | ||
| update={"attributes": new_attributes} | ||
|
Comment on lines
+124
to
+136
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win 🧩 Analysis chain🏁 Script executed: #!/bin/bash
# Verify the parent Agenta config class exposes a `sessions` field wired to SessionsConfig.
rg -n -B5 -A5 'sessions\s*:\s*SessionsConfig' api/oss/src/utils/env.py
rg -n 'class AgentaConfig' -A 30 api/oss/src/utils/env.pyRepository: Agenta-AI/agenta Length of output: 1516 🏁 Script executed: #!/bin/bash
set -euo pipefail
echo "== env.py config definitions/references =="
rg -n 'SessionsConfig|SessionsRecordsConfig|utora|agenta|AGENTA_SESSIONS|SESSIONS' api/oss/src/utils/env.py
echo
echo "== relevant env.py sections =="
sed -n '1,120p' api/oss/src/utils/env.py
sed -n '400,560p' api/oss/src/utils/env.py
echo
echo "== all project references to SessionsConfig/SessionsRecordsConfig =="
rg -n 'SessionsConfig|SessionsRecordsConfig|agenta\.sessions|AGENTA_SESSIONS' api/oss/src || true
echo
echo "== streaming.py usage context =="
sed -n '1,180p' api/oss/src/core/sessions/records/streaming.pyRepository: Agenta-AI/agenta Length of output: 18458 Wire
|
||
| ) | ||
|
|
||
| message = { | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,73 @@ | ||
| """Smart truncation of oversized record bodies (`_truncate_attributes`). | ||
|
|
||
| The legacy path replaces an over-cap body with `{"_truncated": True}`, losing the event's | ||
| type/id and all content. Smart truncation preserves the event shape + partial content so the | ||
| record log stays reconstructable server-side. These pin that contract. | ||
| """ | ||
|
|
||
| from orjson import dumps | ||
|
|
||
| from oss.src.core.sessions.records.streaming import ( | ||
| MAX_ATTRIBUTES_BYTES, | ||
| _TRUNCATION_MARKER, | ||
| _truncate_attributes, | ||
| ) | ||
|
|
||
|
|
||
| def _size(obj) -> int: | ||
| return len(dumps(obj)) | ||
|
|
||
|
|
||
| def test_under_budget_returns_unchanged(): | ||
| attrs = {"type": "message", "text": "hi"} | ||
| out = _truncate_attributes(attrs, MAX_ATTRIBUTES_BYTES, _size(attrs)) | ||
| assert out is attrs # untouched, no _truncated marker | ||
|
|
||
|
|
||
| def test_large_string_field_is_trimmed_but_structure_preserved(): | ||
| big = "x" * (MAX_ATTRIBUTES_BYTES * 2) | ||
| attrs = {"type": "tool_result", "id": "call-1", "output": big} | ||
| out = _truncate_attributes(attrs, MAX_ATTRIBUTES_BYTES, _size(attrs)) | ||
|
|
||
| # Discriminator fields survive (unlike the legacy whole-body drop). | ||
| assert out["type"] == "tool_result" | ||
| assert out["id"] == "call-1" | ||
| # The big field is trimmed + marked, and the whole body now fits the cap. | ||
| assert out["output"].endswith(_TRUNCATION_MARKER) | ||
| assert len(out["output"]) < len(big) | ||
| assert _size(out) <= MAX_ATTRIBUTES_BYTES | ||
| # Metadata records what was trimmed. | ||
| assert out["_truncated"]["fields"] == ["output"] | ||
| assert out["_truncated"]["original_bytes"] == _size(attrs) | ||
|
|
||
|
|
||
| def test_trims_the_largest_of_several_string_fields(): | ||
| attrs = { | ||
| "type": "message", | ||
| "small": "ok", | ||
| "text": "y" * (MAX_ATTRIBUTES_BYTES * 2), | ||
| } | ||
| out = _truncate_attributes(attrs, MAX_ATTRIBUTES_BYTES, _size(attrs)) | ||
| assert out["small"] == "ok" # small field untouched | ||
| assert out["text"].endswith(_TRUNCATION_MARKER) | ||
| assert _size(out) <= MAX_ATTRIBUTES_BYTES | ||
|
|
||
|
|
||
| def test_non_string_bloat_falls_back_to_discriminator_only(): | ||
| # A huge nested structure with no single big string leaf can't be string-trimmed. | ||
| attrs = { | ||
| "type": "tool_call", | ||
| "id": "call-9", | ||
| "input": {str(i): i for i in range(MAX_ATTRIBUTES_BYTES)}, | ||
| } | ||
| out = _truncate_attributes(attrs, MAX_ATTRIBUTES_BYTES, _size(attrs)) | ||
| assert out["type"] == "tool_call" | ||
| assert out["id"] == "call-9" | ||
| assert out["_truncated"] is True | ||
| assert _size(out) <= MAX_ATTRIBUTES_BYTES | ||
|
|
||
|
|
||
| def test_non_dict_attributes_fall_back(): | ||
| huge = "z" * (MAX_ATTRIBUTES_BYTES * 2) | ||
| out = _truncate_attributes(huge, MAX_ATTRIBUTES_BYTES, _size(huge)) | ||
| assert out == {"_truncated": True, "_original_bytes": _size(huge)} |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,55 @@ | ||
| # Last-message-only sends + server-side history reconstruction | ||
|
|
||
| **Branch:** `feat/sessions-last-message-only` (off `feat/sessions-storage-rework` / #5436) | ||
| **Motivation:** AGE-3970 (trace-drawer OOM on long turns) + device-independent continuity. | ||
| Shrink request/trace payloads by sending only the newest user message per turn and | ||
| reconstructing prior conversation server-side from durable session **records**. | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Do we still use caching to reconstruct the prior conversation using the local DB, or are we using only the service side records?
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Only the service-side records. There is no local database read on this path. Concretely: There is a separate mechanism that can look like caching, and it is worth keeping distinct. When a sandbox is still parked from the previous turn, the harness process inside it already holds its own conversation context, so the runner sends only the new text instead of replaying anything. That is the warm path, and it never touches records. The two interact badly today, which is the part I would not have predicted from the design. With last-message-only on, the warm path is rejected on every turn, because the warm check fingerprints the inbound history and the inbound history is now a single message. So in practice you never stay warm, and every turn falls through to the records path. Evidence is in finding 2 of the QA comment.
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. My question was related to the playground and not the runner.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Only server-side records. The runner fetches the session's durable record log ( |
||
|
|
||
| ## Decisions (locked) | ||
| - **Reconstruction lives in the runner (TS)** — at the `buildTurnText`/`priorMessages` seam, | ||
| reusing the runner's records client + session plumbing; a TS↔TS port of the FE | ||
| `transcriptToMessages.ts` folding. (Shared infra with JP/Mahmoud — keep additive/flagged.) | ||
| - **Harden records durability FIRST**, before the FE trusts reconstruction — so reconstructed | ||
| history is authoritative, not lossy. | ||
|
|
||
| ## Current architecture (verified FE / SDK / runner / api) | ||
| Two continuity paths already exist in the runner: | ||
| - **Warm** (pool hit / same harness authored last turn): harness native `resumeSession()` | ||
| supplies history; `sendLastMessageOnly` is ALREADY true (`runtime-contracts.ts:168`). | ||
| - **Cold** (evicted / different harness / cross-device): runner rebuilds the full transcript | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is not an accurate description. Just making sure that the design and the implementation take that into account. There is warm, there is cold, there is cold evicted, and I think the issue happens only when cold evicted and we need to rebuild the whole transcript again.
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You are right, and the QA shows the implementation carries the same gap, not just the write-up. The three states as they actually exist in the runner:
Your instinct is right that the interesting case is the evicted one, and it is worse than the design assumes. The Measured on the three-turn journey: turn 2 went from 1570 ms to 4623 ms and turn 3 from 1399 ms to 4252 ms, with three I will rewrite this section of the design doc against the three real states as part of the fixes.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Agreed, the doc conflated them. Corrected in 4c34ce4: records-based reconstruction supplies model context only in the cold-evicted case; warm and cold-not-evicted resume natively. Reconstruction still runs on warm turns, but only to realign the keep-alive fingerprint. |
||
| from inbound `request.messages` (`transcript.ts:buildTurnText`/`priorMessages`; selected at | ||
| `run-turn.ts:134`). | ||
|
|
||
| The FE always sends full history (`agentRequest.ts:401`) because it can't predict warm/cold. | ||
| Records are **write-only** from the run's perspective: the runner persists every event | ||
| per-turn (`persist.ts` → `POST /sessions/records/ingest`) but never reads them back for model | ||
| context. Records ARE sufficient to reconstruct full ordered history (roles, text, reasoning, | ||
| tool_call/tool_result, approvals) — each carries the whole ACP `AgentEvent`. | ||
|
|
||
| ## Hard constraint: HITL cold-replay | ||
| Resume is message-driven, not interaction-driven. The runner binds tool_result→tool_call by | ||
| scanning `request.messages` (`responder.ts:413,470`) and scans full history for approval | ||
| verdicts (`responder.ts:358`). Durable `session_interactions` rows exist but aren't read to | ||
| rebuild history. ⇒ Dropping history on a cold resume breaks approvals/client-tools/elicitation | ||
| until the runner reconstructs matched tool_call/tool_result pairs from records. **The BE | ||
| reconstructor + durable records are the prerequisite for the FE change.** | ||
|
|
||
| ## Phases | ||
| - **Phase 1 — records durability** (enabler): make `message` + `tool_result` (and `tool_call`, | ||
| `interaction_request`) record persistence reliable end-to-end — reduce the runner's | ||
| fire-and-forget drop (ack/retry-to-DLQ) and fix the 64 KB attribute cap | ||
| (`records/streaming.py`) for tool bodies (raise + spill). Order records by | ||
| `session_turns.turn_index` then `record_index`. | ||
| - **Phase 2 — runner reconstructor**: new `services/runner/src/sessions/` module folds | ||
| records→`ChatMessage[]` (inverse of `buildPersistingEmitter`); generalize | ||
| `sendLastMessageOnly` to fire whenever `session_id` yields reconstructable history; feed | ||
| `buildTurnText`/`priorMessages` + otel `run.start`. Behind a flag; FE unchanged. | ||
| - **Phase 3 — FE last-message-only**: `buildAgentRequest` sends only the trailing user message | ||
| when the session is server-known with records; falls back to full history otherwise. Trace/ | ||
| request payloads shrink → AGE-3970 mitigated. | ||
|
|
||
| ## Open risks tracked | ||
| - Runner is ephemeral: guaranteeing delivery may mean blocking the run on a persist ack | ||
| (latency) or a durable outbox — Phase 1 decides the mechanism. | ||
| - Interaction-row verdict as a backstop source (if a settled answer must survive without the | ||
| last message) — evaluate in Phase 2. | ||
| Original file line number | Diff line number | Diff line change | ||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,50 @@ | ||||||||||||||||||
| /** | ||||||||||||||||||
| * Seam that lets the runner rebuild prior conversation from the durable record log instead of | ||||||||||||||||||
| * trusting a full inbound history — the server side of "client sends only the last message". | ||||||||||||||||||
| * | ||||||||||||||||||
| * Flag-gated (`AGENTA_SESSIONS_RECONSTRUCT`) and a strict no-op until BOTH the flag is on AND the | ||||||||||||||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Did we do a live queue here where we compared the reconstructed session that you get from the records in the runner (like this logic) to the stuff that is sent by the frontend in different scenarios, like:
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please spin off this in a work tree, and let's use the Claude Code subscription with Hi-Co and do a QA using this. FDPR obviously doesn't say that it did it before.
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. No, that comparison had not been done. It has now, and you were right that it is the validation that matters: it found four defects that the pass/fail suite cannot see. Method: one dev stack at this PR's head with all four flags on, then the agent release gate run twice against it, changing only how the client sends history. Run A sends the full conversation, which is today's frontend. Run B sends only the trailing user message, mirroring Both runs are all-PASS on all nine journeys, which is the whole point. The differences only show in the message arrays and the runner log. What it found:
One scenario from your list is still untested: a fresh turn after an approval resume, which is where Codex predicted a duplicated tool-call id. The gate's approve journey stops at the resume, so it never exercised it. Full detail and evidence: QA comment.
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done. Worktree at One thing had to be fixed before the QA could run at all: the runner service deliberately has no Results are in the QA comment. Fixes will land on that branch as a PR stacked on this one, so the fix diff stays readable on its own.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes — that differential QA ran (see the QA comment on this PR). It diffed reconstructed vs FE-sent history across tool calls, approvals, MCP, and multi-turn, and surfaced 4 defects a pass/fail gate can't see. Fixes are in the stacked PR #5500. The flag stays until a clean re-run.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done — the QA above ran that way (worktree at this PR's head, Claude subscription / Hi-Co, all four flags on). Findings + fixes: #5500. |
||||||||||||||||||
| * client actually sent a minimal history: when the client still sends the whole conversation | ||||||||||||||||||
| * (`messages.length > 1`), reconstruction is skipped and behaviour is unchanged. Best-effort — any | ||||||||||||||||||
| * miss (no session, no records, fetch failure) leaves the inbound history untouched. | ||||||||||||||||||
| */ | ||||||||||||||||||
|
|
||||||||||||||||||
| import type { AgentRunRequest } from "../../protocol.ts"; | ||||||||||||||||||
| import { fetchSessionRecords } from "../../sessions/records-query.ts"; | ||||||||||||||||||
| import { reconstructMessages } from "../../sessions/reconstruct.ts"; | ||||||||||||||||||
|
|
||||||||||||||||||
| function reconstructEnabled(): boolean { | ||||||||||||||||||
| return ( | ||||||||||||||||||
| String(process.env.AGENTA_SESSIONS_RECONSTRUCT ?? "").toLowerCase() === "true" | ||||||||||||||||||
| ); | ||||||||||||||||||
| } | ||||||||||||||||||
|
|
||||||||||||||||||
| /** | ||||||||||||||||||
| * Returns a request whose `messages` are `[...reconstructed prior turns, ...inbound]` when | ||||||||||||||||||
| * reconstruction applies, else `null` to keep the inbound history as-is. | ||||||||||||||||||
| * | ||||||||||||||||||
| * MUST be called before the current turn's user message is persisted, so the record log holds | ||||||||||||||||||
| * only prior turns (no duplication of the incoming prompt). | ||||||||||||||||||
| */ | ||||||||||||||||||
| export async function reconstructHistoryIfNeeded( | ||||||||||||||||||
| request: AgentRunRequest, | ||||||||||||||||||
| sessionId: string | undefined, | ||||||||||||||||||
| auth: () => string, | ||||||||||||||||||
| log?: (msg: string) => void, | ||||||||||||||||||
| ): Promise<AgentRunRequest | null> { | ||||||||||||||||||
| if (!reconstructEnabled() || !sessionId) return null; | ||||||||||||||||||
| const inbound = request.messages ?? []; | ||||||||||||||||||
| // The client already sent the conversation — nothing to rebuild. | ||||||||||||||||||
| if (inbound.length > 1) return null; | ||||||||||||||||||
|
Comment on lines
+34
to
+37
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win Require exactly one inbound user message before reconstructing. The current Proposed fix const inbound = request.messages ?? [];
-// The client already sent the conversation — nothing to rebuild.
-if (inbound.length > 1) return null;
+// Reconstruct only for a fresh turn containing exactly its trailing user message.
+if (inbound.length !== 1 || inbound[0]?.role !== "user") return null;📝 Committable suggestion
Suggested change
|
||||||||||||||||||
|
|
||||||||||||||||||
| const records = await fetchSessionRecords(sessionId, auth); | ||||||||||||||||||
| if (!records || records.length === 0) return null; | ||||||||||||||||||
|
|
||||||||||||||||||
| const reconstructed = reconstructMessages(records); | ||||||||||||||||||
| if (reconstructed.length === 0) return null; | ||||||||||||||||||
|
|
||||||||||||||||||
| log?.( | ||||||||||||||||||
| `[reconstruct] session=${sessionId} records=${records.length} ` + | ||||||||||||||||||
| `priorMessages=${reconstructed.length} inbound=${inbound.length}`, | ||||||||||||||||||
| ); | ||||||||||||||||||
| return { ...request, messages: [...reconstructed, ...inbound] }; | ||||||||||||||||||
| } | ||||||||||||||||||
Uh oh!
There was an error while loading. Please reload this page.