From ff9d68bdd90c4ee90bec5b82c5f8390e8e33ffc2 Mon Sep 17 00:00:00 2001
From: Bikkies <cheesenbizkitz@gmail.com>
Date: Thu, 9 Apr 2026 19:00:11 +1000
Subject: [PATCH] fix: re-assemble when afterTurn is unavailable

- Decouple lastSeenLength advancement from afterTurn existence so engines without afterTurn still advance the fence and trigger assemble on new messages
- Use hasNewMessages flag instead of didCallAfterTurn for the assemble gate
- Add test verifying assemble fires on every iteration with new messages when engine lacks afterTurn
---
 .../tool-result-context-guard.test.ts         | 26 +++++++++++++++++++
 .../tool-result-context-guard.ts              | 11 ++++----
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/src/agents/pi-embedded-runner/tool-result-context-guard.test.ts b/src/agents/pi-embedded-runner/tool-result-context-guard.test.ts
index fbe8b5ef46f..598430b1e9f 100644
--- a/src/agents/pi-embedded-runner/tool-result-context-guard.test.ts
+++ b/src/agents/pi-embedded-runner/tool-result-context-guard.test.ts
@@ -463,6 +463,32 @@ describe("installContextEngineLoopHook", () => {
     expect(engine.assemble).toHaveBeenCalledTimes(1);
   });
 
+  it("keeps calling assemble on subsequent iterations when engine lacks afterTurn but new messages arrive", async () => {
+    const agent = makeGuardableAgent();
+    const engine = makeMockEngine({ omitAfterTurn: true });
+    installContextEngineLoopHook({
+      agent,
+      contextEngine: engine,
+      sessionId,
+      sessionKey,
+      sessionFile,
+      tokenBudget,
+      modelId,
+    });
+
+    const batch1 = [makeUser("first"), makeToolResult("call_1", "r1")];
+    await callTransform(agent, batch1);
+    expect(engine.assemble).toHaveBeenCalledTimes(1);
+
+    const batch2 = [...batch1, makeUser("second"), makeToolResult("call_2", "r2")];
+    await callTransform(agent, batch2);
+    expect(engine.assemble).toHaveBeenCalledTimes(2);
+
+    const batch3 = [...batch2, makeUser("third"), makeToolResult("call_3", "r3")];
+    await callTransform(agent, batch3);
+    expect(engine.assemble).toHaveBeenCalledTimes(3);
+  });
+
   it("falls through to the source messages when engine.afterTurn throws", async () => {
     const agent = makeGuardableAgent();
     const engine = makeMockEngine({
diff --git a/src/agents/pi-embedded-runner/tool-result-context-guard.ts b/src/agents/pi-embedded-runner/tool-result-context-guard.ts
index 0a3bdcd6db9..620a9629239 100644
--- a/src/agents/pi-embedded-runner/tool-result-context-guard.ts
+++ b/src/agents/pi-embedded-runner/tool-result-context-guard.ts
@@ -222,9 +222,9 @@ export function installContextEngineLoopHook(params: {
       lastSeenLength = params.getPrePromptMessageCount?.() ?? 0;
     }
 
-    let didCallAfterTurn = false;
+    const hasNewMessages = sourceMessages.length > lastSeenLength;
     try {
-      if (sourceMessages.length > lastSeenLength && typeof contextEngine.afterTurn === "function") {
+      if (hasNewMessages && typeof contextEngine.afterTurn === "function") {
         const prePromptCount = lastSeenLength;
         await contextEngine.afterTurn({
           sessionId,
@@ -234,13 +234,14 @@ export function installContextEngineLoopHook(params: {
           prePromptMessageCount: prePromptCount,
           tokenBudget,
         });
+      }
+      if (hasNewMessages) {
         lastSeenLength = sourceMessages.length;
-        didCallAfterTurn = true;
       }
       // Skip assemble when nothing has changed since the last call AND we
       // already returned an assembled view at least once. The engine's view
-      // cannot have changed without new messages arriving via afterTurn.
-      if (didCallAfterTurn || !hasAssembledBefore) {
+      // cannot have changed without new messages arriving.
+      if (hasNewMessages || !hasAssembledBefore) {
         const assembled = await contextEngine.assemble({
           sessionId,
           sessionKey,