mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 07:17:54 +08:00
core/tools: embed truncate tokens in fetch (offset param)
This commit is contained in:
@@ -16,6 +16,7 @@ from plyngent.tools.net.policy import (
|
||||
parse_fetch_url,
|
||||
soft_confirm_reason,
|
||||
)
|
||||
from plyngent.tools.truncate_token import truncate_with_token
|
||||
|
||||
|
||||
def format_fetch_result(
|
||||
@@ -32,13 +33,20 @@ def format_fetch_result(
|
||||
warnings: list[str],
|
||||
body_kind: str,
|
||||
max_chars: int,
|
||||
offset: int = 0,
|
||||
) -> str:
|
||||
text = body_text
|
||||
char_truncated = False
|
||||
if max_chars >= 1 and len(text) > max_chars:
|
||||
omitted = len(text) - max_chars
|
||||
text = text[:max_chars] + f"\n...[truncated {omitted} characters]"
|
||||
char_truncated = True
|
||||
start = max(0, offset)
|
||||
body = body_text[start:]
|
||||
body, token = truncate_with_token(
|
||||
body,
|
||||
max_chars,
|
||||
kind="http",
|
||||
location=final_url,
|
||||
offset=start,
|
||||
limit=max_chars,
|
||||
total_len=len(body_text),
|
||||
)
|
||||
char_truncated = token is not None
|
||||
warn_line = "; ".join(warnings) if warnings else ""
|
||||
parts = [
|
||||
f"status={status}",
|
||||
@@ -54,7 +62,7 @@ def format_fetch_result(
|
||||
if warn_line:
|
||||
parts.append(f"warnings={warn_line}")
|
||||
parts.append("--- body ---")
|
||||
parts.append(text)
|
||||
parts.append(body)
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
@@ -72,6 +80,7 @@ async def fetch(
|
||||
max_redirects: int = DEFAULT_MAX_REDIRECTS,
|
||||
follow_redirects: bool = True,
|
||||
allow_http_downgrade: bool = False,
|
||||
offset: int = 0,
|
||||
) -> str:
|
||||
"""HTTP request (GET/POST/PUT/DELETE); return status, metadata, and truncated body.
|
||||
|
||||
@@ -122,6 +131,7 @@ async def fetch(
|
||||
warnings=result.warnings,
|
||||
body_kind=result.body_kind,
|
||||
max_chars=max_chars,
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ from plyngent.tools.net.policy import (
|
||||
parse_fetch_url,
|
||||
soft_confirm_reason,
|
||||
)
|
||||
from plyngent.tools.truncate_token import decode_truncate_token
|
||||
from tests.test_tools.helpers import call_async
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -237,6 +238,36 @@ async def test_fetch_truncation_and_binary(workspace: Path, http_server: str) ->
|
||||
assert "binary body omitted" in out_bin
|
||||
|
||||
|
||||
async def test_fetch_char_truncation_embeds_token(workspace: Path, http_server: str) -> None:
|
||||
del workspace
|
||||
port = int(http_server.rsplit(":", 1)[1])
|
||||
grant_private_host("127.0.0.1", port)
|
||||
|
||||
out = await call_async(fetch, f"{http_server}/big", max_chars=1000)
|
||||
assert "truncated=true" in out
|
||||
body = out.split("--- body ---", 1)[-1]
|
||||
assert "TRUNCATE_TOKEN:" in body
|
||||
token_line = next(ln for ln in body.splitlines() if "TRUNCATE_TOKEN:" in ln)
|
||||
token_str = token_line.split("[TRUNCATE_TOKEN: ", 1)[1].rstrip("]")
|
||||
token = decode_truncate_token(token_str)
|
||||
assert token is not None
|
||||
assert token.kind == "http"
|
||||
assert token.offset == 1000 # next chunk starts at max_chars
|
||||
assert token.limit == 1000
|
||||
|
||||
|
||||
async def test_fetch_offset_skips_body(workspace: Path, http_server: str) -> None:
|
||||
del workspace
|
||||
port = int(http_server.rsplit(":", 1)[1])
|
||||
grant_private_host("127.0.0.1", port)
|
||||
|
||||
# /big body is 5000 chars; offset 4900 leaves 100 chars, no token needed.
|
||||
out = await call_async(fetch, f"{http_server}/big", max_chars=200, offset=4900)
|
||||
body = out.split("--- body ---", 1)[-1]
|
||||
assert body.strip() == "x" * 100
|
||||
assert "TRUNCATE_TOKEN:" not in out
|
||||
|
||||
|
||||
async def test_fetch_private_denied_without_grant(workspace: Path, http_server: str) -> None:
|
||||
del workspace
|
||||
# No grant, no policy hook → hard deny (even though server is up).
|
||||
|
||||
Reference in New Issue
Block a user