diff --git a/src/plyngent/tools/truncate_token.py b/src/plyngent/tools/truncate_token.py new file mode 100644 index 0000000..4c67b25 --- /dev/null +++ b/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 diff --git a/tests/test_tools/test_truncate_token.py b/tests/test_tools/test_truncate_token.py new file mode 100644 index 0000000..3a3b215 --- /dev/null +++ b/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