core/tools: unified truncate marker with memory remainder store

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
+Every truncation site (the agent loop's generic tool-result cap,
-cut output short; the model passes the token to ``get_truncated`` to read the
+``run_command`` / ``run_command_batch`` per-stream caps, ``read_file``,
-next chunk without re-requesting the whole source or raising limits.
+``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
+    ``offset``/``limit`` are in characters (http bodies, memory remainders) or
-    line-based via ``read_file``); ``get_truncated`` interprets by ``kind``.
+    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:
+def truncation_marker(max_chars: int, omitted: int, token: TruncateToken) -> str:
-    return (
+    """The single truncation hint: ``[Truncated (N chars max.; M omitted). truncate_token=...]``."""
-        f"\n...[truncated {omitted} chars; more via get_truncated token] "
+    return f"\n[Truncated ({max_chars} chars max.; {omitted} omitted). truncate_token={encode_truncate_token(token)}]"
-        f"[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``
+    Returns ``(bounded_text, next_token)``. A token is embedded whenever more
-    is the full source length. A token is embedded whenever more content remains
+    source content remains after the returned chunk, so ``get_truncated``
-    after the returned chunk, so ``get_truncated`` chains through truncations
+    chains through truncations; each chunk may itself carry a fresh token. The
-    and each chunk may itself carry a fresh token. The token advances by the
+    token advances by the length actually returned (no content is skipped), and
-    length actually returned (no content is skipped), and the marker is budgeted
+    the marker is budgeted so the total returned length stays near
-    so the total returned length stays near ``max_chars`` (the agent loop's own
+    ``max_chars`` (the agent loop's own cap never re-cuts it and never strips
-    cap never re-cuts it and never strips the token).
+    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 = (
+    return text + truncation_marker(max_chars, omitted, token), token
-        f"\n...[more content available; use get_truncated with the token] "
+
-        f"[TRUNCATE_TOKEN: {encode_truncate_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

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 "[Truncated (200 chars max.; 300 omitted). truncate_token=" in out
-    assert f"[TRUNCATE_TOKEN: {encode_truncate_token(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