Report context usage as usage_update

src/acp/session.ts

@@ -27,6 +27,7 @@ import {
   isBashTool
 } from './translate/bash.js'
 import { toolResultToText } from './translate/pi-tools.js'
+import { contextTokens, contextWindowFor } from './translate/usage.js'
 
 type SessionCreateParams = {
   cwd: string
@@ -296,6 +297,11 @@ export class PiAcpSession {
   // before completing a `session/prompt` request.
   private lastEmit: Promise<void> = Promise.resolve()
 
+  // A model's context window, keyed `provider/id`. pi states it in the model
+  // list rather than on the message that spends the tokens, so it is looked up
+  // on first use and kept for the rest of the session.
+  private contextWindows = new Map<string, number | null>()
+
   constructor(opts: {
     sessionId: string
     cwd: string
@@ -463,6 +469,40 @@ export class PiAcpSession {
     })
   }
 
+  /**
+   * Resolve a model's context window, caching the answer — including a null
+   * one, so a model that never reports a window costs a single lookup rather
+   * than one per message.
+   */
+  private async contextWindow(provider: string, id: string): Promise<number | null> {
+    const key = `${provider}/${id}`
+    const cached = this.contextWindows.get(key)
+    if (cached !== undefined) return cached
+
+    let size: number | null = null
+    try {
+      size = contextWindowFor(await this.proc.getAvailableModels(), provider, id)
+    } catch {
+      // A model list we can't read just means no usage reporting for this model.
+    }
+
+    this.contextWindows.set(key, size)
+    return size
+  }
+
+  private emitUsage(message: unknown): void {
+    const m = message as any
+    const used = contextTokens(m?.usage)
+    const provider = String(m?.provider ?? '')
+    const model = String(m?.model ?? '')
+    if (used === null || !provider || !model) return
+
+    void this.contextWindow(provider, model).then(size => {
+      if (size === null) return
+      this.emit({ sessionUpdate: 'usage_update', used, size })
+    })
+  }
+
   private cleanupToolCall(toolCallId: string): void {
     this.currentToolCalls.delete(toolCallId)
     this.fileSnapshots.delete(toolCallId)
@@ -606,6 +646,15 @@ export class PiAcpSession {
           break
         }
 
+        // `done` carries the final assistant message, the one place pi states
+        // both the tokens spent and the model they were spent on. `error`
+        // carries a message too, but pi treats aborted and errored usage as
+        // invalid, so it is deliberately left unreported.
+        if (ame?.type === 'done') {
+          this.emitUsage(ame.message)
+          break
+        }
+
         // Ignore other delta/event types for now.
         break
       }

src/acp/translate/usage.ts

@@ -0,0 +1,35 @@
+/**
+ * ACP's `usage_update` reports context occupancy against the model's window.
+ * pi supplies the two halves from different places: the occupancy rides on
+ * every assistant message, while the window is a property of the model and is
+ * only stated in the model list.
+ */
+
+/**
+ * Tokens currently in context, or null when pi reported nothing usable.
+ * Mirrors pi's own accounting: prefer the provider's native total, and fall
+ * back to summing the components when it omits one.
+ */
+export function contextTokens(usage: unknown): number | null {
+  if (!usage || typeof usage !== 'object') return null
+  const u = usage as any
+
+  const total = Number(u.totalTokens)
+  if (Number.isFinite(total) && total > 0) return total
+
+  const parts = [u.input, u.output, u.cacheRead, u.cacheWrite].map(n => Number(n))
+  if (parts.some(n => !Number.isFinite(n))) return null
+
+  // All-zero usage means the message never reached the model.
+  const sum = parts.reduce((a, b) => a + b, 0)
+  return sum > 0 ? sum : null
+}
+
+/** Context window for one model, read from a `get_available_models` payload. */
+export function contextWindowFor(models: unknown, provider: string, id: string): number | null {
+  const list = Array.isArray((models as any)?.models) ? (models as any).models : []
+  const found = list.find((m: any) => String(m?.provider ?? '') === provider && String(m?.id ?? '') === id)
+
+  const size = Number(found?.contextWindow)
+  return Number.isFinite(size) && size > 0 ? size : null
+}

test/component/session-usage-update.test.ts

@@ -0,0 +1,190 @@
+import test from 'node:test'
+import assert from 'node:assert/strict'
+import { PiAcpSession } from '../../src/acp/session.js'
+import { FakeAgentSideConnection, FakePiRpcProcess, asAgentConn } from '../helpers/fakes.js'
+
+const doneEvent = (usage: unknown, provider = 'openai', model = 'gpt-x') => ({
+  type: 'message_update' as const,
+  assistantMessageEvent: {
+    type: 'done',
+    reason: 'stop',
+    message: { role: 'assistant', provider, model, usage, stopReason: 'stop' }
+  }
+})
+
+test('PiAcpSession: emits usage_update when an assistant message completes', async () => {
+  const conn = new FakeAgentSideConnection()
+  const proc = new FakePiRpcProcess()
+  proc.getAvailableModels = async () => ({
+    models: [{ provider: 'openai', id: 'gpt-x', contextWindow: 128000 }]
+  })
+
+  new PiAcpSession({
+    sessionId: 's1',
+    cwd: process.cwd(),
+    mcpServers: [],
+    proc: proc as any,
+    conn: asAgentConn(conn),
+    fileCommands: []
+  })
+
+  proc.emit(doneEvent({ input: 1000, output: 200, cacheRead: 50, cacheWrite: 0 }) as any)
+
+  await new Promise(r => setTimeout(r, 10))
+
+  assert.equal(conn.updates.length, 1)
+  assert.deepEqual(conn.updates[0]!.update, {
+    sessionUpdate: 'usage_update',
+    used: 1250,
+    size: 128000
+  })
+})
+
+test('PiAcpSession: reports usage only for a completed message', async () => {
+  const conn = new FakeAgentSideConnection()
+  const proc = new FakePiRpcProcess()
+  proc.getAvailableModels = async () => ({
+    models: [{ provider: 'openai', id: 'gpt-x', contextWindow: 128000 }]
+  })
+
+  new PiAcpSession({
+    sessionId: 's1',
+    cwd: process.cwd(),
+    mcpServers: [],
+    proc: proc as any,
+    conn: asAgentConn(conn),
+    fileCommands: []
+  })
+
+  // Both of these carry an assistant message with usage on it, and neither is
+  // a completed one: `text_end` is mid-stream and cumulative, and pi treats an
+  // errored message's usage as invalid.
+  proc.emit({
+    type: 'message_update',
+    assistantMessageEvent: {
+      type: 'text_end',
+      contentIndex: 0,
+      content: 'hi',
+      partial: { role: 'assistant', provider: 'openai', model: 'gpt-x', usage: { totalTokens: 12 } }
+    }
+  } as any)
+  proc.emit({
+    type: 'message_update',
+    assistantMessageEvent: {
+      type: 'error',
+      reason: 'error',
+      error: { role: 'assistant', provider: 'openai', model: 'gpt-x', usage: { totalTokens: 34 } }
+    }
+  } as any)
+  proc.emit(doneEvent({ totalTokens: 4096 }) as any)
+
+  await new Promise(r => setTimeout(r, 10))
+
+  const usage = conn.updates.filter(u => (u.update as any).sessionUpdate === 'usage_update')
+  assert.equal(usage.length, 1)
+  assert.deepEqual(usage[0]!.update, { sessionUpdate: 'usage_update', used: 4096, size: 128000 })
+})
+
+test('PiAcpSession: says nothing when the model reports no context window', async () => {
+  const conn = new FakeAgentSideConnection()
+  const proc = new FakePiRpcProcess()
+  proc.getAvailableModels = async () => ({ models: [{ provider: 'openai', id: 'gpt-x' }] })
+
+  new PiAcpSession({
+    sessionId: 's1',
+    cwd: process.cwd(),
+    mcpServers: [],
+    proc: proc as any,
+    conn: asAgentConn(conn),
+    fileCommands: []
+  })
+
+  proc.emit(doneEvent({ totalTokens: 4096 }) as any)
+
+  await new Promise(r => setTimeout(r, 10))
+
+  assert.deepEqual(conn.updates, [])
+})
+
+test('PiAcpSession: an unreadable model list does not break the turn', async () => {
+  const conn = new FakeAgentSideConnection()
+  const proc = new FakePiRpcProcess()
+  proc.getAvailableModels = async () => {
+    throw new Error('pi is not answering')
+  }
+
+  new PiAcpSession({
+    sessionId: 's1',
+    cwd: process.cwd(),
+    mcpServers: [],
+    proc: proc as any,
+    conn: asAgentConn(conn),
+    fileCommands: []
+  })
+
+  proc.emit(doneEvent({ totalTokens: 4096 }) as any)
+
+  await new Promise(r => setTimeout(r, 10))
+
+  assert.deepEqual(conn.updates, [])
+})
+
+test('PiAcpSession: looks up the model list once, then reuses the window', async () => {
+  const conn = new FakeAgentSideConnection()
+  const proc = new FakePiRpcProcess()
+  let lookups = 0
+  proc.getAvailableModels = async () => {
+    lookups += 1
+    return { models: [{ provider: 'openai', id: 'gpt-x', contextWindow: 128000 }] }
+  }
+
+  new PiAcpSession({
+    sessionId: 's1',
+    cwd: process.cwd(),
+    mcpServers: [],
+    proc: proc as any,
+    conn: asAgentConn(conn),
+    fileCommands: []
+  })
+
+  proc.emit(doneEvent({ totalTokens: 100 }) as any)
+  await new Promise(r => setTimeout(r, 10))
+  proc.emit(doneEvent({ totalTokens: 200 }) as any)
+  await new Promise(r => setTimeout(r, 10))
+
+  assert.equal(lookups, 1)
+  assert.deepEqual(
+    conn.updates.map(u => (u.update as any).used),
+    [100, 200]
+  )
+})
+
+test('PiAcpSession: a second model gets its own window', async () => {
+  const conn = new FakeAgentSideConnection()
+  const proc = new FakePiRpcProcess()
+  proc.getAvailableModels = async () => ({
+    models: [
+      { provider: 'openai', id: 'gpt-x', contextWindow: 128000 },
+      { provider: 'local', id: 'gpt-x', contextWindow: 8192 }
+    ]
+  })
+
+  new PiAcpSession({
+    sessionId: 's1',
+    cwd: process.cwd(),
+    mcpServers: [],
+    proc: proc as any,
+    conn: asAgentConn(conn),
+    fileCommands: []
+  })
+
+  proc.emit(doneEvent({ totalTokens: 100 }, 'openai', 'gpt-x') as any)
+  await new Promise(r => setTimeout(r, 10))
+  proc.emit(doneEvent({ totalTokens: 200 }, 'local', 'gpt-x') as any)
+  await new Promise(r => setTimeout(r, 10))
+
+  assert.deepEqual(
+    conn.updates.map(u => (u.update as any).size),
+    [128000, 8192]
+  )
+})

test/unit/usage.test.ts

@@ -0,0 +1,38 @@
+import test from 'node:test'
+import assert from 'node:assert/strict'
+import { contextTokens, contextWindowFor } from '../../src/acp/translate/usage.js'
+
+test('contextTokens: prefers the provider-reported total', () => {
+  assert.equal(contextTokens({ totalTokens: 1234, input: 1, output: 2, cacheRead: 3, cacheWrite: 4 }), 1234)
+})
+
+test('contextTokens: sums the components when no total is reported', () => {
+  assert.equal(contextTokens({ totalTokens: 0, input: 10, output: 20, cacheRead: 30, cacheWrite: 40 }), 100)
+})
+
+test('contextTokens: reports nothing for absent, malformed or all-zero usage', () => {
+  assert.equal(contextTokens(undefined), null)
+  assert.equal(contextTokens('nope'), null)
+  assert.equal(contextTokens({ totalTokens: 0, input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }), null)
+  // A partial usage object is not silently treated as a smaller one.
+  assert.equal(contextTokens({ input: 1, output: 2 }), null)
+})
+
+test('contextWindowFor: matches on provider and id together', () => {
+  const models = {
+    models: [
+      { provider: 'alpha', id: 'shared', contextWindow: 100 },
+      { provider: 'beta', id: 'shared', contextWindow: 200 }
+    ]
+  }
+
+  assert.equal(contextWindowFor(models, 'beta', 'shared'), 200)
+  assert.equal(contextWindowFor(models, 'gamma', 'shared'), null)
+})
+
+test('contextWindowFor: tolerates a missing, zero or unreadable window', () => {
+  assert.equal(contextWindowFor({ models: [{ provider: 'a', id: 'm' }] }, 'a', 'm'), null)
+  assert.equal(contextWindowFor({ models: [{ provider: 'a', id: 'm', contextWindow: 0 }] }, 'a', 'm'), null)
+  assert.equal(contextWindowFor({}, 'a', 'm'), null)
+  assert.equal(contextWindowFor(undefined, 'a', 'm'), null)
+})