feat: implement bounded LLM request recovery

packages/ui/tui/src/index.ts

@@ -35,7 +35,8 @@ import type { Context } from 'cordis'
 import z from 'schemastery'
 import type { Agent, AgentStatus } from '@deepseek-ai/dsh-agent'
 import type {} from '@deepseek-ai/dsh-agent-loop'
-import type { ContentBlock, StreamChunk } from '@deepseek-ai/dsh-llm'
+import type { ContentBlock, StreamChunk, TokenUsage } from '@deepseek-ai/dsh-llm'
+import type {} from '@deepseek-ai/dsh-llm-retry'
 import { SessionId, type Session, type SessionEvent, type TodoItem } from '@deepseek-ai/dsh-session'
 import type {
   FileDiff,
@@ -612,15 +613,38 @@ function formatCwd(cwd: string | undefined): string {
   return displayText(cwd)
 }
 
-function sessionTokens(session: Session): { input: number; output: number } {
-  let input = 0
-  let output = 0
-  for (const event of session.events) {
-    if (event.type !== 'assistant/message' || event.data.usage === undefined) continue
-    input += event.data.usage.inputTokens
-    output += event.data.usage.outputTokens
+interface SessionTokenTotals {
+  input: number
+  output: number
+  readonly byStep: Map<string, TokenUsage>
+}
+
+function recordTokenUsage(totals: SessionTokenTotals, turn: number, step: number, usage: TokenUsage): void {
+  const key = `${turn}:${step}`
+  const previous = totals.byStep.get(key)
+  if (previous !== undefined) {
+    totals.input -= previous.inputTokens
+    totals.output -= previous.outputTokens
   }
-  return { input, output }
+  totals.byStep.set(key, usage)
+  totals.input += usage.inputTokens
+  totals.output += usage.outputTokens
+}
+
+function recordEventUsage(totals: SessionTokenTotals, event: SessionEvent): void {
+  if (event.type === 'assistant/chunk' && event.data.chunk.type === 'usage') {
+    recordTokenUsage(totals, event.data.turn, event.data.step, event.data.chunk.usage)
+  } else if (event.type === 'assistant/message' && event.data.usage !== undefined) {
+    recordTokenUsage(totals, event.data.turn, event.data.step, event.data.usage)
+  }
+}
+
+function sessionTokens(session: Session): SessionTokenTotals {
+  const totals: SessionTokenTotals = { input: 0, output: 0, byStep: new Map() }
+  for (const event of session.events) {
+    recordEventUsage(totals, event)
+  }
+  return totals
 }
 
 class FooterComponent implements Component {
@@ -899,6 +923,14 @@ export function createTuiChat(
     return card
   }
 
+  const clearStreaming = (): void => {
+    if (streaming === undefined) return
+    const index = chat.children.indexOf(streaming)
+    /* v8 ignore next -- streaming is assigned only after the same component is added, and every removal clears it. */
+    if (index >= 0) chat.children.splice(index, 1)
+    streaming = undefined
+  }
+
   const renderEvent = (event: SessionEvent, options: { addHistory: boolean; renderChunks: boolean }): void => {
     switch (event.type) {
       case 'user/message': {
@@ -941,15 +973,19 @@ export function createTuiChat(
         }
         break
       case 'assistant/message': {
-        if (streaming !== undefined) {
-          const index = chat.children.indexOf(streaming)
-          if (index >= 0) chat.children.splice(index, 1)
-          streaming = undefined
-        }
+        clearStreaming()
         const component = new AssistantMessageComponent(event.data.content, showReasoning, palette, mdTheme)
         if (component.children.length > 0) chat.addChild(component)
         break
       }
+      case 'llm/retry': {
+        clearStreaming()
+        appendNotice(
+          `Retrying model request (${event.data.retry}/${event.data.maxRetries}) in ${event.data.delayMs}ms: ${event.data.failure.message}`,
+          'warning',
+        )
+        break
+      }
       case 'tool/call':
         chat.addChild(new Spacer(1))
         chat.addChild(parsedTool(event))
@@ -970,9 +1006,13 @@ export function createTuiChat(
         todo.update(event.data.todos)
         break
       case 'turn/end':
+        clearStreaming()
         if (event.data.reason.kind === 'error') {
           const key = `${event.data.turn}:${event.data.reason.step}`
-          if (!liveErrors.delete(key)) appendNotice(event.data.reason.message, 'error')
+          const message = 'failure' in event.data.reason
+            ? event.data.reason.failure.message
+            : event.data.reason.message
+          if (!liveErrors.delete(key)) appendNotice(message, 'error')
         } else if (event.data.reason.kind === 'aborted') {
           appendNotice(event.data.reason.reason ?? 'Turn cancelled.', 'warning')
         } else if (event.data.reason.kind === 'max-tokens') {
@@ -1245,10 +1285,7 @@ export function createTuiChat(
 
   const disposeSessionEvents = ctx.on('session/event', (session, event) => {
     if (session !== agent.session) return
-    if (event.type === 'assistant/message' && event.data.usage !== undefined) {
-      tokens.input += event.data.usage.inputTokens
-      tokens.output += event.data.usage.outputTokens
-    }
+    recordEventUsage(tokens, event)
     if ('surfaceOp' in event && typeof event.surfaceOp === 'object') {
       rebuildTranscript(false)
       return