From f0cbfebc73e31f390671b4df70a586b91454f532 Mon Sep 17 00:00:00 2001 From: Sunny Date: Tue, 13 Oct 2026 08:45:28 +0530 Subject: [PATCH] fix(engine): preserve prompt cache prefix with skill hook context --- .../_engine/adapters.py | 42 ++++++++++++- .../engine/tests/test_compaction_notice.py | 61 +++++++++++++++++++ 2 files changed, 100 insertions(+), 3 deletions(-) diff --git a/packages/engine/src/amplifier_agent_engine/_engine/adapters.py b/packages/engine/src/amplifier_agent_engine/_engine/adapters.py index 67dc2fc2..390d11d4 100644 --- a/packages/engine/src/amplifier_agent_engine/_engine/adapters.py +++ b/packages/engine/src/amplifier_agent_engine/_engine/adapters.py @@ -157,6 +157,11 @@ def __init__(self, tool_result_max_bytes: int | None, *, hooks: Any = None) -> N ) # The compaction stats object the notice describes, its index in the view, and the message. self._notice_anchor: tuple[dict[str, Any], int, dict[str, Any]] | None = None + # Request-only skill context is anchored in batches so later conversation + # growth extends, rather than changes, the provider-visible prompt prefix. + self._skill_context_anchors: list[tuple[int, dict[str, Any]]] = [] + self._skill_context_count = 0 + self._skill_context_stats: dict[str, Any] | None = None # The input budget the engine sets for the selection in use, or None for the provider's window. self.model_budget: Callable[[], int | None] = lambda: None self.entry_token: str | None = None @@ -204,9 +209,36 @@ async def get_messages_for_request( ) -> list[dict[str, Any]]: messages = await super().get_messages_for_request(token_budget=token_budget, provider=provider) messages = self._with_compaction_notice(messages) - if self.skill_context: - messages.append({"role": "user", "content": "\n\n".join(self.skill_context)}) - return messages + return self._with_skill_context(messages) + + def _with_skill_context(self, view: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Insert turn-scoped hook context without changing an earlier request prefix.""" + stats = self._last_compaction_stats + if self._skill_context_count > len(self.skill_context) or ( + self._skill_context_anchors and self._skill_context_stats is not stats + ): + self._reset_skill_context_anchors() + + result = list(view) + for index, message in self._skill_context_anchors: + if index > len(result): + self._reset_skill_context_anchors() + return self._with_skill_context(view) + result.insert(index, copy.deepcopy(message)) + + additions = self.skill_context[self._skill_context_count :] + if additions and not self._tail_has_unanswered_tool_calls(result): + message = {"role": "user", "content": "\n\n".join(additions)} + self._skill_context_anchors.append((len(result), message)) + self._skill_context_count = len(self.skill_context) + self._skill_context_stats = stats + result.append(copy.deepcopy(message)) + return result + + def _reset_skill_context_anchors(self) -> None: + self._skill_context_anchors = [] + self._skill_context_count = 0 + self._skill_context_stats = None def _with_compaction_notice(self, view: list[dict[str, Any]]) -> list[dict[str, Any]]: """Insert the notice for the latest compaction at the position it was first placed. @@ -236,10 +268,12 @@ def _with_compaction_notice(self, view: list[dict[str, Any]]) -> list[dict[str, async def set_messages(self, messages: list[dict[str, Any]]) -> None: self._notice_anchor = None + self._reset_skill_context_anchors() await super().set_messages(messages) async def clear(self) -> None: self._notice_anchor = None + self._reset_skill_context_anchors() await super().clear() def _derive_budget(self, token_budget: int | None, provider: Any | None) -> int: @@ -863,6 +897,7 @@ async def execute(self, input: TurnInput, observer: Observer) -> None: self.observer = observer self.skill_hooks.begin(self.inherited_skill_hooks) self.context.skill_context.clear() + self.context._reset_skill_context_anchors() self._call_ids.clear() self.core.coordinator.cancellation.reset() self._turn_start = len(await self.context.get_messages()) @@ -888,6 +923,7 @@ async def execute(self, input: TurnInput, observer: Observer) -> None: finally: self.skill_hooks.clear() self.context.skill_context.clear() + self.context._reset_skill_context_anchors() self.context.turn_active = False self.observer = None diff --git a/packages/engine/tests/test_compaction_notice.py b/packages/engine/tests/test_compaction_notice.py index ffa8e874..9d76ed1b 100644 --- a/packages/engine/tests/test_compaction_notice.py +++ b/packages/engine/tests/test_compaction_notice.py @@ -155,3 +155,64 @@ def thinking(text: str) -> dict[str, Any]: with_notice = [*plain[:4], notice, *plain[4:]] replayed = _bounded_reasoning_replay(with_notice) assert [message for message in replayed if not is_notice(message)] == _bounded_reasoning_replay(plain) + + +async def test_skill_context_preserves_request_prefix_within_turn(): + context = await filled_context() + context.skill_context = ["PreToolUse context"] + + first = await context.get_messages_for_request(token_budget=BUDGET) + stats = context._last_compaction_stats + assert stats is not None + + await context.add_message({"role": "assistant", "content": "Continuing."}) + await context.add_message({"role": "user", "content": "Next action."}) + + second = await context.get_messages_for_request(token_budget=BUDGET) + + assert context._last_compaction_stats is stats + assert wire(second[: len(first)]) == wire(first) + stored = await context.get_messages() + assert not any("PreToolUse context" in str(message) for message in stored) + assert not any(is_notice(message) for message in stored) + + +async def test_additional_skill_hook_context_keeps_previous_request_stable(): + context = await filled_context() + context.skill_context = ["First hook context"] + + first = await context.get_messages_for_request(token_budget=BUDGET) + stats = context._last_compaction_stats + assert stats is not None + + await context.add_message({"role": "assistant", "content": "Tool completed."}) + await context.add_message({"role": "user", "content": "Continue."}) + + context.skill_context.append("Second hook context") + second = await context.get_messages_for_request(token_budget=BUDGET) + + assert context._last_compaction_stats is stats + assert wire(second[: len(first)]) == wire(first) + assert any("Second hook context" in str(message) for message in second) + assert not any("First hook context" in str(message) for message in await context.get_messages()) + assert not any("Second hook context" in str(message) for message in await context.get_messages()) + + +async def test_skill_context_waits_for_an_unanswered_tool_call(): + context = await filled_context() + await tool_round(context, 99, answered=False) + context.skill_context = ["PreToolUse context"] + + view = await context.get_messages_for_request(token_budget=BUDGET) + assert view[-1]["tool_calls"][0]["id"] == "call-99" + assert not any("PreToolUse context" in str(message) for message in view) + + await context.add_message({"role": "tool", "name": "read_file", "tool_call_id": "call-99", "content": "done"}) + view = await context.get_messages_for_request(token_budget=BUDGET) + call_index = next( + index + for index, message in enumerate(view) + if any(call["id"] == "call-99" for call in message.get("tool_calls", [])) + ) + assert view[call_index + 1]["tool_call_id"] == "call-99" + assert view[-1] == {"role": "user", "content": "PreToolUse context"}