fix: carry observed overflow token counts into compaction (#40357)

Merged via squash.

Prepared head SHA: b99eed4329
Co-authored-by: rabsef-bicrym <52549148+rabsef-bicrym@users.noreply.github.com>
Co-authored-by: jalehman <550978+jalehman@users.noreply.github.com>
Reviewed-by: @jalehman

src/context-engine/context-engine.test.ts

@@ -1,5 +1,6 @@
 import type { AgentMessage } from "@mariozechner/pi-agent-core";
-import { describe, expect, it, beforeEach } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { compactEmbeddedPiSessionDirect } from "../agents/pi-embedded-runner/compact.runtime.js";
 // ---------------------------------------------------------------------------
 // We dynamically import the registry so we can get a fresh module per test
 // group when needed.  For most groups we use the shared singleton directly.
@@ -19,6 +20,23 @@ import type {
   IngestResult,
 } from "./types.js";
 
+vi.mock("../agents/pi-embedded-runner/compact.runtime.js", () => ({
+  compactEmbeddedPiSessionDirect: vi.fn(async () => ({
+    ok: true,
+    compacted: false,
+    reason: "mock compaction",
+    result: {
+      summary: "",
+      firstKeptEntryId: "",
+      tokensBefore: 0,
+      tokensAfter: 0,
+      details: undefined,
+    },
+  })),
+}));
+
+const mockedCompactEmbeddedPiSessionDirect = vi.mocked(compactEmbeddedPiSessionDirect);
+
 // ---------------------------------------------------------------------------
 // Helpers
 // ---------------------------------------------------------------------------
@@ -91,6 +109,10 @@ class MockContextEngine implements ContextEngine {
 // ═══════════════════════════════════════════════════════════════════════════
 
 describe("Engine contract tests", () => {
+  beforeEach(() => {
+    mockedCompactEmbeddedPiSessionDirect.mockClear();
+  });
+
   it("a mock engine implementing ContextEngine can be registered and resolved", async () => {
     const factory = () => new MockContextEngine();
     registerContextEngine("mock", factory);
@@ -153,6 +175,25 @@ describe("Engine contract tests", () => {
     // Should complete without error
     await expect(engine.dispose()).resolves.toBeUndefined();
   });
+
+  it("legacy compact preserves runtimeContext currentTokenCount when top-level value is absent", async () => {
+    const engine = new LegacyContextEngine();
+
+    await engine.compact({
+      sessionId: "s1",
+      sessionFile: "/tmp/session.json",
+      runtimeContext: {
+        workspaceDir: "/tmp/workspace",
+        currentTokenCount: 277403,
+      },
+    });
+
+    expect(mockedCompactEmbeddedPiSessionDirect).toHaveBeenCalledWith(
+      expect.objectContaining({
+        currentTokenCount: 277403,
+      }),
+    );
+  });
 });
 
 // ═══════════════════════════════════════════════════════════════════════════

src/context-engine/legacy.ts

@@ -78,6 +78,13 @@ export class LegacyContextEngine implements ContextEngine {
     // set by the caller in run.ts. We spread them and override the fields
     // that come from the ContextEngine compact() signature directly.
     const runtimeContext = params.runtimeContext ?? {};
+    const currentTokenCount =
+      params.currentTokenCount ??
+      (typeof runtimeContext.currentTokenCount === "number" &&
+      Number.isFinite(runtimeContext.currentTokenCount) &&
+      runtimeContext.currentTokenCount > 0
+        ? Math.floor(runtimeContext.currentTokenCount)
+        : undefined);
 
     // eslint-disable-next-line @typescript-eslint/no-explicit-any -- bridge runtimeContext matches CompactEmbeddedPiSessionParams
     const result = await compactEmbeddedPiSessionDirect({
@@ -85,6 +92,7 @@ export class LegacyContextEngine implements ContextEngine {
       sessionId: params.sessionId,
       sessionFile: params.sessionFile,
       tokenBudget: params.tokenBudget,
+      ...(currentTokenCount !== undefined ? { currentTokenCount } : {}),
       force: params.force,
       customInstructions: params.customInstructions,
       workspaceDir: (runtimeContext.workspaceDir as string) ?? process.cwd(),