core/tools: add truncate token module (chainable cursors)

src/plyngent/tools/truncate_token.py

@@ -0,0 +1,85 @@
+"""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.
+"""
+
+from __future__ import annotations
+
+import base64
+from typing import Literal
+
+import msgspec
+
+type TruncateKind = Literal["file", "http"]
+
+
+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``.
+    """
+
+    kind: TruncateKind
+    location: str
+    offset: int
+    limit: int
+
+
+def encode_truncate_token(token: TruncateToken) -> str:
+    """Base64url-encode a token so it is safe inside tool-result text."""
+    payload = msgspec.json.encode(token)
+    return base64.urlsafe_b64encode(payload).decode("ascii")
+
+
+def decode_truncate_token(text: str) -> TruncateToken | None:
+    """Decode a token; None when malformed (model garbage or tampering)."""
+    try:
+        raw = base64.urlsafe_b64decode(text.encode("ascii"))
+        token = msgspec.json.decode(raw, type=TruncateToken)
+    except ValueError, msgspec.DecodeError:
+        return None
+    if token.offset < 0 or token.limit < 1:
+        return None
+    return token
+
+
+def truncate_with_token(
+    text: str,
+    max_chars: int,
+    *,
+    kind: TruncateKind,
+    location: str,
+    offset: int,
+    limit: int,
+    total_len: int,
+) -> tuple[str, TruncateToken | None]:
+    """Bound ``text`` to ``max_chars``; return (bounded_text, next_token).
+
+    ``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 marker is budgeted so the
+    total returned length stays at or under ``max_chars`` (the agent loop's own
+    cap never re-cuts it and never strips the token).
+    """
+    if max_chars < 1:
+        return text, None
+    overflow = len(text) > max_chars
+    more_after = offset + max_chars < total_len
+    if not overflow and not more_after:
+        return text, None
+    token = TruncateToken(kind=kind, location=location, offset=offset + max_chars, limit=limit)
+    tok_text = encode_truncate_token(token)
+    if overflow:
+        omitted = len(text) - max_chars
+        marker = f"\n...[truncated {omitted} chars; more via get_truncated token] [TRUNCATE_TOKEN: {tok_text}]"
+        cut = max_chars - len(marker)
+        if cut < 1:
+            # Marker longer than the budget: emit anyway so chaining never drops.
+            return text[:max_chars] + marker, token
+        return text[:cut] + marker, token
+    marker = f"\n...[more content available; use get_truncated with the token] [TRUNCATE_TOKEN: {tok_text}]"
+    return text + marker, token

tests/test_tools/test_truncate_token.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from plyngent.tools.truncate_token import (
+    TruncateToken,
+    decode_truncate_token,
+    encode_truncate_token,
+    truncate_with_token,
+)
+
+
+def test_token_roundtrip() -> None:
+    token = TruncateToken(kind="file", location="src/a.py", offset=100, limit=32_000)
+    encoded = encode_truncate_token(token)
+    assert decode_truncate_token(encoded) == token
+
+
+def test_decode_garbage_returns_none() -> None:
+    assert decode_truncate_token("not-a-token") is None
+    assert decode_truncate_token("") is None
+    assert decode_truncate_token("!!!") is None
+
+
+def test_decode_rejects_invalid_fields() -> None:
+    bad = TruncateToken(kind="file", location="x", offset=-1, limit=10)
+    assert decode_truncate_token(encode_truncate_token(bad)) is None
+    bad_limit = TruncateToken(kind="http", location="u", offset=0, limit=0)
+    assert decode_truncate_token(encode_truncate_token(bad_limit)) is None
+
+
+def test_truncate_with_token_short_text_untouched() -> None:
+    out, token = truncate_with_token(
+        "hello",
+        32_000,
+        kind="file",
+        location="a.txt",
+        offset=0,
+        limit=32_000,
+        total_len=5,
+    )
+    assert out == "hello"
+    assert token is None
+
+
+def test_truncate_with_token_cuts_and_advances() -> None:
+    text = "x" * 500
+    out, token = truncate_with_token(
+        text,
+        200,
+        kind="file",
+        location="a.txt",
+        offset=0,
+        limit=200,
+        total_len=500,
+    )
+    assert token is not None
+    assert "truncated 300 chars" in out
+    assert f"[TRUNCATE_TOKEN: {encode_truncate_token(token)}]" in out
+    assert token.offset == 200
+    assert token.limit == 200
+    assert len(out) <= 200  # marker fits inside the char budget
+
+
+def test_truncate_with_token_more_after_nonoverflow() -> None:
+    # Segment fits the chunk, but the source continues: still emit a token.
+    text = "y" * 30
+    out, token = truncate_with_token(
+        text,
+        40,
+        kind="file",
+        location="a.txt",
+        offset=0,
+        limit=40,
+        total_len=200,
+    )
+    assert token is not None
+    assert token.offset == 40
+    assert "more content available" in out
+    assert out.startswith("y" * 30)
+
+
+def test_truncate_with_token_end_reached_no_token() -> None:
+    text = "z" * 30
+    out, token = truncate_with_token(
+        text,
+        40,
+        kind="http",
+        location="https://example.com/x",
+        offset=0,
+        limit=40,
+        total_len=30,
+    )
+    assert token is None
+    assert out == text
+
+
+def test_truncate_with_token_nonzero_start() -> None:
+    text = "y" * 500
+    out, token = truncate_with_token(
+        text,
+        200,
+        kind="http",
+        location="https://example.com/x",
+        offset=150,
+        limit=200,
+        total_len=700,
+    )
+    assert token is not None
+    assert token.kind == "http"
+    assert token.offset == 350  # 150 + 200
+    assert "truncated 300 chars" in out
+    assert len(out) <= 200