fix: preserve tool output and session history

Author: elkaix

.changeset/clear-ravens-wave.md

@@ -0,0 +1,5 @@
+---
+"@pymodel/pythinker-code": patch
+---
+
+Fix context usage updates in interactive clients.

.changeset/kind-turtles-grow.md

@@ -0,0 +1,5 @@
+---
+"@pymodel/pythinker-code": patch
+---
+
+Fix session history after steering or undoing a turn.

.changeset/tidy-spiders-fix.md

@@ -0,0 +1,5 @@
+---
+"@pymodel/pythinker-code": patch
+---
+
+Fix loss of large tool outputs in long conversations.

apps/pythinker-code/src/tui/controllers/session-event-handler.ts

@@ -721,9 +721,15 @@ export class SessionEventHandler {
       this.host.state.appState.dynamicWorkflowMode &&
       this.host.state.dynamicWorkflowModeEntry === 'task';
     const patch: Partial<AppState> = {};
-    if (event.contextUsage !== undefined) patch.contextUsage = event.contextUsage;
     if (event.contextTokens !== undefined) patch.contextTokens = event.contextTokens;
     if (event.maxContextTokens !== undefined) patch.maxContextTokens = event.maxContextTokens;
+    if (event.contextUsage !== undefined) {
+      patch.contextUsage = event.contextUsage;
+    } else if (event.contextTokens !== undefined || event.maxContextTokens !== undefined) {
+      const tokens = patch.contextTokens ?? this.host.state.appState.contextTokens;
+      const max = patch.maxContextTokens ?? this.host.state.appState.maxContextTokens;
+      patch.contextUsage = max > 0 ? tokens / max : 0;
+    }
     if (event.planMode !== undefined) patch.planMode = event.planMode;
     if (event.dynamicWorkflowMode !== undefined) patch.dynamicWorkflowMode = event.dynamicWorkflowMode;
     if (event.towerMode !== undefined) patch.towerMode = event.towerMode;

apps/pythinker-code/test/tui/pythinker-tui-message-flow.test.ts

@@ -5041,6 +5041,86 @@ command = "vim"
     expect(stripSgr(renderTranscript(driver))).toContain('LLM not set');
   });
 
+  it('recomputes context usage when a status update carries context tokens without it', async () => {
+    const { driver } = await makeDriver();
+    driver.state.appState.contextTokens = 0;
+    driver.state.appState.maxContextTokens = 1_000_000;
+    driver.state.appState.contextUsage = 0.74;
+
+    driver.sessionEventHandler.handleEvent(
+      {
+        type: 'agent.status.updated',
+        agentId: 'main',
+        sessionId: 'ses-1',
+        contextTokens: 180_000,
+      } as Event,
+      vi.fn(),
+    );
+
+    expect(driver.state.appState.contextTokens).toBe(180_000);
+    expect(driver.state.appState.contextUsage).toBeCloseTo(0.18);
+  });
+
+  it('recomputes context usage when a status update carries max context tokens without it', async () => {
+    const { driver } = await makeDriver();
+    driver.state.appState.contextTokens = 180_000;
+    driver.state.appState.maxContextTokens = 256_000;
+    driver.state.appState.contextUsage = 180_000 / 256_000;
+
+    driver.sessionEventHandler.handleEvent(
+      {
+        type: 'agent.status.updated',
+        agentId: 'main',
+        sessionId: 'ses-1',
+        maxContextTokens: 1_000_000,
+      } as Event,
+      vi.fn(),
+    );
+
+    expect(driver.state.appState.maxContextTokens).toBe(1_000_000);
+    expect(driver.state.appState.contextUsage).toBeCloseTo(0.18);
+  });
+
+  it('keeps an explicit context usage from status updates', async () => {
+    const { driver } = await makeDriver();
+    driver.state.appState.contextTokens = 100;
+    driver.state.appState.maxContextTokens = 1_000_000;
+    driver.state.appState.contextUsage = 0;
+
+    driver.sessionEventHandler.handleEvent(
+      {
+        type: 'agent.status.updated',
+        agentId: 'main',
+        sessionId: 'ses-1',
+        contextTokens: 180_000,
+        maxContextTokens: 1_000_000,
+        contextUsage: 0.42,
+      } as Event,
+      vi.fn(),
+    );
+
+    expect(driver.state.appState.contextUsage).toBe(0.42);
+  });
+
+  it('zeroes context usage when no context window is known', async () => {
+    const { driver } = await makeDriver();
+    driver.state.appState.contextTokens = 180_000;
+    driver.state.appState.maxContextTokens = 0;
+    driver.state.appState.contextUsage = 0.74;
+
+    driver.sessionEventHandler.handleEvent(
+      {
+        type: 'agent.status.updated',
+        agentId: 'main',
+        sessionId: 'ses-1',
+        contextTokens: 190_000,
+      } as Event,
+      vi.fn(),
+    );
+
+    expect(driver.state.appState.contextUsage).toBe(0);
+  });
+
   it('applies the effective thinking effort from status updates', async () => {
     const { driver } = await makeDriver();
 

apps/vscode/src/runtime/event-adapter.ts

@@ -367,7 +367,8 @@ function mapStatusUpdate(
   sdkEvent: Extract<Event, { type: 'agent.status.updated' }>,
 ): MappedLegacyWireEvent {
   const payload: StatusUpdate = {};
-  if (sdkEvent.contextUsage !== undefined) payload.context_usage = sdkEvent.contextUsage;
+  const contextUsage = contextUsageRatio(sdkEvent);
+  if (contextUsage !== undefined) payload.context_usage = contextUsage;
   if (sdkEvent.planMode !== undefined) payload.plan_mode = sdkEvent.planMode;
   const thinkingLevel = (sdkEvent as any).thinkingLevel ?? (sdkEvent as any).thinkingEffort;
   if (thinkingLevel !== undefined) payload.thinking_effort = thinkingLevel;
@@ -419,6 +420,22 @@ function mapSubagentStatus(
   return { state, event: { type: 'SubagentStatus', payload } };
 }
 
+function contextUsageRatio(
+  sdkEvent: Extract<Event, { type: 'agent.status.updated' }>,
+): number | undefined {
+  if (sdkEvent.contextUsage !== undefined) return sdkEvent.contextUsage;
+  const { contextTokens, maxContextTokens } = sdkEvent;
+  if (
+    typeof contextTokens !== 'number' ||
+    typeof maxContextTokens !== 'number' ||
+    !Number.isFinite(contextTokens) ||
+    !Number.isFinite(maxContextTokens)
+  ) {
+    return undefined;
+  }
+  return maxContextTokens > 0 ? contextTokens / maxContextTokens : undefined;
+}
+
 function usageDelta(current: AdapterTokenUsage, previous: AdapterTokenUsage | undefined): TokenUsage {
   return {
     input_other: delta(current.inputOther, previous?.inputOther),

apps/vscode/src/runtime/session-runtime.ts

@@ -172,6 +172,7 @@ export class SessionRuntime {
           model: status.model,
           thinking_effort: status.thinkingEffort,
           plan_mode: status.planMode,
+          context_usage: status.contextUsage,
           permission: status.permission,
         },
         _sessionId: this.id,

apps/vscode/test/event-adapter.test.ts

@@ -295,6 +295,50 @@ describe('event adapter (projects SDK events into the legacy Webview contract)',
     });
   });
 
+  it('derives context usage from the v2 context token pair', () => {
+    const result = adaptSdkEvent(createEventAdapterState(), {
+      type: 'agent.status.updated',
+      sessionId: 'session-1',
+      agentId: 'main',
+      contextTokens: 25_600,
+      maxContextTokens: 256_000,
+    });
+
+    expect(result.event).toEqual({
+      type: 'StatusUpdate',
+      payload: { context_usage: 0.1 },
+      _sessionId: 'session-1',
+    });
+  });
+
+  it('preserves an explicit context usage over the context token pair', () => {
+    const result = adaptSdkEvent(createEventAdapterState(), {
+      type: 'agent.status.updated',
+      sessionId: 'session-1',
+      agentId: 'main',
+      contextUsage: 0,
+      contextTokens: 25_600,
+      maxContextTokens: 256_000,
+    });
+
+    expect(result.event).toMatchObject({
+      type: 'StatusUpdate',
+      payload: { context_usage: 0 },
+    });
+  });
+
+  it('does not derive a ratio from an invalid context capacity', () => {
+    const result = adaptSdkEvent(createEventAdapterState(), {
+      type: 'agent.status.updated',
+      sessionId: 'session-1',
+      agentId: 'main',
+      contextTokens: 25_600,
+      maxContextTokens: 0,
+    });
+
+    expect(result.event).toBeUndefined();
+  });
+
   it('emits only new token usage when SDK status carries cumulative turn usage', () => {
     const first = adaptSdkEvent(createEventAdapterState(), {
       type: 'agent.status.updated',

apps/vscode/test/pythinker-runtime.test.ts

@@ -428,7 +428,13 @@ describe("Pythinker runtime (owns shared SDK sessions for Webviews)", () => {
         type: "StatusUpdate",
         // The permission mode rides along: the chat badge is the only place the
         // user can see which mode a toggle command just landed on.
-        payload: { model: "kimi-test", thinking_effort: "max", plan_mode: true, permission: "manual" },
+        payload: {
+          model: "kimi-test",
+          thinking_effort: "max",
+          plan_mode: true,
+          permission: "manual",
+          context_usage: 0,
+        },
         _sessionId: "saved-1",
       },
       webviewId: "view-1",
@@ -459,7 +465,13 @@ describe("Pythinker runtime (owns shared SDK sessions for Webviews)", () => {
       event: Events.StreamEvent,
       data: {
         type: "StatusUpdate",
-        payload: { model: "kimi-test", thinking_effort: "off", plan_mode: false, permission: "yolo" },
+        payload: {
+          model: "kimi-test",
+          thinking_effort: "off",
+          plan_mode: false,
+          permission: "yolo",
+          context_usage: 0,
+        },
         _sessionId: "saved-1",
       },
       webviewId: "view-1",

apps/vscode/test/session-runtime.test.ts

@@ -218,6 +218,24 @@ function turnEnded(
 }
 
 describe("session runtime (adapts one SDK session for subscribed Webviews)", () => {
+  it("announces the current context usage to a subscribed Webview", async () => {
+    const { runtime, broadcasts } = createRuntime();
+
+    await runtime.announceStatus("view-1");
+
+    expect(streamData(broadcasts)).toContainEqual({
+      type: "StatusUpdate",
+      payload: {
+        model: undefined,
+        thinking_effort: "off",
+        plan_mode: false,
+        permission: "manual",
+        context_usage: 0,
+      },
+      _sessionId: "session-1",
+    });
+  });
+
   it("renders a host-only command without making it a forkable core turn", () => {
     const { runtime, broadcasts } = createRuntime();