diff --git a/src/plyngent/tools/truncate_token.py b/src/plyngent/tools/truncate_token.py index 3c1c2d2..be3af0e 100644 --- a/src/plyngent/tools/truncate_token.py +++ b/src/plyngent/tools/truncate_token.py @@ -1,25 +1,43 @@ """Truncate-token cursors for resuming truncated tool results. -``fetch`` and ``read_file`` append a ``[TRUNCATE_TOKEN: ...]`` marker when they -cut output short; the model passes the token to ``get_truncated`` to read the -next chunk without re-requesting the whole source or raising limits. +Every truncation site (the agent loop's generic tool-result cap, +``run_command`` / ``run_command_batch`` per-stream caps, ``read_file``, +``fetch``, and request-time compact shrinks) appends the same compact marker +when output is cut short:: + + [Truncated (32000 chars max.; 500 omitted). truncate_token=...] + +The model passes the token to ``get_truncated`` to read the next chunk without +re-requesting the whole source or raising limits. Tokens chain: each truncated +chunk carries a fresh token, so ``get_truncated`` keeps resuming until the +source is exhausted (then no marker is emitted). + +Resumable sources (``file`` / ``http``) are re-read by ``location`` + +``offset``. Arbitrary tool output (run_command stdout, todo renders, …) has no +resumable source, so its remainder lives in a short-lived in-memory store +(``kind="memory"``) that is forgotten when the agent process exits — the model +must re-run the tool after that. """ from __future__ import annotations import base64 +import uuid from typing import Literal import msgspec -type TruncateKind = Literal["file", "http"] +type TruncateKind = Literal["file", "http", "memory"] class TruncateToken(msgspec.Struct, frozen=True): """Opaque cursor into a truncated source. - ``offset``/``limit`` are in characters (http bodies) or lines (files is - line-based via ``read_file``); ``get_truncated`` interprets by ``kind``. + ``offset``/``limit`` are in characters (http bodies, memory remainders) or + lines (files is line-based via ``read_file``); ``get_truncated`` interprets + by ``kind``. For ``memory``, ``location`` is a key into the in-memory + remainder store and ``offset``/``limit`` index characters of that stored + remainder. """ kind: TruncateKind @@ -28,6 +46,23 @@ class TruncateToken(msgspec.Struct, frozen=True): limit: int +# In-memory remainders for generic (non-resumable) tool output. Process-lifetime +# only: forgotten when the agent exits. Single-user CLI → one process per agent. +_MEMORY_REMAINDERS: dict[str, str] = {} + + +def store_remainder(text: str) -> str: + """Store truncated-away output; return its opaque store key.""" + key = f"mem-{uuid.uuid4().hex[:12]}" + _MEMORY_REMAINDERS[key] = text + return key + + +def get_remainder(key: str) -> str | None: + """Fetch a stored remainder by key; None when expired/unknown.""" + return _MEMORY_REMAINDERS.get(key) + + def encode_truncate_token(token: TruncateToken) -> str: """Base64url-encode a token so it is safe inside tool-result text.""" payload = msgspec.json.encode(token) @@ -46,11 +81,9 @@ def decode_truncate_token(text: str) -> TruncateToken | None: return token -def _truncated_marker(omitted: int, token: TruncateToken) -> str: - return ( - f"\n...[truncated {omitted} chars; more via get_truncated token] " - f"[TRUNCATE_TOKEN: {encode_truncate_token(token)}]" - ) +def truncation_marker(max_chars: int, omitted: int, token: TruncateToken) -> str: + """The single truncation hint: ``[Truncated (N chars max.; M omitted). truncate_token=...]``.""" + return f"\n[Truncated ({max_chars} chars max.; {omitted} omitted). truncate_token={encode_truncate_token(token)}]" def truncate_with_token( @@ -63,34 +96,55 @@ def truncate_with_token( limit: int, total_len: int, ) -> tuple[str, TruncateToken | None]: - """Bound ``text`` to ``max_chars``; return (bounded_text, next_token). + """Bound ``text`` (a slice of the full source starting at ``offset``) to ``max_chars``. - ``text`` is a slice of the full source starting at ``offset``; ``total_len`` - is the full source length. A token is embedded whenever more content remains - after the returned chunk, so ``get_truncated`` chains through truncations - and each chunk may itself carry a fresh token. The token advances by the - length actually returned (no content is skipped), and the marker is budgeted - so the total returned length stays near ``max_chars`` (the agent loop's own - cap never re-cuts it and never strips the token). + Returns ``(bounded_text, next_token)``. A token is embedded whenever more + source content remains after the returned chunk, so ``get_truncated`` + chains through truncations; each chunk may itself carry a fresh token. The + token advances by the length actually returned (no content is skipped), and + the marker is budgeted so the total returned length stays near + ``max_chars`` (the agent loop's own cap never re-cuts it and never strips + the token). ``omitted`` in the marker counts source characters not yet + delivered. When nothing was cut and the source ends, ``text`` is returned + unchanged with no marker. """ if max_chars < 1: return text, None if len(text) > max_chars: - omitted = len(text) - max_chars # First pass sizes the marker with an offset+max_chars token; re-encode # with the real cut so the cursor lands exactly where content stops. first = TruncateToken(kind=kind, location=location, offset=offset + max_chars, limit=limit) - cut = max_chars - len(_truncated_marker(omitted, first)) + omitted = total_len - (offset + max_chars) + cut = max_chars - len(truncation_marker(max_chars, omitted, first)) if cut < 1: - return text[:max_chars] + _truncated_marker(omitted, first), first + return text[:max_chars] + truncation_marker(max_chars, omitted, first), first token = TruncateToken(kind=kind, location=location, offset=offset + cut, limit=limit) - return text[:cut] + _truncated_marker(omitted, token), token + return text[:cut] + truncation_marker(max_chars, omitted, token), token more_after = offset + len(text) < total_len if not more_after: return text, None + omitted = total_len - (offset + len(text)) token = TruncateToken(kind=kind, location=location, offset=offset + len(text), limit=limit) - marker = ( - f"\n...[more content available; use get_truncated with the token] " - f"[TRUNCATE_TOKEN: {encode_truncate_token(token)}]" + return text + truncation_marker(max_chars, omitted, token), token + + +def truncate_generic(text: str, max_chars: int) -> str: + """Cap arbitrary tool output; embed a memory truncate token (chainable). + + The full remainder is kept in the in-memory store so ``get_truncated`` can + resume it. Returns ``text`` unchanged (no marker) when it fits within + ``max_chars``. + """ + if max_chars < 1 or len(text) <= max_chars: + return text + key = store_remainder(text) + bounded, _ = truncate_with_token( + text, + max_chars, + kind="memory", + location=key, + offset=0, + limit=max_chars, + total_len=len(text), ) - return text + marker, token + return bounded diff --git a/tests/test_tools/test_truncate_token.py b/tests/test_tools/test_truncate_token.py index 6c1c5a6..f2cdd9d 100644 --- a/tests/test_tools/test_truncate_token.py +++ b/tests/test_tools/test_truncate_token.py @@ -4,7 +4,10 @@ from plyngent.tools.truncate_token import ( TruncateToken, decode_truncate_token, encode_truncate_token, + get_remainder, + truncate_generic, truncate_with_token, + truncation_marker, ) @@ -27,6 +30,14 @@ def test_decode_rejects_invalid_fields() -> None: assert decode_truncate_token(encode_truncate_token(bad_limit)) is None +def test_marker_format() -> None: + token = TruncateToken(kind="file", location="a.txt", offset=100, limit=32_000) + marker = truncation_marker(32_000, 500, token) + assert marker.startswith("\n[Truncated (32000 chars max.; 500 omitted). truncate_token=") + assert marker.endswith("]") + assert encode_truncate_token(token) in marker + + def test_truncate_with_token_short_text_untouched() -> None: out, token = truncate_with_token( "hello", @@ -39,6 +50,7 @@ def test_truncate_with_token_short_text_untouched() -> None: ) assert out == "hello" assert token is None + assert "Truncated" not in out # no hint when no truncation def test_truncate_with_token_cuts_and_advances() -> None: @@ -53,10 +65,9 @@ def test_truncate_with_token_cuts_and_advances() -> None: total_len=500, ) assert token is not None - content, _, _ = out.partition("\n...[") + content, _, _ = out.partition("\n[Truncated") assert content.startswith("x") - assert "truncated 300 chars" in out - assert f"[TRUNCATE_TOKEN: {encode_truncate_token(token)}]" in out + assert "[Truncated (200 chars max.; 300 omitted). truncate_token=" in out assert token.offset == len(content) # cursor lands exactly where content stops assert len(out) <= 200 # marker fits inside the char budget @@ -75,7 +86,7 @@ def test_truncate_with_token_more_after_nonoverflow() -> None: ) assert token is not None assert token.offset == 30 # advanced by returned length only - assert "more content available" in out + assert "[Truncated (40 chars max.; 170 omitted). truncate_token=" in out assert out.startswith("y" * 30) @@ -92,6 +103,7 @@ def test_truncate_with_token_end_reached_no_token() -> None: ) assert token is None assert out == text + assert "Truncated" not in out # no hint when the source ends def test_truncate_with_token_nonzero_start() -> None: @@ -107,7 +119,32 @@ def test_truncate_with_token_nonzero_start() -> None: ) assert token is not None assert token.kind == "http" - content, _, _ = out.partition("\n...[") + content, _, _ = out.partition("\n[Truncated") assert token.offset == 150 + len(content) - assert "truncated 300 chars" in out + assert "[Truncated (200 chars max.; 350 omitted). truncate_token=" in out assert len(out) <= 200 + + +def test_truncate_generic_short_untouched() -> None: + assert truncate_generic("tiny", 200) == "tiny" + + +def test_truncate_generic_embeds_memory_token() -> None: + text = "a" * 500 + out = truncate_generic(text, 200) + assert len(out) <= 200 + assert "[Truncated (200 chars max." in out + token_str = out.split("truncate_token=", 1)[1].rstrip("]") + token = decode_truncate_token(token_str) + assert token is not None + assert token.kind == "memory" + # Remainder is the full original text; token offset points into it. + remainder = get_remainder(token.location) + assert remainder == text + assert token.offset < len(text) + + +def test_truncate_generic_no_marker_when_fits() -> None: + out = truncate_generic("a" * 50, 200) + assert out == "a" * 50 + assert "Truncated" not in out