From e3e0cf2080a26e28d3a120157a2dd4c250c1cb26 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Fri, 24 Jul 2026 19:40:37 +0800 Subject: [PATCH] core/agent: append-only directive checkpoints on token bands --- src/plyngent/agent/chat.py | 104 ++++++++++++++++- src/plyngent/agent/directive_checkpoint.py | 109 ++++++++++++++++++ src/plyngent/agent/loop.py | 82 +++++++++++-- tests/test_agent/test_directive_checkpoint.py | 82 +++++++++++++ 4 files changed, 364 insertions(+), 13 deletions(-) create mode 100644 src/plyngent/agent/directive_checkpoint.py create mode 100644 tests/test_agent/test_directive_checkpoint.py diff --git a/src/plyngent/agent/chat.py b/src/plyngent/agent/chat.py index af7624c..25f115a 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -17,6 +17,10 @@ from .budget import ( DEFAULT_TOOL_RESULT_MAX_CHARS, estimate_messages_tokens, ) +from .directive_checkpoint import ( + DEFAULT_DIRECTIVE_REMINDER_TOKENS, + parse_checkpoint_bands, +) from .events import UsageEvent from .loop import DEFAULT_MAX_ROUNDS, run_chat_loop from .todo_nag import ( @@ -128,11 +132,15 @@ class ChatAgent: max_context_tokens: int todo_stack: TodoStack | None todo_nag_strategy: TodoNagStrategy + directive_reminder_tokens: int + directive_reminder_text: str | None messages: list[AnyChatMessage] session_usage: TokenUsage last_turn_usage: TokenUsage last_request_usage: TokenUsage last_turn_rounds: int + peak_prompt_tokens: int + reminder_last_band: int # Index into messages of the first unpersisted message (checkpoint cursor). _persist_from: int @@ -155,6 +163,10 @@ class ChatAgent: max_context_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS, todo_stack: TodoStack | None = None, 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: self.client = client self.model = model @@ -171,13 +183,19 @@ class ChatAgent: self.max_context_tokens = max_context_tokens self.todo_stack = todo_stack 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.session_usage = TokenUsage() self.last_turn_usage = TokenUsage() self.last_request_usage = TokenUsage() 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._ensure_system_prompt() + self._sync_reminder_band_from_messages() @property def pending_retry_text(self) -> str | None: @@ -214,6 +232,53 @@ class ChatAgent: # already pointed past stored messages stay correct after insert. 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( self, messages: Sequence[AnyChatMessage], @@ -233,6 +298,7 @@ class ChatAgent: else: self._persist_from = len(self.messages) if persisted else 0 self._ensure_system_prompt() + self._sync_reminder_band_from_messages() @property def persist_from(self) -> int: @@ -246,6 +312,15 @@ class ChatAgent: raise RuntimeError(msg) loaded = await self.memory.list_messages(self.session_id) 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: """Attach a memory session id; optionally load existing messages.""" @@ -295,6 +370,15 @@ class ChatAgent: end = committed_prefix_end(self.messages, user_index) 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]: """Run the tool loop for an already-appended user message. @@ -340,17 +424,29 @@ class ChatAgent: max_context_tokens=self.max_context_tokens, todo_stack=self.todo_stack, 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): turn_rounds += 1 last_request = event.usage turn_usage = turn_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 - # 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) - if commit_end > self._persist_from: - await self._persist_range(self._persist_from, commit_end) + end = self._developer_tail_end(commit_end) + if end > self._persist_from: + await self._persist_range(self._persist_from, end) yield event completed = True except BaseException: @@ -363,6 +459,7 @@ class ChatAgent: self.last_turn_rounds = turn_rounds if 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]: """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, todo_stack=None, todo_nag_strategy="none", + directive_reminder_tokens=0, messages=history, ) async for event in aside.run(text): diff --git a/src/plyngent/agent/directive_checkpoint.py b/src/plyngent/agent/directive_checkpoint.py new file mode 100644 index 0000000..26fb0db --- /dev/null +++ b/src/plyngent/agent/directive_checkpoint.py @@ -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 diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index 26e5547..c60464f 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -26,6 +26,10 @@ from .budget import ( estimate_messages_tokens, truncate_tool_result, ) +from .directive_checkpoint import ( + DEFAULT_DIRECTIVE_REMINDER_TOKENS, + inject_directive_checkpoints, +) from .events import ( AgentEvent, AssistantMessageEvent, @@ -312,7 +316,36 @@ def _last_assistant(messages: list[AnyChatMessage], pre_len: int) -> AssistantCh 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, 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, todo_stack: TodoStack | None = None, 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]: """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 (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. + + 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 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 sent_estimate_tokens: int | None = None todo_review_injected = False + last_band = max(0, reminder_last_band) while True: while rounds_used < allowance: @@ -372,15 +415,42 @@ async def run_chat_loop( # noqa: C901 — multi-phase tool loop ) pre_len = len(messages) + last_usage_event: UsageEvent | None = None async for event in _assistant_round(client, param, messages, stream=stream): if isinstance(event, UsageEvent): # Next rounds scale char-estimates by real/resolved prompt size. prompt_tokens_hint = event.usage.prompt_tokens sent_estimate_tokens = sent_est + last_usage_event = event yield event + assistant = _last_assistant(messages, pre_len) 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: todo_review_injected = True 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: continue 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})" if on_limit is not None and await _call_on_limit(on_limit, reason): diff --git a/tests/test_agent/test_directive_checkpoint.py b/tests/test_agent/test_directive_checkpoint.py new file mode 100644 index 0000000..52dc759 --- /dev/null +++ b/tests/test_agent/test_directive_checkpoint.py @@ -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 == []