mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-25 08:04:57 +08:00
Compare commits
8 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| f3cefccefc | |||
| 6b8874b26d | |||
| e3e0cf2080 | |||
| f473fed708 | |||
| 6cf8902c57 | |||
| f5a667567d | |||
| 05fe1b4070 | |||
| d74b03923f |
@@ -87,7 +87,8 @@ Module-level `@tool` handlers. Call `set_workspace_root()` before use.
|
|||||||
- **`vcs`**: read-only VCS tools (`vcs_kind` / `vcs_status` / `vcs_diff` / `vcs_log` / `vcs_branch`) via `VcsBackend` protocol; **git** implemented; detectors are pluggable for other systems.
|
- **`vcs`**: read-only VCS tools (`vcs_kind` / `vcs_status` / `vcs_diff` / `vcs_log` / `vcs_branch`) via `VcsBackend` protocol; **git** implemented; detectors are pluggable for other systems.
|
||||||
- **`chat`**: human prompts as tools — `ask_user_line` / `ask_user_choice` / `ask_user_form` (shared `prompting` core).
|
- **`chat`**: human prompts as tools — `ask_user_line` / `ask_user_choice` / `ask_user_form` (shared `prompting` core).
|
||||||
- **`todo`**: LIFO stack of **task groups** (not a queue of tasks) — `todo_push` creates one group of siblings; `todo_pop` removes the whole top group; `todo_update` by id; DFS breakdown push[T1,T2]→push[T1.1…]→pop→push[T2.1…]; stored on session row; open items = unfinished work (`[TODO OPEN WORK]`); all-terminal non-empty = hygiene (`[TODO HYGIENE]`); nag channel via `[agent] todo_nag_strategy` (`developer`/`user` = prose nag; `synthetic_tool` = forged `todo_list` call + **real** `stack.render()` result + ToolCall/Result events; `none`).
|
- **`todo`**: LIFO stack of **task groups** (not a queue of tasks) — `todo_push` creates one group of siblings; `todo_pop` removes the whole top group; `todo_update` by id; DFS breakdown push[T1,T2]→push[T1.1…]→pop→push[T2.1…]; stored on session row; open items = unfinished work (`[TODO OPEN WORK]`); all-terminal non-empty = hygiene (`[TODO HYGIENE]`); nag channel via `[agent] todo_nag_strategy` (`developer`/`user` = prose nag; `synthetic_tool` = forged `todo_list` call + **real** `stack.render()` result + ToolCall/Result events; `none`).
|
||||||
- **`DEFAULT_TOOLS`**: file + process + vcs + chat + todo tool list for a `ToolRegistry`.
|
- **`net` / `fetch`**: HTTP GET/POST/PUT/DELETE via niquests (isolated session, not the LLM client). Manual redirects with per-hop SSRF checks (`asyncio` `getaddrinfo`). Private/loopback hosts need instance policy grant (CLI timed confirm; **not** YOLO). Public cleartext HTTP and mutating methods soft-confirm (`YOLO`/`TRUSTABLE`). HTTPS→HTTP redirects blocked unless `allow_http_downgrade`. `user_agent` tool arg (or headers) is never replaced by the default when provided. Caps: body bytes/chars, request body size.
|
||||||
|
- **`DEFAULT_TOOLS` / catalog**: file + process + vcs + chat + todo + net (`fetch`) for local surface.
|
||||||
|
|
||||||
### Prompting (`prompting.py`)
|
### Prompting (`prompting.py`)
|
||||||
|
|
||||||
|
|||||||
@@ -243,7 +243,7 @@ User messages are saved immediately. On API error or Ctrl+C, partial assistant/t
|
|||||||
|
|
||||||
## Tools (when enabled)
|
## Tools (when enabled)
|
||||||
|
|
||||||
Default registry: file ops (including `tree` with default noise-dir skips), `run_command` / PTY (POSIX openpty; Windows ConPTY via pywinpty), read-only VCS (git), human prompts (`ask_user_line` / `ask_user_choice` / `ask_user_form`), and todo stack tools (`todo_list` / `todo_push` / `todo_pop` / `todo_update` / `todo_clear`).
|
Default registry: file ops (including `tree` with default noise-dir skips), `run_command` / PTY (POSIX openpty; Windows ConPTY via pywinpty), read-only VCS (git), HTTP `fetch` (GET/POST/PUT/DELETE via niquests; private/loopback hosts need a human policy grant, not YOLO), human prompts (`ask_user_line` / `ask_user_choice` / `ask_user_form`), and todo stack tools (`todo_list` / `todo_push` / `todo_pop` / `todo_update` / `todo_clear`).
|
||||||
|
|
||||||
Safety defaults:
|
Safety defaults:
|
||||||
|
|
||||||
|
|||||||
+1
-1
@@ -21,7 +21,7 @@
|
|||||||
- router: Multi-source capability routing (Phase H; not implemented).
|
- router: Multi-source capability routing (Phase H; not implemented).
|
||||||
- config: Plyngent configuration center (TOML), including ``[plugins]``.
|
- config: Plyngent configuration center (TOML), including ``[plugins]``.
|
||||||
- agent: Tool loop, streaming, usage, compact; `@tool` / tags / registry.
|
- agent: Tool loop, streaming, usage, compact; `@tool` / tags / registry.
|
||||||
- tools: Workspace file/process/VCS/chat/todo tools; catalog; plugins;
|
- tools: Workspace file/process/VCS/chat/todo/net tools; catalog; plugins;
|
||||||
instance/session context and views.
|
instance/session context and views.
|
||||||
- prompting: Shared ask/choose/form for CLI and tools.
|
- prompting: Shared ask/choose/form for CLI and tools.
|
||||||
- cli: Click entry, slash registry, REPL, one-shot chat.
|
- cli: Click entry, slash registry, REPL, one-shot chat.
|
||||||
|
|||||||
@@ -41,6 +41,9 @@ max_context_tokens = 200000
|
|||||||
# developer (default) | user | synthetic_tool | none
|
# developer (default) | user | synthetic_tool | none
|
||||||
# (legacy system is treated as developer)
|
# (legacy system is treated as developer)
|
||||||
# todo_nag_strategy = "developer"
|
# todo_nag_strategy = "developer"
|
||||||
|
# Append-only developer playbook checkpoints every N prompt_tokens (0 = off).
|
||||||
|
# directive_reminder_tokens = 100000
|
||||||
|
# directive_reminder_text = ""
|
||||||
|
|
||||||
# Optional: custom compact/summarisation prompts (empty = built-in defaults).
|
# Optional: custom compact/summarisation prompts (empty = built-in defaults).
|
||||||
# compact_system_prompt = ""
|
# compact_system_prompt = ""
|
||||||
|
|||||||
+101
-3
@@ -17,6 +17,10 @@ from .budget import (
|
|||||||
DEFAULT_TOOL_RESULT_MAX_CHARS,
|
DEFAULT_TOOL_RESULT_MAX_CHARS,
|
||||||
estimate_messages_tokens,
|
estimate_messages_tokens,
|
||||||
)
|
)
|
||||||
|
from .directive_checkpoint import (
|
||||||
|
DEFAULT_DIRECTIVE_REMINDER_TOKENS,
|
||||||
|
parse_checkpoint_bands,
|
||||||
|
)
|
||||||
from .events import UsageEvent
|
from .events import UsageEvent
|
||||||
from .loop import DEFAULT_MAX_ROUNDS, run_chat_loop
|
from .loop import DEFAULT_MAX_ROUNDS, run_chat_loop
|
||||||
from .todo_nag import (
|
from .todo_nag import (
|
||||||
@@ -128,11 +132,15 @@ class ChatAgent:
|
|||||||
max_context_tokens: int
|
max_context_tokens: int
|
||||||
todo_stack: TodoStack | None
|
todo_stack: TodoStack | None
|
||||||
todo_nag_strategy: TodoNagStrategy
|
todo_nag_strategy: TodoNagStrategy
|
||||||
|
directive_reminder_tokens: int
|
||||||
|
directive_reminder_text: str | None
|
||||||
messages: list[AnyChatMessage]
|
messages: list[AnyChatMessage]
|
||||||
session_usage: TokenUsage
|
session_usage: TokenUsage
|
||||||
last_turn_usage: TokenUsage
|
last_turn_usage: TokenUsage
|
||||||
last_request_usage: TokenUsage
|
last_request_usage: TokenUsage
|
||||||
last_turn_rounds: int
|
last_turn_rounds: int
|
||||||
|
peak_prompt_tokens: int
|
||||||
|
reminder_last_band: int
|
||||||
# Index into messages of the first unpersisted message (checkpoint cursor).
|
# Index into messages of the first unpersisted message (checkpoint cursor).
|
||||||
_persist_from: int
|
_persist_from: int
|
||||||
|
|
||||||
@@ -155,6 +163,10 @@ class ChatAgent:
|
|||||||
max_context_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
|
max_context_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
|
||||||
todo_stack: TodoStack | None = None,
|
todo_stack: TodoStack | None = None,
|
||||||
todo_nag_strategy: str | TodoNagStrategy = DEFAULT_TODO_NAG_STRATEGY,
|
todo_nag_strategy: str | TodoNagStrategy = DEFAULT_TODO_NAG_STRATEGY,
|
||||||
|
directive_reminder_tokens: int = DEFAULT_DIRECTIVE_REMINDER_TOKENS,
|
||||||
|
directive_reminder_text: str | None = None,
|
||||||
|
peak_prompt_tokens: int = 0,
|
||||||
|
reminder_last_band: int = 0,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.client = client
|
self.client = client
|
||||||
self.model = model
|
self.model = model
|
||||||
@@ -171,13 +183,19 @@ class ChatAgent:
|
|||||||
self.max_context_tokens = max_context_tokens
|
self.max_context_tokens = max_context_tokens
|
||||||
self.todo_stack = todo_stack
|
self.todo_stack = todo_stack
|
||||||
self.todo_nag_strategy = parse_todo_nag_strategy(str(todo_nag_strategy))
|
self.todo_nag_strategy = parse_todo_nag_strategy(str(todo_nag_strategy))
|
||||||
|
self.directive_reminder_tokens = max(0, int(directive_reminder_tokens))
|
||||||
|
text = directive_reminder_text.strip() if directive_reminder_text else ""
|
||||||
|
self.directive_reminder_text = text or None
|
||||||
self.messages = list(messages) if messages is not None else []
|
self.messages = list(messages) if messages is not None else []
|
||||||
self.session_usage = TokenUsage()
|
self.session_usage = TokenUsage()
|
||||||
self.last_turn_usage = TokenUsage()
|
self.last_turn_usage = TokenUsage()
|
||||||
self.last_request_usage = TokenUsage()
|
self.last_request_usage = TokenUsage()
|
||||||
self.last_turn_rounds = 0
|
self.last_turn_rounds = 0
|
||||||
|
self.peak_prompt_tokens = max(0, int(peak_prompt_tokens))
|
||||||
|
self.reminder_last_band = max(0, int(reminder_last_band))
|
||||||
self._persist_from = len(self.messages)
|
self._persist_from = len(self.messages)
|
||||||
self._ensure_system_prompt()
|
self._ensure_system_prompt()
|
||||||
|
self._sync_reminder_band_from_messages()
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def pending_retry_text(self) -> str | None:
|
def pending_retry_text(self) -> str | None:
|
||||||
@@ -214,6 +232,53 @@ class ChatAgent:
|
|||||||
# already pointed past stored messages stay correct after insert.
|
# already pointed past stored messages stay correct after insert.
|
||||||
self._persist_from = min(len(self.messages), self._persist_from + 1)
|
self._persist_from = min(len(self.messages), self._persist_from + 1)
|
||||||
|
|
||||||
|
def _sync_reminder_band_from_messages(self) -> None:
|
||||||
|
"""Raise :attr:`reminder_last_band` to match durable checkpoint markers."""
|
||||||
|
from_history = parse_checkpoint_bands(self.messages)
|
||||||
|
self.reminder_last_band = max(self.reminder_last_band, from_history)
|
||||||
|
|
||||||
|
def apply_session_context_usage(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
last_prompt_tokens: int | None = None,
|
||||||
|
peak_prompt_tokens: int | None = None,
|
||||||
|
last_completion_tokens: int | None = None,
|
||||||
|
usage_source: str | None = None,
|
||||||
|
reminder_last_band: int | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Hydrate usage / reminder band from a session row (resume)."""
|
||||||
|
if peak_prompt_tokens is not None:
|
||||||
|
self.peak_prompt_tokens = max(self.peak_prompt_tokens, int(peak_prompt_tokens))
|
||||||
|
if reminder_last_band is not None:
|
||||||
|
self.reminder_last_band = max(self.reminder_last_band, int(reminder_last_band))
|
||||||
|
if last_prompt_tokens is not None and last_prompt_tokens > 0:
|
||||||
|
source = usage_source or "api"
|
||||||
|
self.last_request_usage = TokenUsage(
|
||||||
|
prompt_tokens=int(last_prompt_tokens),
|
||||||
|
completion_tokens=int(last_completion_tokens or 0),
|
||||||
|
total_tokens=int(last_prompt_tokens) + int(last_completion_tokens or 0),
|
||||||
|
source=source,
|
||||||
|
)
|
||||||
|
self.peak_prompt_tokens = max(self.peak_prompt_tokens, int(last_prompt_tokens))
|
||||||
|
self._sync_reminder_band_from_messages()
|
||||||
|
|
||||||
|
async def _persist_context_usage(self) -> None:
|
||||||
|
if self.memory is None or self.session_id is None:
|
||||||
|
return
|
||||||
|
last = self.last_request_usage
|
||||||
|
_ = await self.memory.update_session_context_usage(
|
||||||
|
self.session_id,
|
||||||
|
last_prompt_tokens=last.prompt_tokens if not last.is_zero() else None,
|
||||||
|
peak_prompt_tokens=self.peak_prompt_tokens or None,
|
||||||
|
last_completion_tokens=last.completion_tokens if not last.is_zero() else None,
|
||||||
|
usage_source=last.source if not last.is_zero() else None,
|
||||||
|
reminder_last_band=self.reminder_last_band,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _on_reminder_band(self, band: int) -> None:
|
||||||
|
self.reminder_last_band = max(self.reminder_last_band, band)
|
||||||
|
await self._persist_context_usage()
|
||||||
|
|
||||||
def replace_messages(
|
def replace_messages(
|
||||||
self,
|
self,
|
||||||
messages: Sequence[AnyChatMessage],
|
messages: Sequence[AnyChatMessage],
|
||||||
@@ -233,6 +298,7 @@ class ChatAgent:
|
|||||||
else:
|
else:
|
||||||
self._persist_from = len(self.messages) if persisted else 0
|
self._persist_from = len(self.messages) if persisted else 0
|
||||||
self._ensure_system_prompt()
|
self._ensure_system_prompt()
|
||||||
|
self._sync_reminder_band_from_messages()
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def persist_from(self) -> int:
|
def persist_from(self) -> int:
|
||||||
@@ -246,6 +312,15 @@ class ChatAgent:
|
|||||||
raise RuntimeError(msg)
|
raise RuntimeError(msg)
|
||||||
loaded = await self.memory.list_messages(self.session_id)
|
loaded = await self.memory.list_messages(self.session_id)
|
||||||
self.replace_messages(loaded, persisted=True)
|
self.replace_messages(loaded, persisted=True)
|
||||||
|
row = await self.memory.get_session(self.session_id)
|
||||||
|
if row is not None:
|
||||||
|
self.apply_session_context_usage(
|
||||||
|
last_prompt_tokens=row.last_prompt_tokens,
|
||||||
|
peak_prompt_tokens=row.peak_prompt_tokens,
|
||||||
|
last_completion_tokens=row.last_completion_tokens,
|
||||||
|
usage_source=row.usage_source,
|
||||||
|
reminder_last_band=row.reminder_last_band,
|
||||||
|
)
|
||||||
|
|
||||||
async def bind_session(self, session_id: int, *, load: bool = True) -> None:
|
async def bind_session(self, session_id: int, *, load: bool = True) -> None:
|
||||||
"""Attach a memory session id; optionally load existing messages."""
|
"""Attach a memory session id; optionally load existing messages."""
|
||||||
@@ -295,6 +370,15 @@ class ChatAgent:
|
|||||||
end = committed_prefix_end(self.messages, user_index)
|
end = committed_prefix_end(self.messages, user_index)
|
||||||
del self.messages[end:]
|
del self.messages[end:]
|
||||||
|
|
||||||
|
def _developer_tail_end(self, start: int) -> int:
|
||||||
|
"""Extend *start* through trailing developer checkpoint messages."""
|
||||||
|
from plyngent.lmproto.openai_compatible.model import DeveloperChatMessage
|
||||||
|
|
||||||
|
end = start
|
||||||
|
while end < len(self.messages) and isinstance(self.messages[end], DeveloperChatMessage):
|
||||||
|
end += 1
|
||||||
|
return end
|
||||||
|
|
||||||
async def _run_from_user_message(self, user_msg: UserChatMessage) -> AsyncIterator[AgentEvent]:
|
async def _run_from_user_message(self, user_msg: UserChatMessage) -> AsyncIterator[AgentEvent]:
|
||||||
"""Run the tool loop for an already-appended user message.
|
"""Run the tool loop for an already-appended user message.
|
||||||
|
|
||||||
@@ -340,17 +424,29 @@ class ChatAgent:
|
|||||||
max_context_tokens=self.max_context_tokens,
|
max_context_tokens=self.max_context_tokens,
|
||||||
todo_stack=self.todo_stack,
|
todo_stack=self.todo_stack,
|
||||||
todo_nag_strategy=self.todo_nag_strategy,
|
todo_nag_strategy=self.todo_nag_strategy,
|
||||||
|
directive_reminder_tokens=self.directive_reminder_tokens,
|
||||||
|
directive_reminder_text=self.directive_reminder_text,
|
||||||
|
reminder_last_band=self.reminder_last_band,
|
||||||
|
on_reminder_band=self._on_reminder_band,
|
||||||
):
|
):
|
||||||
if isinstance(event, UsageEvent):
|
if isinstance(event, UsageEvent):
|
||||||
turn_rounds += 1
|
turn_rounds += 1
|
||||||
last_request = event.usage
|
last_request = event.usage
|
||||||
turn_usage = turn_usage.add(event.usage)
|
turn_usage = turn_usage.add(event.usage)
|
||||||
self.session_usage = self.session_usage.add(event.usage)
|
self.session_usage = self.session_usage.add(event.usage)
|
||||||
|
self.last_request_usage = event.usage
|
||||||
|
self.peak_prompt_tokens = max(
|
||||||
|
self.peak_prompt_tokens,
|
||||||
|
event.usage.prompt_tokens,
|
||||||
|
)
|
||||||
|
await self._persist_context_usage()
|
||||||
# After a full tool batch, commit prefix so failures do not
|
# After a full tool batch, commit prefix so failures do not
|
||||||
# discard work that already had external effects.
|
# discard work that already had external effects. Trailing
|
||||||
|
# developer checkpoints after that batch are committed too.
|
||||||
commit_end = committed_prefix_end(self.messages, user_index)
|
commit_end = committed_prefix_end(self.messages, user_index)
|
||||||
if commit_end > self._persist_from:
|
end = self._developer_tail_end(commit_end)
|
||||||
await self._persist_range(self._persist_from, commit_end)
|
if end > self._persist_from:
|
||||||
|
await self._persist_range(self._persist_from, end)
|
||||||
yield event
|
yield event
|
||||||
completed = True
|
completed = True
|
||||||
except BaseException:
|
except BaseException:
|
||||||
@@ -363,6 +459,7 @@ class ChatAgent:
|
|||||||
self.last_turn_rounds = turn_rounds
|
self.last_turn_rounds = turn_rounds
|
||||||
if self._persist_from < len(self.messages):
|
if self._persist_from < len(self.messages):
|
||||||
await self._persist_range(self._persist_from, len(self.messages))
|
await self._persist_range(self._persist_from, len(self.messages))
|
||||||
|
await self._persist_context_usage()
|
||||||
|
|
||||||
async def run(self, user_text: str) -> AsyncIterator[AgentEvent]:
|
async def run(self, user_text: str) -> AsyncIterator[AgentEvent]:
|
||||||
"""Append a user message (persist immediately), run the tool loop, yield events."""
|
"""Append a user message (persist immediately), run the tool loop, yield events."""
|
||||||
@@ -461,6 +558,7 @@ class ChatAgent:
|
|||||||
max_context_tokens=self.max_context_tokens,
|
max_context_tokens=self.max_context_tokens,
|
||||||
todo_stack=None,
|
todo_stack=None,
|
||||||
todo_nag_strategy="none",
|
todo_nag_strategy="none",
|
||||||
|
directive_reminder_tokens=0,
|
||||||
messages=history,
|
messages=history,
|
||||||
)
|
)
|
||||||
async for event in aside.run(text):
|
async for event in aside.run(text):
|
||||||
|
|||||||
@@ -0,0 +1,109 @@
|
|||||||
|
"""Append-only developer playbook checkpoints at token-band crossings."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from plyngent.lmproto.openai_compatible.model import DeveloperChatMessage
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Sequence
|
||||||
|
|
||||||
|
from plyngent.lmproto.openai_compatible.model import AnyChatMessage
|
||||||
|
|
||||||
|
# 0 = disabled. Default matches long-session soft drift without spamming short chats.
|
||||||
|
DEFAULT_DIRECTIVE_REMINDER_TOKENS = 100_000
|
||||||
|
|
||||||
|
DEFAULT_DIRECTIVE_REMINDER_TEXT = """\
|
||||||
|
Tool playbook still applies (see system). Hard constraints:
|
||||||
|
- Prefer file tools over shell; several `run_command` calls may run in parallel; \
|
||||||
|
use `run_command_batch` for ordered pipelines.
|
||||||
|
- `edit_replace`: fix match / `max_replaces`; `read_file` with_lineno before `edit_lineno`.
|
||||||
|
- Prefer `fetch` for HTTP(S); private/LAN hosts need human policy allow (not YOLO).
|
||||||
|
- PTY secrets only via `ask_into_pty`; denylists and confirms still apply.
|
||||||
|
- Todo stack: open items mean unfinished work.
|
||||||
|
"""
|
||||||
|
|
||||||
|
_BAND_MARKER = re.compile(
|
||||||
|
r"\[DIRECTIVE CHECKPOINT band=(\d+)\b",
|
||||||
|
re.IGNORECASE,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def checkpoint_body(
|
||||||
|
band: int,
|
||||||
|
*,
|
||||||
|
tokens: int,
|
||||||
|
source: str,
|
||||||
|
reminder_text: str | None = None,
|
||||||
|
) -> str:
|
||||||
|
"""Build a durable developer checkpoint message body for *band*."""
|
||||||
|
playbook = (reminder_text if reminder_text is not None else DEFAULT_DIRECTIVE_REMINDER_TEXT).strip()
|
||||||
|
header = f"[DIRECTIVE CHECKPOINT band={band} tokens≈{tokens} source={source}]"
|
||||||
|
if not playbook:
|
||||||
|
return header
|
||||||
|
return f"{header}\n{playbook}"
|
||||||
|
|
||||||
|
|
||||||
|
def parse_checkpoint_bands(messages: Sequence[AnyChatMessage]) -> int:
|
||||||
|
"""Return the highest checkpoint band found in durable history (0 if none)."""
|
||||||
|
highest = 0
|
||||||
|
for msg in messages:
|
||||||
|
if not isinstance(msg, DeveloperChatMessage):
|
||||||
|
continue
|
||||||
|
match = _BAND_MARKER.search(msg.content)
|
||||||
|
if match is None:
|
||||||
|
continue
|
||||||
|
highest = max(highest, int(match.group(1)))
|
||||||
|
return highest
|
||||||
|
|
||||||
|
|
||||||
|
def bands_to_fire(*, last_fired_band: int, current_band: int) -> list[int]:
|
||||||
|
"""Inclusive bands to append so markers stay monotonic (fill gaps)."""
|
||||||
|
if current_band <= last_fired_band:
|
||||||
|
return []
|
||||||
|
return list(range(last_fired_band + 1, current_band + 1))
|
||||||
|
|
||||||
|
|
||||||
|
def token_band(prompt_tokens: int, interval: int) -> int:
|
||||||
|
"""Band index for *prompt_tokens* (0 = below first threshold)."""
|
||||||
|
if interval < 1 or prompt_tokens < 1:
|
||||||
|
return 0
|
||||||
|
return prompt_tokens // interval
|
||||||
|
|
||||||
|
|
||||||
|
def inject_directive_checkpoints(
|
||||||
|
messages: list[AnyChatMessage],
|
||||||
|
*,
|
||||||
|
prompt_tokens: int,
|
||||||
|
source: str,
|
||||||
|
interval: int,
|
||||||
|
last_fired_band: int,
|
||||||
|
reminder_text: str | None = None,
|
||||||
|
) -> tuple[int, list[DeveloperChatMessage]]:
|
||||||
|
"""Append developer checkpoints for newly crossed bands.
|
||||||
|
|
||||||
|
Returns ``(new_last_fired_band, appended_messages)``. Does nothing when
|
||||||
|
*interval* < 1 or no new bands are crossed. Append-only (never edits prior
|
||||||
|
checkpoints) so prefix caching can keep a stable history prefix.
|
||||||
|
"""
|
||||||
|
if interval < 1:
|
||||||
|
return last_fired_band, []
|
||||||
|
current = token_band(prompt_tokens, interval)
|
||||||
|
to_fire = bands_to_fire(last_fired_band=last_fired_band, current_band=current)
|
||||||
|
if not to_fire:
|
||||||
|
return last_fired_band, []
|
||||||
|
appended: list[DeveloperChatMessage] = []
|
||||||
|
for band in to_fire:
|
||||||
|
msg = DeveloperChatMessage(
|
||||||
|
content=checkpoint_body(
|
||||||
|
band,
|
||||||
|
tokens=prompt_tokens,
|
||||||
|
source=source,
|
||||||
|
reminder_text=reminder_text,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
messages.append(msg)
|
||||||
|
appended.append(msg)
|
||||||
|
return to_fire[-1], appended
|
||||||
+72
-10
@@ -26,6 +26,10 @@ from .budget import (
|
|||||||
estimate_messages_tokens,
|
estimate_messages_tokens,
|
||||||
truncate_tool_result,
|
truncate_tool_result,
|
||||||
)
|
)
|
||||||
|
from .directive_checkpoint import (
|
||||||
|
DEFAULT_DIRECTIVE_REMINDER_TOKENS,
|
||||||
|
inject_directive_checkpoints,
|
||||||
|
)
|
||||||
from .events import (
|
from .events import (
|
||||||
AgentEvent,
|
AgentEvent,
|
||||||
AssistantMessageEvent,
|
AssistantMessageEvent,
|
||||||
@@ -312,7 +316,36 @@ def _last_assistant(messages: list[AnyChatMessage], pre_len: int) -> AssistantCh
|
|||||||
return last
|
return last
|
||||||
|
|
||||||
|
|
||||||
async def run_chat_loop( # noqa: C901 — multi-phase tool loop
|
async def _maybe_inject_directive_checkpoints(
|
||||||
|
messages: list[AnyChatMessage],
|
||||||
|
*,
|
||||||
|
usage_event: UsageEvent | None,
|
||||||
|
interval: int,
|
||||||
|
last_band: int,
|
||||||
|
reminder_text: str | None,
|
||||||
|
on_reminder_band: Callable[[int], Awaitable[None] | None] | None,
|
||||||
|
) -> int:
|
||||||
|
"""Append durable checkpoints after a usage sample; return updated last band."""
|
||||||
|
if usage_event is None or interval < 1:
|
||||||
|
return last_band
|
||||||
|
new_band, appended = inject_directive_checkpoints(
|
||||||
|
messages,
|
||||||
|
prompt_tokens=usage_event.usage.prompt_tokens,
|
||||||
|
source=usage_event.usage.source,
|
||||||
|
interval=interval,
|
||||||
|
last_fired_band=last_band,
|
||||||
|
reminder_text=reminder_text,
|
||||||
|
)
|
||||||
|
if not appended:
|
||||||
|
return last_band
|
||||||
|
if on_reminder_band is not None:
|
||||||
|
maybe = on_reminder_band(new_band)
|
||||||
|
if inspect.isawaitable(maybe):
|
||||||
|
await maybe
|
||||||
|
return new_band
|
||||||
|
|
||||||
|
|
||||||
|
async def run_chat_loop( # noqa: C901, PLR0912 — multi-phase tool loop
|
||||||
client: ChatClient,
|
client: ChatClient,
|
||||||
messages: list[AnyChatMessage],
|
messages: list[AnyChatMessage],
|
||||||
*,
|
*,
|
||||||
@@ -327,6 +360,10 @@ async def run_chat_loop( # noqa: C901 — multi-phase tool loop
|
|||||||
max_context_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
|
max_context_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
|
||||||
todo_stack: TodoStack | None = None,
|
todo_stack: TodoStack | None = None,
|
||||||
todo_nag_strategy: TodoNagStrategy = DEFAULT_TODO_NAG_STRATEGY,
|
todo_nag_strategy: TodoNagStrategy = DEFAULT_TODO_NAG_STRATEGY,
|
||||||
|
directive_reminder_tokens: int = DEFAULT_DIRECTIVE_REMINDER_TOKENS,
|
||||||
|
directive_reminder_text: str | None = None,
|
||||||
|
reminder_last_band: int = 0,
|
||||||
|
on_reminder_band: Callable[[int], Awaitable[None] | None] | None = None,
|
||||||
) -> AsyncIterator[AgentEvent]:
|
) -> AsyncIterator[AgentEvent]:
|
||||||
"""Multi-round chat/tool loop; mutates ``messages`` in place and yields events.
|
"""Multi-round chat/tool loop; mutates ``messages`` in place and yields events.
|
||||||
|
|
||||||
@@ -338,6 +375,11 @@ async def run_chat_loop( # noqa: C901 — multi-phase tool loop
|
|||||||
When *todo_stack* is set and still needs review after a natural stop
|
When *todo_stack* is set and still needs review after a natural stop
|
||||||
(open items, or non-empty stack untouched this turn), injects a review nag
|
(open items, or non-empty stack untouched this turn), injects a review nag
|
||||||
(channel from *todo_nag_strategy*) once so the model reconciles unfinished work.
|
(channel from *todo_nag_strategy*) once so the model reconciles unfinished work.
|
||||||
|
|
||||||
|
After each usage sample, may append durable developer directive checkpoints
|
||||||
|
when *prompt_tokens* crosses bands of *directive_reminder_tokens* (0 disables).
|
||||||
|
*reminder_last_band* is the highest band already injected (history/DB).
|
||||||
|
*on_reminder_band* is notified with the new last band after appends.
|
||||||
"""
|
"""
|
||||||
tool_items: Sequence[AnyToolItem] | None = None
|
tool_items: Sequence[AnyToolItem] | None = None
|
||||||
if tools is not None and len(tools) > 0:
|
if tools is not None and len(tools) > 0:
|
||||||
@@ -349,6 +391,7 @@ async def run_chat_loop( # noqa: C901 — multi-phase tool loop
|
|||||||
prompt_tokens_hint: int | None = None
|
prompt_tokens_hint: int | None = None
|
||||||
sent_estimate_tokens: int | None = None
|
sent_estimate_tokens: int | None = None
|
||||||
todo_review_injected = False
|
todo_review_injected = False
|
||||||
|
last_band = max(0, reminder_last_band)
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
while rounds_used < allowance:
|
while rounds_used < allowance:
|
||||||
@@ -372,15 +415,42 @@ async def run_chat_loop( # noqa: C901 — multi-phase tool loop
|
|||||||
)
|
)
|
||||||
|
|
||||||
pre_len = len(messages)
|
pre_len = len(messages)
|
||||||
|
last_usage_event: UsageEvent | None = None
|
||||||
async for event in _assistant_round(client, param, messages, stream=stream):
|
async for event in _assistant_round(client, param, messages, stream=stream):
|
||||||
if isinstance(event, UsageEvent):
|
if isinstance(event, UsageEvent):
|
||||||
# Next rounds scale char-estimates by real/resolved prompt size.
|
# Next rounds scale char-estimates by real/resolved prompt size.
|
||||||
prompt_tokens_hint = event.usage.prompt_tokens
|
prompt_tokens_hint = event.usage.prompt_tokens
|
||||||
sent_estimate_tokens = sent_est
|
sent_estimate_tokens = sent_est
|
||||||
|
last_usage_event = event
|
||||||
yield event
|
yield event
|
||||||
|
|
||||||
assistant = _last_assistant(messages, pre_len)
|
assistant = _last_assistant(messages, pre_len)
|
||||||
tool_calls = assistant.tool_calls
|
tool_calls = assistant.tool_calls
|
||||||
if tool_calls is UNSET or not tool_calls or tools is None:
|
has_tools = tool_calls is not UNSET and bool(tool_calls) and tools is not None
|
||||||
|
if has_tools:
|
||||||
|
assert tools is not None
|
||||||
|
assert tool_calls is not UNSET
|
||||||
|
async for event in _execute_tool_calls(
|
||||||
|
tools,
|
||||||
|
tool_calls,
|
||||||
|
messages,
|
||||||
|
max_result_chars=max_tool_result_chars,
|
||||||
|
parallel=parallel_tools,
|
||||||
|
):
|
||||||
|
yield event
|
||||||
|
|
||||||
|
# After tools (or text-only assistant): durable checkpoints stay out of
|
||||||
|
# assistant→tool batches so commit/retry structure remains valid.
|
||||||
|
last_band = await _maybe_inject_directive_checkpoints(
|
||||||
|
messages,
|
||||||
|
usage_event=last_usage_event,
|
||||||
|
interval=directive_reminder_tokens,
|
||||||
|
last_band=last_band,
|
||||||
|
reminder_text=directive_reminder_text,
|
||||||
|
on_reminder_band=on_reminder_band,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not has_tools:
|
||||||
if todo_stack is not None and todo_stack.needs_review() and not todo_review_injected:
|
if todo_stack is not None and todo_stack.needs_review() and not todo_review_injected:
|
||||||
todo_review_injected = True
|
todo_review_injected = True
|
||||||
injected, nag_events = inject_todo_nag_for_stack_with_events(
|
injected, nag_events = inject_todo_nag_for_stack_with_events(
|
||||||
@@ -394,14 +464,6 @@ async def run_chat_loop( # noqa: C901 — multi-phase tool loop
|
|||||||
if injected:
|
if injected:
|
||||||
continue
|
continue
|
||||||
return
|
return
|
||||||
async for event in _execute_tool_calls(
|
|
||||||
tools,
|
|
||||||
tool_calls,
|
|
||||||
messages,
|
|
||||||
max_result_chars=max_tool_result_chars,
|
|
||||||
parallel=parallel_tools,
|
|
||||||
):
|
|
||||||
yield event
|
|
||||||
|
|
||||||
reason = f"tool loop reached {allowance} rounds (used {rounds_used})"
|
reason = f"tool loop reached {allowance} rounds (used {rounds_used})"
|
||||||
if on_limit is not None and await _call_on_limit(on_limit, reason):
|
if on_limit is not None and await _call_on_limit(on_limit, reason):
|
||||||
|
|||||||
+10
-1
@@ -300,9 +300,14 @@ async def _run_chat( # noqa: C901, PLR0912, PLR0915 — chat orchestration
|
|||||||
# Path denylist and policy confirm live on instance.workspace (no process bag).
|
# Path denylist and policy confirm live on instance.workspace (no process bag).
|
||||||
state.instance_state.workspace.path_denylist = tuple(store.agent_config.path_denylist or ())
|
state.instance_state.workspace.path_denylist = tuple(store.agent_config.path_denylist or ())
|
||||||
if interactive:
|
if interactive:
|
||||||
from plyngent.cli.limits import prompt_policy_command_confirm
|
from plyngent.cli.limits import prompt_policy_command_confirm, prompt_policy_fetch_confirm
|
||||||
|
from plyngent.tools.net import set_fetch_policy_confirm_hook
|
||||||
|
|
||||||
state.instance_state.workspace.policy_confirm_hook = prompt_policy_command_confirm
|
state.instance_state.workspace.policy_confirm_hook = prompt_policy_command_confirm
|
||||||
|
set_fetch_policy_confirm_hook(
|
||||||
|
prompt_policy_fetch_confirm,
|
||||||
|
instance=state.instance_state,
|
||||||
|
)
|
||||||
# Seed cache if we already fetched; else warm in background for Tab.
|
# Seed cache if we already fetched; else warm in background for Tab.
|
||||||
if remote_ids is not None:
|
if remote_ids is not None:
|
||||||
state.seed_remote_models(remote_ids)
|
state.seed_remote_models(remote_ids)
|
||||||
@@ -334,6 +339,10 @@ async def _run_chat( # noqa: C901, PLR0912, PLR0915 — chat orchestration
|
|||||||
if isinstance(state_obj, ReplState):
|
if isinstance(state_obj, ReplState):
|
||||||
state_obj.instance_state.workspace.policy_confirm_hook = None
|
state_obj.instance_state.workspace.policy_confirm_hook = None
|
||||||
state_obj.instance_state.workspace.policy_allowed_commands.clear()
|
state_obj.instance_state.workspace.policy_allowed_commands.clear()
|
||||||
|
from plyngent.tools.net import clear_private_grants, set_fetch_policy_confirm_hook
|
||||||
|
|
||||||
|
set_fetch_policy_confirm_hook(None, instance=state_obj.instance_state)
|
||||||
|
clear_private_grants(instance=state_obj.instance_state)
|
||||||
await state_obj.instance_state.shutdown()
|
await state_obj.instance_state.shutdown()
|
||||||
else:
|
else:
|
||||||
# No ReplState: only PTY class cleanup (temps require instance allowlist).
|
# No ReplState: only PTY class cleanup (temps require instance allowlist).
|
||||||
|
|||||||
@@ -232,6 +232,51 @@ def prompt_policy_command_confirm(
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def prompt_policy_fetch_confirm(
|
||||||
|
host: str,
|
||||||
|
port: int,
|
||||||
|
url: str,
|
||||||
|
timeout_seconds: float,
|
||||||
|
) -> bool:
|
||||||
|
"""Ask whether to allow fetch to a private/loopback host (timed).
|
||||||
|
|
||||||
|
Independent of YOLO: always prompts when interactive. Default is **deny**.
|
||||||
|
Timeout, cancel, or non-interactive → False. On approve, the host:port is
|
||||||
|
granted for the rest of this process (see net grants).
|
||||||
|
"""
|
||||||
|
reason = (
|
||||||
|
f"fetch targets private/loopback host {host!r} port {port}\n"
|
||||||
|
f"url: {url}\n"
|
||||||
|
f"Allow this host:port for the rest of this process?\n"
|
||||||
|
f"(timeout {timeout_seconds:g}s defaults to DENY; not skipped by YOLO)"
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
with pause_task_cancel_for_prompt():
|
||||||
|
backend = get_prompt_backend()
|
||||||
|
if not backend.is_interactive():
|
||||||
|
return False
|
||||||
|
backend.echo()
|
||||||
|
backend.secho(format_tool_confirm_box("policy", reason), fg="yellow")
|
||||||
|
backend.echo()
|
||||||
|
grant = f"{host}:{port}"
|
||||||
|
prompt = f"[policy] allow fetch {grant!r}? [y/N] (timeout {timeout_seconds:g}s): "
|
||||||
|
with contextlib.suppress(OSError):
|
||||||
|
_ = sys.stderr.write(prompt)
|
||||||
|
_ = sys.stderr.flush()
|
||||||
|
raw = _read_yes_no_line_with_timeout(timeout_seconds)
|
||||||
|
if raw is None:
|
||||||
|
with contextlib.suppress(OSError, NonInteractiveError):
|
||||||
|
backend.secho(
|
||||||
|
f"[policy] timed out after {timeout_seconds:g}s — denied",
|
||||||
|
fg="red",
|
||||||
|
err=True,
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
return raw.strip().lower() in {"y", "yes"}
|
||||||
|
except NonInteractiveError, KeyboardInterrupt, EOFError:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
async def prompt_confirm_tool_async(name: str, args: Mapping[str, object], reason: str) -> bool | str:
|
async def prompt_confirm_tool_async(name: str, args: Mapping[str, object], reason: str) -> bool | str:
|
||||||
"""Async confirm: True allow, False deny, str = deny with user comment."""
|
"""Async confirm: True allow, False deny, str = deny with user comment."""
|
||||||
del args
|
del args
|
||||||
|
|||||||
@@ -544,6 +544,9 @@ def status_cmd(state: ReplState) -> None:
|
|||||||
f"context_tokens={ctx_tilde}{ctx_tokens}/{ctx_budget} ({ctx_tag}) "
|
f"context_tokens={ctx_tilde}{ctx_tokens}/{ctx_budget} ({ctx_tag}) "
|
||||||
f"context_chars={ctx_chars} "
|
f"context_chars={ctx_chars} "
|
||||||
f"tool_result_max={state.agent.max_tool_result_chars}\n"
|
f"tool_result_max={state.agent.max_tool_result_chars}\n"
|
||||||
|
f"peak_prompt_tokens={state.agent.peak_prompt_tokens} "
|
||||||
|
f"reminder_band={state.agent.reminder_last_band} "
|
||||||
|
f"reminder_every={state.agent.directive_reminder_tokens or 'off'}\n"
|
||||||
f"last_request={last_req.format_line()}\n"
|
f"last_request={last_req.format_line()}\n"
|
||||||
f"usage_last_turn={last_u.format_line(billed=True)} "
|
f"usage_last_turn={last_u.format_line(billed=True)} "
|
||||||
f"rounds={last_rounds}\n"
|
f"rounds={last_rounds}\n"
|
||||||
|
|||||||
@@ -236,7 +236,14 @@ class ReplState:
|
|||||||
agent_cfg.tool_directives,
|
agent_cfg.tool_directives,
|
||||||
)
|
)
|
||||||
on_limit = prompt_continue_limit_async if self.interactive_limits else None
|
on_limit = prompt_continue_limit_async if self.interactive_limits else None
|
||||||
return ChatAgent(
|
peak = 0
|
||||||
|
band = 0
|
||||||
|
last_req = None
|
||||||
|
if hasattr(self, "agent"):
|
||||||
|
peak = self.agent.peak_prompt_tokens
|
||||||
|
band = self.agent.reminder_last_band
|
||||||
|
last_req = self.agent.last_request_usage
|
||||||
|
agent = ChatAgent(
|
||||||
self.client,
|
self.client,
|
||||||
model=self.model,
|
model=self.model,
|
||||||
tools=self._tool_registry(),
|
tools=self._tool_registry(),
|
||||||
@@ -251,7 +258,14 @@ class ReplState:
|
|||||||
max_context_tokens=agent_cfg.max_context_tokens,
|
max_context_tokens=agent_cfg.max_context_tokens,
|
||||||
todo_stack=self.todo_stack,
|
todo_stack=self.todo_stack,
|
||||||
todo_nag_strategy=agent_cfg.todo_nag_strategy,
|
todo_nag_strategy=agent_cfg.todo_nag_strategy,
|
||||||
|
directive_reminder_tokens=agent_cfg.directive_reminder_tokens,
|
||||||
|
directive_reminder_text=agent_cfg.directive_reminder_text or None,
|
||||||
|
peak_prompt_tokens=peak,
|
||||||
|
reminder_last_band=band,
|
||||||
)
|
)
|
||||||
|
if last_req is not None and not last_req.is_zero():
|
||||||
|
agent.last_request_usage = last_req
|
||||||
|
return agent
|
||||||
|
|
||||||
def rebuild_client(self) -> None:
|
def rebuild_client(self) -> None:
|
||||||
"""Recreate client and agent after provider/model/tools change."""
|
"""Recreate client and agent after provider/model/tools change."""
|
||||||
|
|||||||
@@ -29,6 +29,14 @@ raise `max_replaces` or narrow `old_string`.
|
|||||||
for ordered pipelines (`pipe_out` / `mix_stderr` as needed).
|
for ordered pipelines (`pipe_out` / `mix_stderr` as needed).
|
||||||
- Prefer `vcs_*` for status/diff/log/branch when enough.
|
- Prefer `vcs_*` for status/diff/log/branch when enough.
|
||||||
|
|
||||||
|
### Network
|
||||||
|
- Prefer `fetch` (GET/POST/PUT/DELETE) for HTTP(S) docs/APIs over curl/wget via shell.
|
||||||
|
- Set `user_agent` (or a User-Agent header) when the remote expects a specific client; \
|
||||||
|
otherwise a small default is used. Never rely on shell to spoof identity.
|
||||||
|
- Private/loopback/LAN hosts need an explicit human policy allow (not skipped by YOLO).
|
||||||
|
- Prefer provider hosted search (e.g. web_search) for open-ended research; use `fetch` for known URLs.
|
||||||
|
- Respect denials, size limits, and HTTPS→HTTP redirect blocks; do not re-fetch the same URL repeatedly.
|
||||||
|
|
||||||
### Humans & safety
|
### Humans & safety
|
||||||
- Prefer `ask_user_line` / `ask_user_choice` / `ask_user_form` over waiting for the next free-form turn.
|
- Prefer `ask_user_line` / `ask_user_choice` / `ask_user_form` over waiting for the next free-form turn.
|
||||||
- Overwrites, deletes, shells, and risky ops may require human confirm; hard denylists are not skipped by YOLO.
|
- Overwrites, deletes, shells, and risky ops may require human confirm; hard denylists are not skipped by YOLO.
|
||||||
@@ -91,6 +99,12 @@ class AgentConfig(Struct, omit_defaults=True):
|
|||||||
# developer | user | synthetic_tool | none (legacy "system" → developer)
|
# developer | user | synthetic_tool | none (legacy "system" → developer)
|
||||||
todo_nag_strategy: str = "developer"
|
todo_nag_strategy: str = "developer"
|
||||||
|
|
||||||
|
# Append-only developer playbook checkpoints when last prompt_tokens crosses
|
||||||
|
# N * this interval (0 = off). See agent/directive_checkpoint.py.
|
||||||
|
directive_reminder_tokens: int = 100_000
|
||||||
|
# Optional short checklist body (empty = built-in hard-constraint list).
|
||||||
|
directive_reminder_text: str = ""
|
||||||
|
|
||||||
# Compact / summarisation prompts (empty = use built-in defaults).
|
# Compact / summarisation prompts (empty = use built-in defaults).
|
||||||
compact_system_prompt: str = ""
|
compact_system_prompt: str = ""
|
||||||
compact_user_prefix: str = ""
|
compact_user_prefix: str = ""
|
||||||
|
|||||||
@@ -34,6 +34,14 @@ class Session(PlyngentBase):
|
|||||||
model: Mapped[str | None] = mapped_column(String(256), nullable=True)
|
model: Mapped[str | None] = mapped_column(String(256), nullable=True)
|
||||||
# Todo/task stack JSON for multi-step sub-tasks (optional).
|
# Todo/task stack JSON for multi-step sub-tasks (optional).
|
||||||
todo_stack: Mapped[dict[str, object] | None] = mapped_column(JSON(), nullable=True)
|
todo_stack: Mapped[dict[str, object] | None] = mapped_column(JSON(), nullable=True)
|
||||||
|
# Last model request context size (API prompt_tokens preferred).
|
||||||
|
last_prompt_tokens: Mapped[int | None] = mapped_column(nullable=True)
|
||||||
|
peak_prompt_tokens: Mapped[int | None] = mapped_column(nullable=True)
|
||||||
|
last_completion_tokens: Mapped[int | None] = mapped_column(nullable=True)
|
||||||
|
# "api" | "estimate" | None when never recorded.
|
||||||
|
usage_source: Mapped[str | None] = mapped_column(String(16), nullable=True)
|
||||||
|
# Highest directive-reminder band already injected (append-only checkpoints).
|
||||||
|
reminder_last_band: Mapped[int | None] = mapped_column(nullable=True)
|
||||||
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
|
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
|
||||||
updated_at: Mapped[datetime] = mapped_column(
|
updated_at: Mapped[datetime] = mapped_column(
|
||||||
DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
|
DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
|
||||||
|
|||||||
@@ -69,6 +69,7 @@ class MemoryStore:
|
|||||||
await conn.run_sync(_migrate_session_workspace)
|
await conn.run_sync(_migrate_session_workspace)
|
||||||
await conn.run_sync(_migrate_session_llm)
|
await conn.run_sync(_migrate_session_llm)
|
||||||
await conn.run_sync(_migrate_session_todo_stack)
|
await conn.run_sync(_migrate_session_todo_stack)
|
||||||
|
await conn.run_sync(_migrate_session_context_usage)
|
||||||
|
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
"""Dispose the underlying engine."""
|
"""Dispose the underlying engine."""
|
||||||
@@ -234,6 +235,37 @@ class MemoryStore:
|
|||||||
await session.refresh(row)
|
await session.refresh(row)
|
||||||
return row
|
return row
|
||||||
|
|
||||||
|
async def update_session_context_usage(
|
||||||
|
self,
|
||||||
|
sid: int,
|
||||||
|
*,
|
||||||
|
last_prompt_tokens: int | None = None,
|
||||||
|
peak_prompt_tokens: int | None = None,
|
||||||
|
last_completion_tokens: int | None = None,
|
||||||
|
usage_source: str | None = None,
|
||||||
|
reminder_last_band: int | None = None,
|
||||||
|
) -> Session:
|
||||||
|
"""Update remembered context size / directive-reminder band (omit to leave)."""
|
||||||
|
async with self._session_factory() as session:
|
||||||
|
row = await session.get(Session, sid)
|
||||||
|
if row is None:
|
||||||
|
msg = f"session not found: {sid}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
if last_prompt_tokens is not None:
|
||||||
|
row.last_prompt_tokens = int(last_prompt_tokens)
|
||||||
|
if peak_prompt_tokens is not None:
|
||||||
|
row.peak_prompt_tokens = int(peak_prompt_tokens)
|
||||||
|
if last_completion_tokens is not None:
|
||||||
|
row.last_completion_tokens = int(last_completion_tokens)
|
||||||
|
if usage_source is not None:
|
||||||
|
row.usage_source = usage_source
|
||||||
|
if reminder_last_band is not None:
|
||||||
|
row.reminder_last_band = int(reminder_last_band)
|
||||||
|
row.updated_at = datetime.now(UTC)
|
||||||
|
await session.commit()
|
||||||
|
await session.refresh(row)
|
||||||
|
return row
|
||||||
|
|
||||||
async def rename_session(self, sid: int, name: str) -> Session:
|
async def rename_session(self, sid: int, name: str) -> Session:
|
||||||
"""Rename a session (max 64 characters, non-empty after strip)."""
|
"""Rename a session (max 64 characters, non-empty after strip)."""
|
||||||
cleaned = name.strip()
|
cleaned = name.strip()
|
||||||
@@ -345,3 +377,22 @@ def _migrate_session_todo_stack(sync_conn: object) -> None:
|
|||||||
columns = _session_columns(sync_conn)
|
columns = _session_columns(sync_conn)
|
||||||
if "todo_stack" not in columns:
|
if "todo_stack" not in columns:
|
||||||
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN todo_stack JSON"))
|
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN todo_stack JSON"))
|
||||||
|
|
||||||
|
|
||||||
|
def _migrate_session_context_usage(sync_conn: object) -> None:
|
||||||
|
"""Add session context usage + directive reminder band columns."""
|
||||||
|
from sqlalchemy.engine import Connection
|
||||||
|
|
||||||
|
if not isinstance(sync_conn, Connection):
|
||||||
|
return
|
||||||
|
columns = _session_columns(sync_conn)
|
||||||
|
if "last_prompt_tokens" not in columns:
|
||||||
|
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN last_prompt_tokens INTEGER"))
|
||||||
|
if "peak_prompt_tokens" not in columns:
|
||||||
|
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN peak_prompt_tokens INTEGER"))
|
||||||
|
if "last_completion_tokens" not in columns:
|
||||||
|
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN last_completion_tokens INTEGER"))
|
||||||
|
if "usage_source" not in columns:
|
||||||
|
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN usage_source VARCHAR(16)"))
|
||||||
|
if "reminder_last_band" not in columns:
|
||||||
|
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN reminder_last_band INTEGER"))
|
||||||
|
|||||||
@@ -28,6 +28,11 @@ from .file import move_path as move_path
|
|||||||
from .file import read_file as read_file
|
from .file import read_file as read_file
|
||||||
from .file import tree as tree
|
from .file import tree as tree
|
||||||
from .file import write_file as write_file
|
from .file import write_file as write_file
|
||||||
|
from .net import NET_TOOLS as NET_TOOLS
|
||||||
|
from .net import clear_private_grants as clear_private_grants
|
||||||
|
from .net import fetch as fetch
|
||||||
|
from .net import grant_private_host as grant_private_host
|
||||||
|
from .net import set_fetch_policy_confirm_hook as set_fetch_policy_confirm_hook
|
||||||
from .plugins import load_plugin_tools as load_plugin_tools
|
from .plugins import load_plugin_tools as load_plugin_tools
|
||||||
from .process import PROCESS_TOOLS as PROCESS_TOOLS
|
from .process import PROCESS_TOOLS as PROCESS_TOOLS
|
||||||
from .process import ask_into_pty as ask_into_pty
|
from .process import ask_into_pty as ask_into_pty
|
||||||
|
|||||||
@@ -213,11 +213,19 @@ def _ensure_builtin_definitions_registered(catalog: ToolCatalog) -> None:
|
|||||||
"""
|
"""
|
||||||
from plyngent.tools.chat import CHAT_TOOLS
|
from plyngent.tools.chat import CHAT_TOOLS
|
||||||
from plyngent.tools.file import FILE_TOOLS
|
from plyngent.tools.file import FILE_TOOLS
|
||||||
|
from plyngent.tools.net import NET_TOOLS
|
||||||
from plyngent.tools.process import PROCESS_TOOLS
|
from plyngent.tools.process import PROCESS_TOOLS
|
||||||
from plyngent.tools.todo import TODO_TOOLS
|
from plyngent.tools.todo import TODO_TOOLS
|
||||||
from plyngent.tools.vcs import VCS_TOOLS
|
from plyngent.tools.vcs import VCS_TOOLS
|
||||||
|
|
||||||
for definition in (*FILE_TOOLS, *PROCESS_TOOLS, *VCS_TOOLS, *CHAT_TOOLS, *TODO_TOOLS):
|
for definition in (
|
||||||
|
*FILE_TOOLS,
|
||||||
|
*PROCESS_TOOLS,
|
||||||
|
*VCS_TOOLS,
|
||||||
|
*CHAT_TOOLS,
|
||||||
|
*TODO_TOOLS,
|
||||||
|
*NET_TOOLS,
|
||||||
|
):
|
||||||
if catalog.get(definition.name) is None:
|
if catalog.get(definition.name) is None:
|
||||||
catalog.register(definition, source=_BUILTIN_SOURCE)
|
catalog.register(definition, source=_BUILTIN_SOURCE)
|
||||||
|
|
||||||
@@ -235,6 +243,7 @@ def register_builtin_tools(*, force: bool = False) -> ToolCatalog:
|
|||||||
# Group packages pull their leaf modules (FILE_TOOLS etc. still exported).
|
# Group packages pull their leaf modules (FILE_TOOLS etc. still exported).
|
||||||
import plyngent.tools.chat as _chat_tools
|
import plyngent.tools.chat as _chat_tools
|
||||||
import plyngent.tools.file as _file_tools
|
import plyngent.tools.file as _file_tools
|
||||||
|
import plyngent.tools.net as _net_tools
|
||||||
import plyngent.tools.process as _process_tools
|
import plyngent.tools.process as _process_tools
|
||||||
import plyngent.tools.temp_workspace as _temp_workspace_tools
|
import plyngent.tools.temp_workspace as _temp_workspace_tools
|
||||||
import plyngent.tools.todo as _todo_tools
|
import plyngent.tools.todo as _todo_tools
|
||||||
@@ -243,6 +252,7 @@ def register_builtin_tools(*, force: bool = False) -> ToolCatalog:
|
|||||||
_ = (
|
_ = (
|
||||||
_chat_tools,
|
_chat_tools,
|
||||||
_file_tools,
|
_file_tools,
|
||||||
|
_net_tools,
|
||||||
_process_tools,
|
_process_tools,
|
||||||
_temp_workspace_tools,
|
_temp_workspace_tools,
|
||||||
_todo_tools,
|
_todo_tools,
|
||||||
|
|||||||
@@ -233,12 +233,26 @@ def _open_pty_reason(args: Mapping[str, object]) -> str | None:
|
|||||||
return _shell_or_dash_c_reason(argv, via="open_pty")
|
return _shell_or_dash_c_reason(argv, via="open_pty")
|
||||||
|
|
||||||
|
|
||||||
|
def _fetch_reason(args: Mapping[str, object]) -> str | None:
|
||||||
|
"""Soft-confirm cleartext HTTP and mutating fetch methods (not private-host policy)."""
|
||||||
|
from plyngent.tools.net.fetch import fetch_soft_reason
|
||||||
|
|
||||||
|
method_obj = args.get("method", "GET")
|
||||||
|
url_obj = args.get("url", "")
|
||||||
|
body_obj = args.get("body")
|
||||||
|
method = method_obj if isinstance(method_obj, str) else "GET"
|
||||||
|
url = url_obj if isinstance(url_obj, str) else ""
|
||||||
|
body = body_obj if isinstance(body_obj, str) else None
|
||||||
|
return fetch_soft_reason(method, url, body)
|
||||||
|
|
||||||
|
|
||||||
def classify_danger(name: str, args: Mapping[str, object]) -> str | None: # noqa: PLR0911
|
def classify_danger(name: str, args: Mapping[str, object]) -> str | None: # noqa: PLR0911
|
||||||
"""Return a short reason if ``name``/``args`` need user confirm, else ``None``.
|
"""Return a short reason if ``name``/``args`` need user confirm, else ``None``.
|
||||||
|
|
||||||
Hard denylists (paths/commands) still raise independently. This only covers
|
Hard denylists (paths/commands) still raise independently. This only covers
|
||||||
soft confirms for mutating tools and risky shell/REPL launches
|
soft confirms for mutating tools and risky shell/REPL launches
|
||||||
(interactive shells and ``python -c`` / ``bash -c`` one-liners).
|
(interactive shells and ``python -c`` / ``bash -c`` one-liners).
|
||||||
|
Private/loopback fetch targets use a separate policy grant (not YOLO).
|
||||||
"""
|
"""
|
||||||
if name == "delete_path":
|
if name == "delete_path":
|
||||||
return _delete_path_reason(args)
|
return _delete_path_reason(args)
|
||||||
@@ -254,4 +268,6 @@ def classify_danger(name: str, args: Mapping[str, object]) -> str | None: # noq
|
|||||||
return _run_command_batch_reason(args)
|
return _run_command_batch_reason(args)
|
||||||
if name == "open_pty":
|
if name == "open_pty":
|
||||||
return _open_pty_reason(args)
|
return _open_pty_reason(args)
|
||||||
|
if name == "fetch":
|
||||||
|
return _fetch_reason(args)
|
||||||
return None
|
return None
|
||||||
|
|||||||
@@ -0,0 +1,8 @@
|
|||||||
|
from .fetch import fetch as fetch
|
||||||
|
from .grants import clear_private_grants as clear_private_grants
|
||||||
|
from .grants import grant_private_host as grant_private_host
|
||||||
|
from .grants import set_fetch_policy_confirm_hook as set_fetch_policy_confirm_hook
|
||||||
|
|
||||||
|
NET_TOOLS = [
|
||||||
|
fetch,
|
||||||
|
]
|
||||||
@@ -0,0 +1,367 @@
|
|||||||
|
"""Async HTTP fetch helper (niquests) with manual redirects and body caps."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
import niquests
|
||||||
|
|
||||||
|
from plyngent.tools.net.grants import ensure_host_allowed
|
||||||
|
from plyngent.tools.net.policy import (
|
||||||
|
DEFAULT_MAX_BYTES,
|
||||||
|
DEFAULT_MAX_REDIRECTS,
|
||||||
|
DEFAULT_TIMEOUT_SECONDS,
|
||||||
|
FetchPolicyError,
|
||||||
|
is_https_to_http_downgrade,
|
||||||
|
parse_fetch_url,
|
||||||
|
resolve_redirect_url,
|
||||||
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Mapping
|
||||||
|
|
||||||
|
_REDIRECT_STATUS = frozenset({301, 302, 303, 307, 308})
|
||||||
|
_TEXT_CONTENT_HINTS = (
|
||||||
|
"text/",
|
||||||
|
"application/json",
|
||||||
|
"application/xml",
|
||||||
|
"application/javascript",
|
||||||
|
"application/xhtml",
|
||||||
|
"application/x-www-form-urlencoded",
|
||||||
|
"application/graphql",
|
||||||
|
"application/problem+json",
|
||||||
|
"application/ld+json",
|
||||||
|
"+json",
|
||||||
|
"+xml",
|
||||||
|
)
|
||||||
|
_BINARY_CONTENT_HINTS = ("octet-stream", "image/", "audio/", "video/")
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(slots=True)
|
||||||
|
class FetchResult:
|
||||||
|
status: int
|
||||||
|
final_url: str
|
||||||
|
content_type: str
|
||||||
|
body_text: str
|
||||||
|
body_bytes: int
|
||||||
|
truncated: bool
|
||||||
|
redirects: int
|
||||||
|
method: str
|
||||||
|
security: str
|
||||||
|
warnings: list[str] = field(default_factory=list)
|
||||||
|
body_kind: str = "text" # text | binary | empty
|
||||||
|
|
||||||
|
|
||||||
|
def _content_type(headers: Mapping[str, str] | object | None) -> str:
|
||||||
|
if headers is None:
|
||||||
|
return ""
|
||||||
|
get = getattr(headers, "get", None)
|
||||||
|
if not callable(get):
|
||||||
|
return ""
|
||||||
|
value = get("Content-Type") or get("content-type") or ""
|
||||||
|
return str(value)
|
||||||
|
|
||||||
|
|
||||||
|
def _looks_text(content_type: str, sample: bytes) -> bool:
|
||||||
|
lower = content_type.lower()
|
||||||
|
if any(hint in lower for hint in _TEXT_CONTENT_HINTS):
|
||||||
|
return True
|
||||||
|
if not sample:
|
||||||
|
return True
|
||||||
|
if b"\x00" in sample[:512]:
|
||||||
|
return False
|
||||||
|
return not any(hint in lower for hint in _BINARY_CONTENT_HINTS)
|
||||||
|
|
||||||
|
|
||||||
|
def _decode_body(data: bytes, content_type: str) -> str:
|
||||||
|
charset = "utf-8"
|
||||||
|
lower = content_type.lower()
|
||||||
|
if "charset=" in lower:
|
||||||
|
part = lower.split("charset=", 1)[1]
|
||||||
|
charset = part.split(";")[0].strip().strip('"') or "utf-8"
|
||||||
|
try:
|
||||||
|
return data.decode(charset, errors="replace")
|
||||||
|
except LookupError:
|
||||||
|
return data.decode("utf-8", errors="replace")
|
||||||
|
|
||||||
|
|
||||||
|
def _security_label(*, original_scheme: str, final_scheme: str, warnings: list[str]) -> str:
|
||||||
|
if any("https-to-http" in w for w in warnings):
|
||||||
|
return "https-to-http-redirect"
|
||||||
|
if final_scheme == "http" or original_scheme == "http":
|
||||||
|
return "cleartext-http"
|
||||||
|
return "https"
|
||||||
|
|
||||||
|
|
||||||
|
def _location_header(resp: object) -> str | None:
|
||||||
|
resp_headers = getattr(resp, "headers", None)
|
||||||
|
if resp_headers is None:
|
||||||
|
return None
|
||||||
|
value = resp_headers.get("Location") or resp_headers.get("location")
|
||||||
|
return None if value is None else str(value)
|
||||||
|
|
||||||
|
|
||||||
|
def _apply_redirect_method(status: int, method: str, body: bytes | None) -> tuple[str, bytes | None]:
|
||||||
|
if status in {302, 303} and method != "GET":
|
||||||
|
return "GET", None
|
||||||
|
return method, body
|
||||||
|
|
||||||
|
|
||||||
|
def _note_redirect_security(
|
||||||
|
*,
|
||||||
|
previous_scheme: str,
|
||||||
|
next_url: str,
|
||||||
|
next_scheme: str,
|
||||||
|
hop: int,
|
||||||
|
allow_http_downgrade: bool,
|
||||||
|
warnings: list[str],
|
||||||
|
) -> None:
|
||||||
|
if is_https_to_http_downgrade(previous_scheme, next_scheme):
|
||||||
|
if not allow_http_downgrade:
|
||||||
|
msg = f"blocked HTTPS→HTTP redirect to {next_url} (set allow_http_downgrade=true to override)"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
warnings.append(f"https-to-http-redirect hop {hop}: {next_url}")
|
||||||
|
elif next_scheme == "http":
|
||||||
|
warnings.append(f"cleartext HTTP at hop {hop}: {next_url}")
|
||||||
|
|
||||||
|
|
||||||
|
def _body_payload(data: bytes, content_type: str) -> tuple[str, str]:
|
||||||
|
"""Return (body_kind, body_text)."""
|
||||||
|
if not data:
|
||||||
|
return "empty", ""
|
||||||
|
if _looks_text(content_type, data):
|
||||||
|
return "text", _decode_body(data, content_type)
|
||||||
|
preview = data[:64].hex()
|
||||||
|
return (
|
||||||
|
"binary",
|
||||||
|
f"(binary body omitted; content_type={content_type!r}; bytes={len(data)}; hex_prefix={preview})",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _empty_redirect_result(
|
||||||
|
*,
|
||||||
|
status: int,
|
||||||
|
current_url: str,
|
||||||
|
content_type: str,
|
||||||
|
redirects: int,
|
||||||
|
method: str,
|
||||||
|
original_scheme: str,
|
||||||
|
scheme: str,
|
||||||
|
warnings: list[str],
|
||||||
|
) -> FetchResult:
|
||||||
|
return FetchResult(
|
||||||
|
status=status,
|
||||||
|
final_url=current_url,
|
||||||
|
content_type=content_type,
|
||||||
|
body_text="",
|
||||||
|
body_bytes=0,
|
||||||
|
truncated=False,
|
||||||
|
redirects=redirects,
|
||||||
|
method=method,
|
||||||
|
security=_security_label(
|
||||||
|
original_scheme=original_scheme,
|
||||||
|
final_scheme=scheme,
|
||||||
|
warnings=warnings,
|
||||||
|
),
|
||||||
|
warnings=warnings,
|
||||||
|
body_kind="empty",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def http_fetch(
|
||||||
|
*,
|
||||||
|
method: str,
|
||||||
|
url: str,
|
||||||
|
headers: Mapping[str, str],
|
||||||
|
body: bytes | None,
|
||||||
|
timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
|
||||||
|
max_bytes: int = DEFAULT_MAX_BYTES,
|
||||||
|
max_redirects: int = DEFAULT_MAX_REDIRECTS,
|
||||||
|
follow_redirects: bool = True,
|
||||||
|
allow_http_downgrade: bool = False,
|
||||||
|
) -> FetchResult:
|
||||||
|
"""Perform *method* on *url* with SSRF checks on each hop."""
|
||||||
|
if timeout_seconds <= 0:
|
||||||
|
msg = "timeout_seconds must be > 0"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
if max_bytes < 1:
|
||||||
|
msg = "max_bytes must be >= 1"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
if max_redirects < 0:
|
||||||
|
msg = "max_redirects must be >= 0"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
|
||||||
|
current_url = parse_fetch_url(url).url
|
||||||
|
original_scheme = parse_fetch_url(current_url).scheme
|
||||||
|
warnings: list[str] = []
|
||||||
|
redirects = 0
|
||||||
|
active_method = method
|
||||||
|
active_body = body
|
||||||
|
timeout = float(timeout_seconds)
|
||||||
|
|
||||||
|
async with niquests.AsyncSession() as session:
|
||||||
|
while True:
|
||||||
|
await ensure_host_allowed(current_url)
|
||||||
|
parsed = parse_fetch_url(current_url)
|
||||||
|
try:
|
||||||
|
resp = await session.request(
|
||||||
|
active_method,
|
||||||
|
current_url,
|
||||||
|
headers=dict(headers),
|
||||||
|
data=active_body,
|
||||||
|
timeout=timeout,
|
||||||
|
allow_redirects=False,
|
||||||
|
stream=True,
|
||||||
|
)
|
||||||
|
except niquests.RequestException as exc:
|
||||||
|
msg = f"HTTP request failed: {exc}"
|
||||||
|
raise FetchPolicyError(msg) from exc
|
||||||
|
|
||||||
|
status = int(getattr(resp, "status_code", 0) or 0)
|
||||||
|
content_type = _content_type(getattr(resp, "headers", {}) or {})
|
||||||
|
is_redirect = status in _REDIRECT_STATUS
|
||||||
|
|
||||||
|
if follow_redirects and is_redirect and redirects < max_redirects:
|
||||||
|
location = _location_header(resp)
|
||||||
|
await _close_response(resp)
|
||||||
|
if not location:
|
||||||
|
return _empty_redirect_result(
|
||||||
|
status=status,
|
||||||
|
current_url=current_url,
|
||||||
|
content_type=content_type,
|
||||||
|
redirects=redirects,
|
||||||
|
method=active_method,
|
||||||
|
original_scheme=original_scheme,
|
||||||
|
scheme=parsed.scheme,
|
||||||
|
warnings=warnings,
|
||||||
|
)
|
||||||
|
next_url = resolve_redirect_url(current_url, location)
|
||||||
|
next_parsed = parse_fetch_url(next_url)
|
||||||
|
_note_redirect_security(
|
||||||
|
previous_scheme=parsed.scheme,
|
||||||
|
next_url=next_url,
|
||||||
|
next_scheme=next_parsed.scheme,
|
||||||
|
hop=redirects + 1,
|
||||||
|
allow_http_downgrade=allow_http_downgrade,
|
||||||
|
warnings=warnings,
|
||||||
|
)
|
||||||
|
redirects += 1
|
||||||
|
current_url = next_url
|
||||||
|
active_method, active_body = _apply_redirect_method(status, active_method, active_body)
|
||||||
|
continue
|
||||||
|
|
||||||
|
if follow_redirects and is_redirect and redirects >= max_redirects:
|
||||||
|
await _close_response(resp)
|
||||||
|
msg = f"too many redirects (max {max_redirects})"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
|
||||||
|
data, truncated = await _read_capped(resp, max_bytes=max_bytes)
|
||||||
|
await _close_response(resp)
|
||||||
|
|
||||||
|
final_scheme = parse_fetch_url(current_url).scheme
|
||||||
|
if final_scheme == "http":
|
||||||
|
warnings.append(f"cleartext HTTP final_url={current_url}")
|
||||||
|
body_kind, body_text = _body_payload(data, content_type)
|
||||||
|
return FetchResult(
|
||||||
|
status=status,
|
||||||
|
final_url=current_url,
|
||||||
|
content_type=content_type,
|
||||||
|
body_text=body_text,
|
||||||
|
body_bytes=len(data),
|
||||||
|
truncated=truncated,
|
||||||
|
redirects=redirects,
|
||||||
|
method=active_method,
|
||||||
|
security=_security_label(
|
||||||
|
original_scheme=original_scheme,
|
||||||
|
final_scheme=final_scheme,
|
||||||
|
warnings=warnings,
|
||||||
|
),
|
||||||
|
warnings=list(dict.fromkeys(warnings)),
|
||||||
|
body_kind=body_kind,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _append_chunk(
|
||||||
|
chunks: list[bytes],
|
||||||
|
total: int,
|
||||||
|
data: bytes,
|
||||||
|
*,
|
||||||
|
max_bytes: int,
|
||||||
|
) -> tuple[int, bool]:
|
||||||
|
"""Append *data* under *max_bytes*; return (new_total, truncated)."""
|
||||||
|
if total + len(data) > max_bytes:
|
||||||
|
need = max_bytes - total
|
||||||
|
if need > 0:
|
||||||
|
chunks.append(data[:need])
|
||||||
|
total += need
|
||||||
|
return total, True
|
||||||
|
chunks.append(data)
|
||||||
|
return total + len(data), False
|
||||||
|
|
||||||
|
|
||||||
|
async def _consume_async_stream(stream: Any, *, max_bytes: int) -> tuple[bytes, bool]:
|
||||||
|
chunks: list[bytes] = []
|
||||||
|
total = 0
|
||||||
|
async for chunk in stream:
|
||||||
|
if not chunk:
|
||||||
|
continue
|
||||||
|
data = chunk if isinstance(chunk, bytes) else bytes(chunk)
|
||||||
|
total, truncated = _append_chunk(chunks, total, data, max_bytes=max_bytes)
|
||||||
|
if truncated:
|
||||||
|
return b"".join(chunks), True
|
||||||
|
return b"".join(chunks), False
|
||||||
|
|
||||||
|
|
||||||
|
def _consume_sync_stream(stream: Any, *, max_bytes: int) -> tuple[bytes, bool]:
|
||||||
|
chunks: list[bytes] = []
|
||||||
|
total = 0
|
||||||
|
for chunk in stream:
|
||||||
|
if not chunk:
|
||||||
|
continue
|
||||||
|
data = chunk if isinstance(chunk, bytes) else bytes(chunk)
|
||||||
|
total, truncated = _append_chunk(chunks, total, data, max_bytes=max_bytes)
|
||||||
|
if truncated:
|
||||||
|
return b"".join(chunks), True
|
||||||
|
return b"".join(chunks), False
|
||||||
|
|
||||||
|
|
||||||
|
async def _read_full_content(resp: object, *, max_bytes: int) -> tuple[bytes, bool]:
|
||||||
|
content = getattr(resp, "content", None)
|
||||||
|
if content is not None and hasattr(content, "__await__"):
|
||||||
|
raw = await content
|
||||||
|
data = raw if isinstance(raw, bytes) else bytes(raw)
|
||||||
|
elif isinstance(content, bytes):
|
||||||
|
data = content
|
||||||
|
else:
|
||||||
|
data = b""
|
||||||
|
if len(data) > max_bytes:
|
||||||
|
return data[:max_bytes], True
|
||||||
|
return data, False
|
||||||
|
|
||||||
|
|
||||||
|
async def _read_capped(resp: object, *, max_bytes: int) -> tuple[bytes, bool]:
|
||||||
|
"""Read at most *max_bytes* from a streamed response.
|
||||||
|
|
||||||
|
niquests ``AsyncResponse.iter_content`` is async and returns an async
|
||||||
|
generator after await; ``content`` is an async property (coroutine).
|
||||||
|
"""
|
||||||
|
iter_content = getattr(resp, "iter_content", None)
|
||||||
|
if callable(iter_content):
|
||||||
|
stream_obj: Any = iter_content(chunk_size=65_536)
|
||||||
|
if hasattr(stream_obj, "__await__"):
|
||||||
|
stream_obj = await stream_obj
|
||||||
|
if hasattr(stream_obj, "__aiter__"):
|
||||||
|
return await _consume_async_stream(stream_obj, max_bytes=max_bytes)
|
||||||
|
if hasattr(stream_obj, "__iter__"):
|
||||||
|
return _consume_sync_stream(stream_obj, max_bytes=max_bytes)
|
||||||
|
return await _read_full_content(resp, max_bytes=max_bytes)
|
||||||
|
|
||||||
|
|
||||||
|
async def _close_response(resp: object) -> None:
|
||||||
|
close = getattr(resp, "close", None)
|
||||||
|
if close is None:
|
||||||
|
return
|
||||||
|
result: Any = close()
|
||||||
|
if hasattr(result, "__await__"):
|
||||||
|
await result
|
||||||
@@ -0,0 +1,140 @@
|
|||||||
|
"""Model-facing ``fetch`` tool (HTTP GET/POST/PUT/DELETE)."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from plyngent.agent import ToolTag, tool
|
||||||
|
from plyngent.tools.net.client import http_fetch
|
||||||
|
from plyngent.tools.net.policy import (
|
||||||
|
DEFAULT_MAX_BODY_CHARS_IN,
|
||||||
|
DEFAULT_MAX_BYTES,
|
||||||
|
DEFAULT_MAX_CHARS,
|
||||||
|
DEFAULT_MAX_REDIRECTS,
|
||||||
|
DEFAULT_TIMEOUT_SECONDS,
|
||||||
|
FetchPolicyError,
|
||||||
|
normalize_method,
|
||||||
|
normalize_request_headers,
|
||||||
|
parse_fetch_url,
|
||||||
|
soft_confirm_reason,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def format_fetch_result(
|
||||||
|
*,
|
||||||
|
status: int,
|
||||||
|
final_url: str,
|
||||||
|
content_type: str,
|
||||||
|
body_text: str,
|
||||||
|
body_bytes: int,
|
||||||
|
truncated: bool,
|
||||||
|
redirects: int,
|
||||||
|
method: str,
|
||||||
|
security: str,
|
||||||
|
warnings: list[str],
|
||||||
|
body_kind: str,
|
||||||
|
max_chars: int,
|
||||||
|
) -> 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
|
||||||
|
warn_line = "; ".join(warnings) if warnings else ""
|
||||||
|
parts = [
|
||||||
|
f"status={status}",
|
||||||
|
f"method={method}",
|
||||||
|
f"final_url={final_url}",
|
||||||
|
f"content_type={content_type}",
|
||||||
|
f"body_kind={body_kind}",
|
||||||
|
f"bytes={body_bytes}",
|
||||||
|
f"truncated={'true' if truncated or char_truncated else 'false'}",
|
||||||
|
f"redirects={redirects}",
|
||||||
|
f"security={security}",
|
||||||
|
]
|
||||||
|
if warn_line:
|
||||||
|
parts.append(f"warnings={warn_line}")
|
||||||
|
parts.append("--- body ---")
|
||||||
|
parts.append(text)
|
||||||
|
return "\n".join(parts)
|
||||||
|
|
||||||
|
|
||||||
|
@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO | ToolTag.TRUSTABLE)
|
||||||
|
async def fetch(
|
||||||
|
url: str,
|
||||||
|
*,
|
||||||
|
method: str = "GET",
|
||||||
|
headers: dict[str, str] | None = None,
|
||||||
|
body: str | None = None,
|
||||||
|
user_agent: str | None = None,
|
||||||
|
timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
|
||||||
|
max_bytes: int = DEFAULT_MAX_BYTES,
|
||||||
|
max_chars: int = DEFAULT_MAX_CHARS,
|
||||||
|
max_redirects: int = DEFAULT_MAX_REDIRECTS,
|
||||||
|
follow_redirects: bool = True,
|
||||||
|
allow_http_downgrade: bool = False,
|
||||||
|
) -> str:
|
||||||
|
"""HTTP request (GET/POST/PUT/DELETE); return status, metadata, and truncated body.
|
||||||
|
|
||||||
|
Use for public docs/APIs or (after human policy allow) local/LAN servers.
|
||||||
|
``user_agent`` sets User-Agent when provided and takes precedence over a
|
||||||
|
User-Agent entry in ``headers``. A headers-only User-Agent is kept as-is.
|
||||||
|
When both omit UA, a small default is used.
|
||||||
|
|
||||||
|
Private/loopback/link-local hosts require an explicit human policy grant
|
||||||
|
(not skipped by YOLO). HTTPS→HTTP redirects are blocked unless
|
||||||
|
``allow_http_downgrade`` is true.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
verb = normalize_method(method)
|
||||||
|
parsed = parse_fetch_url(url)
|
||||||
|
hdrs = normalize_request_headers(headers, user_agent=user_agent)
|
||||||
|
body_bytes: bytes | None = None
|
||||||
|
if body is not None:
|
||||||
|
if len(body) > DEFAULT_MAX_BODY_CHARS_IN:
|
||||||
|
return f"error: request body too large ({len(body)} chars; max {DEFAULT_MAX_BODY_CHARS_IN})"
|
||||||
|
body_bytes = body.encode("utf-8")
|
||||||
|
result = await http_fetch(
|
||||||
|
method=verb,
|
||||||
|
url=parsed.url,
|
||||||
|
headers=hdrs,
|
||||||
|
body=body_bytes,
|
||||||
|
timeout_seconds=timeout_seconds,
|
||||||
|
max_bytes=max_bytes,
|
||||||
|
max_redirects=max_redirects,
|
||||||
|
follow_redirects=follow_redirects,
|
||||||
|
allow_http_downgrade=allow_http_downgrade,
|
||||||
|
)
|
||||||
|
except FetchPolicyError as exc:
|
||||||
|
return f"error: {exc}"
|
||||||
|
except Exception as exc: # noqa: BLE001 — tool surface returns error text
|
||||||
|
return f"error: fetch failed: {exc}"
|
||||||
|
|
||||||
|
return format_fetch_result(
|
||||||
|
status=result.status,
|
||||||
|
final_url=result.final_url,
|
||||||
|
content_type=result.content_type,
|
||||||
|
body_text=result.body_text,
|
||||||
|
body_bytes=result.body_bytes,
|
||||||
|
truncated=result.truncated,
|
||||||
|
redirects=result.redirects,
|
||||||
|
method=result.method,
|
||||||
|
security=result.security,
|
||||||
|
warnings=result.warnings,
|
||||||
|
body_kind=result.body_kind,
|
||||||
|
max_chars=max_chars,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# Re-export for danger classifier without importing client.
|
||||||
|
def fetch_soft_reason(method: str, url: str, body: str | None) -> str | None:
|
||||||
|
try:
|
||||||
|
verb = normalize_method(method)
|
||||||
|
parsed = parse_fetch_url(url)
|
||||||
|
except FetchPolicyError:
|
||||||
|
return None
|
||||||
|
return soft_confirm_reason(
|
||||||
|
method=verb,
|
||||||
|
url=parsed.url,
|
||||||
|
scheme=parsed.scheme,
|
||||||
|
body_present=bool(body),
|
||||||
|
)
|
||||||
@@ -0,0 +1,127 @@
|
|||||||
|
"""Instance-scoped private-fetch host grants (not YOLO-skippable)."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING, cast
|
||||||
|
|
||||||
|
from plyngent.tools.net.policy import FetchPolicyError, HostClass, assess_host, grant_key, parse_fetch_url
|
||||||
|
from plyngent.tools.workspace import WorkspaceError, require_bound_instance
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Callable, Sequence
|
||||||
|
|
||||||
|
from plyngent.tools.context import InstanceState
|
||||||
|
|
||||||
|
# (host, port, url, timeout_seconds) -> True to allow for this process/instance.
|
||||||
|
type FetchPolicyConfirmHook = Callable[[str, int, str, float], bool]
|
||||||
|
|
||||||
|
|
||||||
|
def get_private_grants(instance: InstanceState | None = None) -> set[str]:
|
||||||
|
"""Return the live grant key set (host:port strings)."""
|
||||||
|
inst = instance if instance is not None else require_bound_instance()
|
||||||
|
bag_obj = inst.extras.get("fetch_private_grants")
|
||||||
|
if isinstance(bag_obj, set):
|
||||||
|
return cast("set[str]", bag_obj)
|
||||||
|
empty: set[str] = set()
|
||||||
|
inst.extras["fetch_private_grants"] = empty
|
||||||
|
return empty
|
||||||
|
|
||||||
|
|
||||||
|
def grant_private_host(host: str, port: int, *, instance: InstanceState | None = None) -> str:
|
||||||
|
"""Record a private host:port grant; returns the grant key."""
|
||||||
|
key = grant_key(host, port)
|
||||||
|
get_private_grants(instance).add(key)
|
||||||
|
return key
|
||||||
|
|
||||||
|
|
||||||
|
def clear_private_grants(*, instance: InstanceState | None = None) -> None:
|
||||||
|
get_private_grants(instance).clear()
|
||||||
|
|
||||||
|
|
||||||
|
def has_private_grant(host: str, port: int, *, instance: InstanceState | None = None) -> bool:
|
||||||
|
return grant_key(host, port) in get_private_grants(instance)
|
||||||
|
|
||||||
|
|
||||||
|
def get_fetch_policy_confirm_hook(instance: InstanceState | None = None) -> FetchPolicyConfirmHook | None:
|
||||||
|
inst = instance if instance is not None else require_bound_instance()
|
||||||
|
hook = inst.extras.get("fetch_policy_confirm_hook")
|
||||||
|
if hook is None or not callable(hook):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def _as_hook(host: str, port: int, url: str, timeout_seconds: float) -> bool:
|
||||||
|
return bool(hook(host, port, url, timeout_seconds))
|
||||||
|
|
||||||
|
return _as_hook
|
||||||
|
|
||||||
|
|
||||||
|
def set_fetch_policy_confirm_hook(
|
||||||
|
hook: FetchPolicyConfirmHook | None,
|
||||||
|
*,
|
||||||
|
instance: InstanceState | None = None,
|
||||||
|
) -> None:
|
||||||
|
inst = instance if instance is not None else require_bound_instance()
|
||||||
|
inst.extras["fetch_policy_confirm_hook"] = hook
|
||||||
|
|
||||||
|
|
||||||
|
async def ensure_host_allowed(
|
||||||
|
url: str,
|
||||||
|
*,
|
||||||
|
policy_timeout_seconds: float | None = None,
|
||||||
|
instance: InstanceState | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Raise :class:`FetchPolicyError` if *url*'s host may not be contacted.
|
||||||
|
|
||||||
|
Public hosts pass. Forbidden (metadata) always fail. Private/loopback
|
||||||
|
requires an instance grant or a successful policy confirm hook (never YOLO).
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
inst = instance if instance is not None else require_bound_instance()
|
||||||
|
except WorkspaceError as exc:
|
||||||
|
msg = f"instance state is not bound for fetch policy: {exc}"
|
||||||
|
raise FetchPolicyError(msg) from exc
|
||||||
|
|
||||||
|
parsed = parse_fetch_url(url)
|
||||||
|
assessment = await assess_host(parsed.host, parsed.port)
|
||||||
|
if assessment.classification is HostClass.PUBLIC:
|
||||||
|
return
|
||||||
|
if assessment.classification is HostClass.FORBIDDEN:
|
||||||
|
msg = f"fetch blocked forbidden host {parsed.host!r} (addresses: {', '.join(assessment.addresses)})"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
|
||||||
|
# PRIVATE
|
||||||
|
if has_private_grant(parsed.host, parsed.port, instance=inst):
|
||||||
|
return
|
||||||
|
|
||||||
|
from plyngent.tools.workspace import get_policy_confirm_timeout
|
||||||
|
|
||||||
|
timeout = (
|
||||||
|
float(policy_timeout_seconds) if policy_timeout_seconds is not None else float(get_policy_confirm_timeout())
|
||||||
|
)
|
||||||
|
hook = get_fetch_policy_confirm_hook(inst)
|
||||||
|
if hook is None:
|
||||||
|
msg = (
|
||||||
|
f"fetch blocked private/loopback host {parsed.host}:{parsed.port} "
|
||||||
|
f"({assessment.reason}; no policy confirm hook / non-interactive deny)"
|
||||||
|
)
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
try:
|
||||||
|
allowed = bool(hook(parsed.host, parsed.port, parsed.url, timeout))
|
||||||
|
except Exception as exc:
|
||||||
|
msg = f"fetch blocked private/loopback host {parsed.host}:{parsed.port} (confirm failed: {exc})"
|
||||||
|
raise FetchPolicyError(msg) from exc
|
||||||
|
if not allowed:
|
||||||
|
msg = (
|
||||||
|
f"fetch blocked private/loopback host {parsed.host}:{parsed.port} "
|
||||||
|
f"(user declined or timed out after {timeout:g}s; not skipped by YOLO)"
|
||||||
|
)
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
_ = grant_private_host(parsed.host, parsed.port, instance=inst)
|
||||||
|
|
||||||
|
|
||||||
|
def format_grant_preview(host: str, port: int, url: str) -> Sequence[str]:
|
||||||
|
return (
|
||||||
|
f"host: {host}",
|
||||||
|
f"port: {port}",
|
||||||
|
f"url: {url}",
|
||||||
|
f"grant_key: {grant_key(host, port)}",
|
||||||
|
)
|
||||||
@@ -0,0 +1,286 @@
|
|||||||
|
"""URL / host policy for the fetch tool (SSRF, private grants, cleartext)."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import ipaddress
|
||||||
|
import socket
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from enum import Enum
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
from urllib.parse import urljoin, urlparse, urlunparse
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Mapping, Sequence
|
||||||
|
|
||||||
|
ALLOWED_METHODS: frozenset[str] = frozenset({"GET", "POST", "PUT", "DELETE"})
|
||||||
|
ALLOWED_SCHEMES: frozenset[str] = frozenset({"http", "https"})
|
||||||
|
|
||||||
|
# Hop-by-hop / identity headers the model must not set (User-Agent is allowed).
|
||||||
|
_FORBIDDEN_REQUEST_HEADERS: frozenset[str] = frozenset(
|
||||||
|
{
|
||||||
|
"host",
|
||||||
|
"content-length",
|
||||||
|
"transfer-encoding",
|
||||||
|
"connection",
|
||||||
|
"keep-alive",
|
||||||
|
"upgrade",
|
||||||
|
"te",
|
||||||
|
"trailer",
|
||||||
|
"proxy-authorization",
|
||||||
|
"proxy-authenticate",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
DEFAULT_USER_AGENT = "plyngent-fetch/0.2"
|
||||||
|
DEFAULT_MAX_REDIRECTS = 5
|
||||||
|
DEFAULT_TIMEOUT_SECONDS = 30.0
|
||||||
|
DEFAULT_MAX_BYTES = 1_000_000
|
||||||
|
DEFAULT_MAX_CHARS = 32_000
|
||||||
|
DEFAULT_MAX_BODY_CHARS_IN = 256_000 # request body size limit (tool arg)
|
||||||
|
|
||||||
|
# Cloud / link-local style targets never granted via policy UI.
|
||||||
|
_NEVER_ALLOW_NETWORKS: tuple[ipaddress.IPv4Network | ipaddress.IPv6Network, ...] = (
|
||||||
|
ipaddress.ip_network("169.254.169.254/32"), # AWS/GCP-style metadata (IPv4)
|
||||||
|
ipaddress.ip_network("fd00:ec2::254/128"), # AWS IMDS IPv6
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class HostClass(Enum):
|
||||||
|
PUBLIC = "public"
|
||||||
|
PRIVATE = "private" # loopback, RFC1918, ULA, link-local, etc. (grantable)
|
||||||
|
FORBIDDEN = "forbidden" # metadata / never allow
|
||||||
|
|
||||||
|
|
||||||
|
class FetchPolicyError(ValueError):
|
||||||
|
"""Hard fetch policy violation (returned to the model as error text)."""
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class ParsedFetchUrl:
|
||||||
|
"""Normalized URL pieces used for policy and the HTTP client."""
|
||||||
|
|
||||||
|
url: str
|
||||||
|
scheme: str
|
||||||
|
host: str
|
||||||
|
port: int
|
||||||
|
path_query: str # path + optional ?query (no fragment)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class HostAssessment:
|
||||||
|
host: str
|
||||||
|
port: int
|
||||||
|
classification: HostClass
|
||||||
|
addresses: tuple[str, ...]
|
||||||
|
reason: str
|
||||||
|
|
||||||
|
|
||||||
|
def normalize_method(method: str) -> str:
|
||||||
|
upper = method.strip().upper()
|
||||||
|
if upper not in ALLOWED_METHODS:
|
||||||
|
allowed = ", ".join(sorted(ALLOWED_METHODS))
|
||||||
|
msg = f"method must be one of {allowed}; got {method!r}"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
return upper
|
||||||
|
|
||||||
|
|
||||||
|
def parse_fetch_url(url: str) -> ParsedFetchUrl:
|
||||||
|
raw = url.strip()
|
||||||
|
if not raw:
|
||||||
|
msg = "url must not be empty"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
parsed = urlparse(raw)
|
||||||
|
scheme = (parsed.scheme or "").lower()
|
||||||
|
if scheme not in ALLOWED_SCHEMES:
|
||||||
|
msg = f"url scheme must be http or https; got {scheme or '(none)'!r}"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
if not parsed.hostname:
|
||||||
|
msg = "url must include a hostname"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
host = parsed.hostname
|
||||||
|
# urlparse keeps brackets out of hostname for IPv6.
|
||||||
|
port = parsed.port
|
||||||
|
if port is None:
|
||||||
|
port = 443 if scheme == "https" else 80
|
||||||
|
# Drop fragment; rebuild without params quirks.
|
||||||
|
path = parsed.path or "/"
|
||||||
|
query = parsed.query
|
||||||
|
path_query = path if not query else f"{path}?{query}"
|
||||||
|
# Prefer normalized form (no fragment).
|
||||||
|
normalized = urlunparse((scheme, parsed.netloc, path, "", query, ""))
|
||||||
|
return ParsedFetchUrl(
|
||||||
|
url=normalized,
|
||||||
|
scheme=scheme,
|
||||||
|
host=host,
|
||||||
|
port=port,
|
||||||
|
path_query=path_query,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_redirect_url(current: str, location: str) -> str:
|
||||||
|
"""Resolve a redirect Location against *current* and re-parse for safety."""
|
||||||
|
if not location or not location.strip():
|
||||||
|
msg = "redirect Location is empty"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
joined = urljoin(current, location.strip())
|
||||||
|
return parse_fetch_url(joined).url
|
||||||
|
|
||||||
|
|
||||||
|
def _ip_classification(address: str) -> HostClass:
|
||||||
|
try:
|
||||||
|
ip = ipaddress.ip_address(address)
|
||||||
|
except ValueError:
|
||||||
|
return HostClass.FORBIDDEN
|
||||||
|
for network in _NEVER_ALLOW_NETWORKS:
|
||||||
|
if ip in network:
|
||||||
|
return HostClass.FORBIDDEN
|
||||||
|
# IPv6 unique-local / link-local / loopback / unspecified
|
||||||
|
if ip.is_loopback or ip.is_link_local or ip.is_private or ip.is_reserved or ip.is_multicast or ip.is_unspecified:
|
||||||
|
return HostClass.PRIVATE
|
||||||
|
return HostClass.PUBLIC
|
||||||
|
|
||||||
|
|
||||||
|
def classify_ip_strings(addresses: Sequence[str]) -> HostClass:
|
||||||
|
"""Worst-class wins: FORBIDDEN > PRIVATE > PUBLIC."""
|
||||||
|
if not addresses:
|
||||||
|
return HostClass.FORBIDDEN
|
||||||
|
worst = HostClass.PUBLIC
|
||||||
|
for addr in addresses:
|
||||||
|
kind = _ip_classification(addr)
|
||||||
|
if kind is HostClass.FORBIDDEN:
|
||||||
|
return HostClass.FORBIDDEN
|
||||||
|
if kind is HostClass.PRIVATE:
|
||||||
|
worst = HostClass.PRIVATE
|
||||||
|
return worst
|
||||||
|
|
||||||
|
|
||||||
|
async def resolve_host_addresses(host: str) -> tuple[str, ...]:
|
||||||
|
"""Resolve *host* via the running loop's ``getaddrinfo`` (async)."""
|
||||||
|
# Literal IPs: no DNS.
|
||||||
|
try:
|
||||||
|
ip = ipaddress.ip_address(host)
|
||||||
|
except ValueError:
|
||||||
|
ip = None
|
||||||
|
if ip is not None:
|
||||||
|
return (str(ip),)
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
try:
|
||||||
|
infos = await loop.getaddrinfo(
|
||||||
|
host,
|
||||||
|
None,
|
||||||
|
type=socket.SOCK_STREAM,
|
||||||
|
proto=socket.IPPROTO_TCP,
|
||||||
|
)
|
||||||
|
except socket.gaierror as exc:
|
||||||
|
msg = f"DNS resolution failed for {host!r}: {exc}"
|
||||||
|
raise FetchPolicyError(msg) from exc
|
||||||
|
addrs: list[str] = []
|
||||||
|
seen: set[str] = set()
|
||||||
|
for info in infos:
|
||||||
|
sockaddr = info[4]
|
||||||
|
if not sockaddr:
|
||||||
|
continue
|
||||||
|
addr = str(sockaddr[0])
|
||||||
|
if addr not in seen:
|
||||||
|
seen.add(addr)
|
||||||
|
addrs.append(addr)
|
||||||
|
if not addrs:
|
||||||
|
msg = f"DNS resolution returned no addresses for {host!r}"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
return tuple(addrs)
|
||||||
|
|
||||||
|
|
||||||
|
async def assess_host(host: str, port: int) -> HostAssessment:
|
||||||
|
"""Classify *host* after resolution (literal IP or DNS)."""
|
||||||
|
addresses = await resolve_host_addresses(host)
|
||||||
|
classification = classify_ip_strings(addresses)
|
||||||
|
if classification is HostClass.FORBIDDEN:
|
||||||
|
reason = f"host {host!r} resolves to a forbidden address ({', '.join(addresses)})"
|
||||||
|
elif classification is HostClass.PRIVATE:
|
||||||
|
reason = f"host {host!r} is loopback/private/link-local ({', '.join(addresses)})"
|
||||||
|
else:
|
||||||
|
reason = f"host {host!r} is public ({', '.join(addresses)})"
|
||||||
|
return HostAssessment(
|
||||||
|
host=host,
|
||||||
|
port=port,
|
||||||
|
classification=classification,
|
||||||
|
addresses=addresses,
|
||||||
|
reason=reason,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def grant_key(host: str, port: int) -> str:
|
||||||
|
"""Stable key for instance-scoped private fetch grants."""
|
||||||
|
return f"{host.lower()}:{port}"
|
||||||
|
|
||||||
|
|
||||||
|
def normalize_request_headers(
|
||||||
|
headers: Mapping[str, str] | None,
|
||||||
|
*,
|
||||||
|
user_agent: str | None = None,
|
||||||
|
) -> dict[str, str]:
|
||||||
|
"""Build outbound headers.
|
||||||
|
|
||||||
|
* Non-empty *user_agent* wins and replaces any User-Agent in *headers*.
|
||||||
|
* Else keep a User-Agent already present in *headers* (never replaced by default).
|
||||||
|
* When both omit UA, set :data:`DEFAULT_USER_AGENT`.
|
||||||
|
* Forbidden hop-by-hop headers raise :class:`FetchPolicyError`.
|
||||||
|
"""
|
||||||
|
out: dict[str, str] = {}
|
||||||
|
if headers:
|
||||||
|
for raw_key, raw_value in headers.items():
|
||||||
|
key = raw_key.strip()
|
||||||
|
if not key:
|
||||||
|
msg = "header name must not be empty"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
lower = key.lower()
|
||||||
|
if lower in _FORBIDDEN_REQUEST_HEADERS:
|
||||||
|
msg = f"header {key!r} is not allowed"
|
||||||
|
raise FetchPolicyError(msg)
|
||||||
|
out[key] = raw_value
|
||||||
|
|
||||||
|
def _has_user_agent(mapping: Mapping[str, str]) -> bool:
|
||||||
|
return any(k.lower() == "user-agent" for k in mapping)
|
||||||
|
|
||||||
|
# Tool-call User-Agent is never replaced by a host default. Prefer the
|
||||||
|
# dedicated ``user_agent`` argument when non-empty; else keep headers; else default.
|
||||||
|
if user_agent is not None and user_agent.strip():
|
||||||
|
out = {key: value for key, value in out.items() if key.lower() != "user-agent"}
|
||||||
|
out["User-Agent"] = user_agent.strip()
|
||||||
|
elif not _has_user_agent(out):
|
||||||
|
out["User-Agent"] = DEFAULT_USER_AGENT
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def is_https_to_http_downgrade(previous_scheme: str, next_scheme: str) -> bool:
|
||||||
|
return previous_scheme.lower() == "https" and next_scheme.lower() == "http"
|
||||||
|
|
||||||
|
|
||||||
|
def soft_confirm_reason(
|
||||||
|
*,
|
||||||
|
method: str,
|
||||||
|
url: str,
|
||||||
|
scheme: str,
|
||||||
|
body_present: bool,
|
||||||
|
) -> str | None:
|
||||||
|
"""Soft-confirm (YOLO-eligible only when reason is None or caller tags allow).
|
||||||
|
|
||||||
|
Public cleartext HTTP and any mutating method get a soft reason.
|
||||||
|
HTTPS GET/HEAD-like GET with no body: no soft reason (still subject to caps).
|
||||||
|
"""
|
||||||
|
parts: list[str] = [f"fetch: {method} {url}"]
|
||||||
|
risky = False
|
||||||
|
if scheme.lower() == "http":
|
||||||
|
parts.append("cleartext HTTP (not HTTPS)")
|
||||||
|
risky = True
|
||||||
|
if method in {"POST", "PUT", "DELETE"}:
|
||||||
|
parts.append(f"mutating method {method}")
|
||||||
|
if body_present:
|
||||||
|
parts.append("request body present")
|
||||||
|
risky = True
|
||||||
|
if not risky:
|
||||||
|
return None
|
||||||
|
return "\n ".join(parts)
|
||||||
@@ -0,0 +1,82 @@
|
|||||||
|
"""Directive checkpoint bands and append-only inject."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from plyngent.agent.directive_checkpoint import (
|
||||||
|
bands_to_fire,
|
||||||
|
inject_directive_checkpoints,
|
||||||
|
parse_checkpoint_bands,
|
||||||
|
token_band,
|
||||||
|
)
|
||||||
|
from plyngent.lmproto.openai_compatible.model import (
|
||||||
|
AnyChatMessage,
|
||||||
|
DeveloperChatMessage,
|
||||||
|
UserChatMessage,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_token_band() -> None:
|
||||||
|
assert token_band(0, 100_000) == 0
|
||||||
|
assert token_band(99_999, 100_000) == 0
|
||||||
|
assert token_band(100_000, 100_000) == 1
|
||||||
|
assert token_band(250_000, 100_000) == 2
|
||||||
|
assert token_band(100_000, 0) == 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_bands_to_fire_fill_gaps() -> None:
|
||||||
|
assert bands_to_fire(last_fired_band=0, current_band=0) == []
|
||||||
|
assert bands_to_fire(last_fired_band=0, current_band=1) == [1]
|
||||||
|
assert bands_to_fire(last_fired_band=0, current_band=2) == [1, 2]
|
||||||
|
assert bands_to_fire(last_fired_band=2, current_band=2) == []
|
||||||
|
assert bands_to_fire(last_fired_band=1, current_band=3) == [2, 3]
|
||||||
|
|
||||||
|
|
||||||
|
def test_inject_append_only_and_parse() -> None:
|
||||||
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
||||||
|
band, appended = inject_directive_checkpoints(
|
||||||
|
messages,
|
||||||
|
prompt_tokens=100_000,
|
||||||
|
source="api",
|
||||||
|
interval=100_000,
|
||||||
|
last_fired_band=0,
|
||||||
|
)
|
||||||
|
assert band == 1
|
||||||
|
assert len(appended) == 1
|
||||||
|
assert isinstance(messages[-1], DeveloperChatMessage)
|
||||||
|
assert "band=1" in messages[-1].content
|
||||||
|
assert parse_checkpoint_bands(messages) == 1
|
||||||
|
|
||||||
|
band2, appended2 = inject_directive_checkpoints(
|
||||||
|
messages,
|
||||||
|
prompt_tokens=250_000,
|
||||||
|
source="api",
|
||||||
|
interval=100_000,
|
||||||
|
last_fired_band=band,
|
||||||
|
)
|
||||||
|
assert band2 == 2
|
||||||
|
assert len(appended2) == 1
|
||||||
|
assert parse_checkpoint_bands(messages) == 2
|
||||||
|
|
||||||
|
# Same band: no re-fire
|
||||||
|
band3, appended3 = inject_directive_checkpoints(
|
||||||
|
messages,
|
||||||
|
prompt_tokens=250_000,
|
||||||
|
source="api",
|
||||||
|
interval=100_000,
|
||||||
|
last_fired_band=band2,
|
||||||
|
)
|
||||||
|
assert band3 == 2
|
||||||
|
assert appended3 == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_inject_disabled() -> None:
|
||||||
|
messages: list[AnyChatMessage] = []
|
||||||
|
band, appended = inject_directive_checkpoints(
|
||||||
|
messages,
|
||||||
|
prompt_tokens=500_000,
|
||||||
|
source="api",
|
||||||
|
interval=0,
|
||||||
|
last_fired_band=0,
|
||||||
|
)
|
||||||
|
assert band == 0
|
||||||
|
assert appended == []
|
||||||
@@ -34,6 +34,27 @@ async def test_create_session_uses_default_user(store: MemoryStore) -> None:
|
|||||||
assert session.uid == user.uid
|
assert session.uid == user.uid
|
||||||
|
|
||||||
|
|
||||||
|
async def test_session_context_usage_roundtrip(store: MemoryStore) -> None:
|
||||||
|
session = await store.create_session(name="usage")
|
||||||
|
row = await store.update_session_context_usage(
|
||||||
|
session.sid,
|
||||||
|
last_prompt_tokens=120_000,
|
||||||
|
peak_prompt_tokens=150_000,
|
||||||
|
last_completion_tokens=400,
|
||||||
|
usage_source="api",
|
||||||
|
reminder_last_band=1,
|
||||||
|
)
|
||||||
|
assert row.last_prompt_tokens == 120_000
|
||||||
|
assert row.peak_prompt_tokens == 150_000
|
||||||
|
assert row.last_completion_tokens == 400
|
||||||
|
assert row.usage_source == "api"
|
||||||
|
assert row.reminder_last_band == 1
|
||||||
|
loaded = await store.get_session(session.sid)
|
||||||
|
assert loaded is not None
|
||||||
|
assert loaded.peak_prompt_tokens == 150_000
|
||||||
|
assert loaded.reminder_last_band == 1
|
||||||
|
|
||||||
|
|
||||||
async def test_append_and_list_messages(store: MemoryStore) -> None:
|
async def test_append_and_list_messages(store: MemoryStore) -> None:
|
||||||
session = await store.create_session()
|
session = await store.create_session()
|
||||||
user_msg = UserChatMessage(content="hello")
|
user_msg = UserChatMessage(content="hello")
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ from plyngent.tools import default_tool_definitions, register_builtin_tools
|
|||||||
from plyngent.tools.catalog import ToolCatalog, ToolSource, catalog_scope, get_catalog, registration_source
|
from plyngent.tools.catalog import ToolCatalog, ToolSource, catalog_scope, get_catalog, registration_source
|
||||||
from plyngent.tools.chat import CHAT_TOOLS
|
from plyngent.tools.chat import CHAT_TOOLS
|
||||||
from plyngent.tools.file import FILE_TOOLS
|
from plyngent.tools.file import FILE_TOOLS
|
||||||
|
from plyngent.tools.net import NET_TOOLS
|
||||||
from plyngent.tools.process import PROCESS_TOOLS
|
from plyngent.tools.process import PROCESS_TOOLS
|
||||||
from plyngent.tools.todo import TODO_TOOLS
|
from plyngent.tools.todo import TODO_TOOLS
|
||||||
from plyngent.tools.vcs import VCS_TOOLS
|
from plyngent.tools.vcs import VCS_TOOLS
|
||||||
@@ -17,7 +18,7 @@ from plyngent.tools.vcs import VCS_TOOLS
|
|||||||
def test_default_tool_names_match_group_lists() -> None:
|
def test_default_tool_names_match_group_lists() -> None:
|
||||||
register_builtin_tools()
|
register_builtin_tools()
|
||||||
selected = default_tool_definitions(surface="local")
|
selected = default_tool_definitions(surface="local")
|
||||||
groups = [*FILE_TOOLS, *PROCESS_TOOLS, *VCS_TOOLS, *CHAT_TOOLS, *TODO_TOOLS]
|
groups = [*FILE_TOOLS, *PROCESS_TOOLS, *VCS_TOOLS, *CHAT_TOOLS, *TODO_TOOLS, *NET_TOOLS]
|
||||||
assert sorted(t.name for t in selected) == sorted(t.name for t in groups)
|
assert sorted(t.name for t in selected) == sorted(t.name for t in groups)
|
||||||
assert len(selected) == len(groups)
|
assert len(selected) == len(groups)
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,240 @@
|
|||||||
|
"""Tests for tools.net.fetch (policy, UA, methods, private grants)."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import threading
|
||||||
|
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||||
|
from typing import TYPE_CHECKING, override
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from plyngent.tools.danger import classify_danger
|
||||||
|
from plyngent.tools.net import fetch, grant_private_host
|
||||||
|
from plyngent.tools.net.policy import (
|
||||||
|
DEFAULT_USER_AGENT,
|
||||||
|
FetchPolicyError,
|
||||||
|
HostClass,
|
||||||
|
classify_ip_strings,
|
||||||
|
normalize_method,
|
||||||
|
normalize_request_headers,
|
||||||
|
parse_fetch_url,
|
||||||
|
soft_confirm_reason,
|
||||||
|
)
|
||||||
|
from tests.test_tools.helpers import call_async
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Iterator
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
class _Handler(BaseHTTPRequestHandler):
|
||||||
|
"""Minimal echo server for fetch tests."""
|
||||||
|
|
||||||
|
@override
|
||||||
|
def log_message(self, format: str, *args: object) -> None:
|
||||||
|
del format, args
|
||||||
|
|
||||||
|
def _read_body(self) -> bytes:
|
||||||
|
length = int(self.headers.get("Content-Length") or "0")
|
||||||
|
if length <= 0:
|
||||||
|
return b""
|
||||||
|
return self.rfile.read(length)
|
||||||
|
|
||||||
|
def _send(self, code: int, body: bytes, *, content_type: str = "text/plain; charset=utf-8") -> None:
|
||||||
|
self.send_response(code)
|
||||||
|
self.send_header("Content-Type", content_type)
|
||||||
|
self.send_header("Content-Length", str(len(body)))
|
||||||
|
self.end_headers()
|
||||||
|
self.wfile.write(body)
|
||||||
|
|
||||||
|
def do_GET(self) -> None:
|
||||||
|
if self.path.startswith("/redirect-private"):
|
||||||
|
self.send_response(302)
|
||||||
|
self.send_header("Location", "http://127.0.0.1:9/nope")
|
||||||
|
self.end_headers()
|
||||||
|
return
|
||||||
|
if self.path.startswith("/redirect-loop"):
|
||||||
|
self.send_response(302)
|
||||||
|
self.send_header("Location", "/redirect-loop")
|
||||||
|
self.end_headers()
|
||||||
|
return
|
||||||
|
if self.path.startswith("/big"):
|
||||||
|
self._send(200, b"x" * 5000)
|
||||||
|
return
|
||||||
|
if self.path.startswith("/bin"):
|
||||||
|
self._send(200, b"\x00\x01\x02\xffbinary", content_type="application/octet-stream")
|
||||||
|
return
|
||||||
|
if self.path.startswith("/ua"):
|
||||||
|
ua = (self.headers.get("User-Agent") or "").encode()
|
||||||
|
self._send(200, ua)
|
||||||
|
return
|
||||||
|
self._send(200, f"GET {self.path}".encode())
|
||||||
|
|
||||||
|
def do_POST(self) -> None:
|
||||||
|
body = self._read_body()
|
||||||
|
self._send(201, b"POST:" + body)
|
||||||
|
|
||||||
|
def do_PUT(self) -> None:
|
||||||
|
body = self._read_body()
|
||||||
|
self._send(200, b"PUT:" + body)
|
||||||
|
|
||||||
|
def do_DELETE(self) -> None:
|
||||||
|
self._send(204, b"")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def http_server() -> Iterator[str]:
|
||||||
|
server = HTTPServer(("127.0.0.1", 0), _Handler)
|
||||||
|
port = server.server_address[1]
|
||||||
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||||
|
thread.start()
|
||||||
|
try:
|
||||||
|
yield f"http://127.0.0.1:{port}"
|
||||||
|
finally:
|
||||||
|
server.shutdown()
|
||||||
|
thread.join(timeout=5)
|
||||||
|
|
||||||
|
|
||||||
|
def test_normalize_method_and_url() -> None:
|
||||||
|
assert normalize_method("get") == "GET"
|
||||||
|
assert normalize_method("POST") == "POST"
|
||||||
|
with pytest.raises(FetchPolicyError):
|
||||||
|
normalize_method("PATCH")
|
||||||
|
with pytest.raises(FetchPolicyError):
|
||||||
|
parse_fetch_url("file:///etc/passwd")
|
||||||
|
with pytest.raises(FetchPolicyError):
|
||||||
|
parse_fetch_url("ftp://example.com/")
|
||||||
|
parsed = parse_fetch_url("https://Example.COM:8443/a?b=1#frag")
|
||||||
|
assert parsed.scheme == "https"
|
||||||
|
assert parsed.host.lower() == "example.com"
|
||||||
|
assert parsed.port == 8443
|
||||||
|
assert "frag" not in parsed.url
|
||||||
|
|
||||||
|
|
||||||
|
def test_user_agent_not_overridden_by_default() -> None:
|
||||||
|
hdrs = normalize_request_headers({"User-Agent": "ModelClient/1.0"})
|
||||||
|
assert hdrs["User-Agent"] == "ModelClient/1.0"
|
||||||
|
# dedicated arg wins over headers
|
||||||
|
hdrs2 = normalize_request_headers({"User-Agent": "From-Headers"}, user_agent="From-Arg")
|
||||||
|
assert hdrs2["User-Agent"] == "From-Arg"
|
||||||
|
# default only when omitted
|
||||||
|
hdrs3 = normalize_request_headers(None)
|
||||||
|
assert hdrs3["User-Agent"] == DEFAULT_USER_AGENT
|
||||||
|
with pytest.raises(FetchPolicyError):
|
||||||
|
normalize_request_headers({"Host": "evil.example"})
|
||||||
|
|
||||||
|
|
||||||
|
def test_classify_ips() -> None:
|
||||||
|
assert classify_ip_strings(["8.8.8.8"]) is HostClass.PUBLIC
|
||||||
|
assert classify_ip_strings(["127.0.0.1"]) is HostClass.PRIVATE
|
||||||
|
assert classify_ip_strings(["192.168.1.1"]) is HostClass.PRIVATE
|
||||||
|
assert classify_ip_strings(["169.254.169.254"]) is HostClass.FORBIDDEN
|
||||||
|
assert classify_ip_strings(["8.8.8.8", "10.0.0.1"]) is HostClass.PRIVATE
|
||||||
|
|
||||||
|
|
||||||
|
def test_soft_confirm_reason_matrix() -> None:
|
||||||
|
assert soft_confirm_reason(method="GET", url="https://ex.com/", scheme="https", body_present=False) is None
|
||||||
|
http_reason = soft_confirm_reason(method="GET", url="http://ex.com/", scheme="http", body_present=False)
|
||||||
|
assert http_reason is not None and "cleartext" in http_reason
|
||||||
|
post = soft_confirm_reason(method="POST", url="https://ex.com/", scheme="https", body_present=True)
|
||||||
|
assert post is not None and "POST" in post
|
||||||
|
|
||||||
|
|
||||||
|
def test_classify_danger_fetch() -> None:
|
||||||
|
assert classify_danger("fetch", {"url": "https://example.com/", "method": "GET"}) is None
|
||||||
|
reason = classify_danger("fetch", {"url": "http://example.com/", "method": "GET"})
|
||||||
|
assert reason is not None and "cleartext" in reason
|
||||||
|
reason2 = classify_danger("fetch", {"url": "https://example.com/api", "method": "DELETE"})
|
||||||
|
assert reason2 is not None and "DELETE" in reason2
|
||||||
|
|
||||||
|
|
||||||
|
async def test_fetch_get_post_put_delete(workspace: Path, http_server: str) -> None:
|
||||||
|
del workspace
|
||||||
|
base = http_server
|
||||||
|
# Grant loopback for this process/instance (fixture binds InstanceState).
|
||||||
|
grant_private_host("127.0.0.1", int(base.rsplit(":", 1)[1]))
|
||||||
|
|
||||||
|
out = await call_async(fetch, f"{base}/hello")
|
||||||
|
assert "status=200" in out
|
||||||
|
assert "GET /hello" in out
|
||||||
|
assert "security=cleartext-http" in out
|
||||||
|
|
||||||
|
out_post = await call_async(fetch, f"{base}/echo", method="POST", body="hi")
|
||||||
|
assert "status=201" in out_post
|
||||||
|
assert "POST:hi" in out_post
|
||||||
|
|
||||||
|
out_put = await call_async(fetch, f"{base}/echo", method="PUT", body="x")
|
||||||
|
assert "PUT:x" in out_put
|
||||||
|
|
||||||
|
out_del = await call_async(fetch, f"{base}/x", method="DELETE")
|
||||||
|
assert "status=204" in out_del
|
||||||
|
|
||||||
|
|
||||||
|
async def test_fetch_user_agent_passthrough(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}/ua", user_agent="AgentUA/9")
|
||||||
|
assert "AgentUA/9" in out
|
||||||
|
|
||||||
|
out2 = await call_async(
|
||||||
|
fetch,
|
||||||
|
f"{http_server}/ua",
|
||||||
|
headers={"User-Agent": "HeaderUA/1"},
|
||||||
|
)
|
||||||
|
assert "HeaderUA/1" in out2
|
||||||
|
|
||||||
|
out3 = await call_async(
|
||||||
|
fetch,
|
||||||
|
f"{http_server}/ua",
|
||||||
|
headers={"User-Agent": "HeaderUA/1"},
|
||||||
|
user_agent="ArgUA/2",
|
||||||
|
)
|
||||||
|
assert "ArgUA/2" in out3
|
||||||
|
assert "HeaderUA/1" not in out3.split("--- body ---", 1)[-1]
|
||||||
|
|
||||||
|
|
||||||
|
async def test_fetch_truncation_and_binary(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_bytes=100)
|
||||||
|
assert "truncated=true" in out
|
||||||
|
assert "bytes=100" in out
|
||||||
|
|
||||||
|
out_bin = await call_async(fetch, f"{http_server}/bin")
|
||||||
|
assert "body_kind=binary" in out_bin
|
||||||
|
assert "binary body omitted" in out_bin
|
||||||
|
|
||||||
|
|
||||||
|
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).
|
||||||
|
out = await call_async(fetch, f"{http_server}/hello")
|
||||||
|
assert out.startswith("error:")
|
||||||
|
assert "private" in out or "loopback" in out
|
||||||
|
|
||||||
|
|
||||||
|
async def test_fetch_forbidden_metadata(workspace: Path) -> None:
|
||||||
|
del workspace
|
||||||
|
out = await call_async(fetch, "http://169.254.169.254/latest/meta-data/")
|
||||||
|
assert out.startswith("error:")
|
||||||
|
assert "forbidden" in out
|
||||||
|
|
||||||
|
|
||||||
|
async def test_fetch_bad_scheme(workspace: Path) -> None:
|
||||||
|
del workspace
|
||||||
|
out = await call_async(fetch, "file:///etc/passwd")
|
||||||
|
assert out.startswith("error:")
|
||||||
|
assert "scheme" in out
|
||||||
|
|
||||||
|
|
||||||
|
async def test_fetch_redirect_loop_capped(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}/redirect-loop", max_redirects=3)
|
||||||
|
assert out.startswith("error:")
|
||||||
|
assert "redirect" in out.lower()
|
||||||
Reference in New Issue
Block a user