core/agent: append-only directive checkpoints on token bands

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:
+                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
             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):

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

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):

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 == []