fix(session): align token count and cost with TUI

lua/opencode/state/renderer.lua

@@ -25,7 +25,7 @@ end
 
 ---@param cost number
 function M.set_cost(cost)
-  if not cost or cost <= 0 then
+  if cost == nil then
     return
   end
   return store.set('cost', cost)
@@ -36,23 +36,11 @@ function M.set_tokens_count(count)
   return store.set('tokens_count', count)
 end
 
----@param tokens_count number
----@param cost number
-function M.set_stats(tokens_count, cost)
-  return store.batch(function()
-    store.set('tokens_count', tokens_count)
-    if cost and cost > 0 then
-      store.set('cost', cost)
-    end
-  end)
-end
-
 function M.reset()
   return store.batch(function()
     store.set('messages', {})
     store.set('current_message', nil)
     store.set('tokens_count', 0)
-    store.set('cost', 0)
     store.set('pending_permissions', {})
   end)
 end

lua/opencode/state/session.lua

@@ -20,6 +20,7 @@ function M.set_active(session)
     store.set('restore_points', {})
     store.set('last_sent_context', nil)
     store.set('user_message_count', {})
+    store.set('cost', type(session) == 'table' and session.cost or 0)
     return store.set('active_session', session)
   end)
   session_tabs.sync()
@@ -38,6 +39,7 @@ function M.clear_active()
     store.set('restore_points', {})
     store.set('last_sent_context', nil)
     store.set('user_message_count', {})
+    store.set('cost', 0)
     return store.set('active_session', nil)
   end)
   session_tabs.sync()

lua/opencode/types.lua

@@ -135,6 +135,8 @@
 ---@field agent string|nil
 ---@field model { id: string, providerID: string, variant?: string }|nil
 ---@field directory? string
+---@field cost? number
+---@field tokens? { input: number, output: number, reasoning: number, cache: { read: number, write: number } }
 ---@field revert? SessionRevertInfo
 ---@field share? SessionShareInfo
 

lua/opencode/ui/renderer/events.lua

@@ -111,20 +111,22 @@ local function replay_orphan_parts(message_id, revert_index)
   end
 end
 
----Update token/cost stats in state from a message
+---@param tokens MessageTokenCount
+---@return number
+local function total_tokens(tokens)
+  return tokens.input + tokens.output + (tokens.reasoning or 0) + tokens.cache.read + tokens.cache.write
+end
+
+---Update token stats in state from a message
 ---@param message OpencodeMessage
 local function update_stats(message)
   if not state.current_model and message.info.providerID and message.info.providerID ~= '' then
     state.model.set_model(message.info.providerID .. '/' .. message.info.modelID)
   end
 
   local tokens = message.info.tokens
-  if tokens and tokens.input > 0 and message.info.cost and type(message.info.cost) == 'number' then
-    state.renderer.set_stats(tokens.input + tokens.output + tokens.cache.read + tokens.cache.write, message.info.cost)
-  elseif tokens and tokens.input > 0 then
-    state.renderer.set_tokens_count(tokens.input + tokens.output + tokens.cache.read + tokens.cache.write)
-  elseif message.info.cost and type(message.info.cost) == 'number' then
-    state.renderer.set_cost(message.info.cost)
+  if tokens and tokens.output > 0 then
+    state.renderer.set_tokens_count(total_tokens(tokens))
   end
 end
 
@@ -413,13 +415,9 @@ function M.on_part_updated(properties, revert_index)
   message.parts[existing_part_index or #message.parts + 1] = part
 
   if part.type == 'step-start' or part.type == 'step-finish' then
-    if part.type == 'step-finish' and part.tokens then
+    if part.type == 'step-finish' and part.tokens and part.tokens.output > 0 then
       local tokens = part.tokens
-      if tokens.input > 0 and part.cost and type(part.cost) == 'number' then
-        state.renderer.set_stats(tokens.input + tokens.output + tokens.cache.read + tokens.cache.write, part.cost)
-      elseif tokens.input > 0 then
-        state.renderer.set_tokens_count(tokens.input + tokens.output + tokens.cache.read + tokens.cache.write)
-      end
+      state.renderer.set_tokens_count(total_tokens(tokens))
     end
     return
   end
@@ -550,6 +548,8 @@ function M.on_session_updated(properties)
     state.store.set_raw('active_session', updated_session)
   end
 
+  state.renderer.set_cost(updated_session.cost or 0)
+
   if revert_changed then
     local real_messages = vim.tbl_filter(function(msg)
       return not (msg.info and msg.info.id and msg.info.id:match('^__opencode_'))

tests/unit/token_count_spec.lua

@@ -0,0 +1,90 @@
+local assert = require('luassert')
+local stub = require('luassert.stub')
+
+local state = require('opencode.state')
+local events = require('opencode.ui.renderer.events')
+local flush = require('opencode.ui.renderer.flush')
+local renderer = require('opencode.ui.renderer')
+local ctx = require('opencode.ui.renderer.ctx')
+
+local function assistant_message(id, tokens)
+  return {
+    info = {
+      id = id,
+      role = 'assistant',
+      sessionID = 'ses_1',
+      providerID = 'test-provider',
+      modelID = 'test-model',
+      tokens = tokens,
+    },
+    parts = {},
+  }
+end
+
+local REASONING_TOKENS = { input = 1000, output = 200, reasoning = 300, cache = { read = 50, write = 25 } }
+local REASONING_TOTAL = 1575
+
+describe('context token count', function()
+  local scroll_stub
+  local reconcile_stub
+  local mark_dirty_stub
+
+  before_each(function()
+    scroll_stub = stub(renderer, 'scroll_to_bottom')
+    reconcile_stub = stub(renderer, 'reconcile_rendered_message_limit')
+    mark_dirty_stub = stub(flush, 'mark_message_dirty')
+
+    ctx:reset()
+    state.session.set_active({ id = 'ses_1' })
+    state.renderer.set_messages({})
+    state.renderer.set_tokens_count(0)
+    state.model.clear()
+  end)
+
+  after_each(function()
+    scroll_stub:revert()
+    reconcile_stub:revert()
+    mark_dirty_stub:revert()
+
+    ctx:reset()
+    state.session.clear_active()
+    state.renderer.set_messages({})
+    state.renderer.set_tokens_count(0)
+    state.model.clear()
+  end)
+
+  it('includes reasoning tokens in the context total', function()
+    events.on_message_updated(assistant_message('msg_1', REASONING_TOKENS))
+
+    assert.equal(REASONING_TOTAL, state.tokens_count)
+  end)
+
+  it('ignores responses without output tokens, matching the TUI', function()
+    events.on_message_updated(assistant_message('msg_1', REASONING_TOKENS))
+    assert.equal(REASONING_TOTAL, state.tokens_count)
+
+    events.on_message_updated(
+      assistant_message('msg_2', { input = 100, output = 0, reasoning = 500, cache = { read = 0, write = 0 } })
+    )
+
+    assert.equal(REASONING_TOTAL, state.tokens_count)
+  end)
+
+  it('derives the same total from step-finish parts', function()
+    local message =
+      assistant_message('msg_1', { input = 0, output = 0, reasoning = 0, cache = { read = 0, write = 0 } })
+    state.renderer.set_messages({ message })
+
+    events.on_part_updated({
+      part = {
+        id = 'part_1',
+        messageID = 'msg_1',
+        sessionID = 'ses_1',
+        type = 'step-finish',
+        tokens = REASONING_TOKENS,
+      },
+    })
+
+    assert.equal(REASONING_TOTAL, state.tokens_count)
+  end)
+end)