From be70e46a618ba52f6a80a93abc811a11ef4ca860 Mon Sep 17 00:00:00 2001 From: disrupted Date: Sat, 26 Sep 2026 13:29:32 +0200 Subject: [PATCH 1/2] fix: source session cost from API --- lua/opencode/state/renderer.lua | 14 +------------- lua/opencode/state/session.lua | 2 ++ lua/opencode/types.lua | 2 ++ lua/opencode/ui/renderer/events.lua | 18 ++++++------------ 4 files changed, 11 insertions(+), 25 deletions(-) diff --git a/lua/opencode/state/renderer.lua b/lua/opencode/state/renderer.lua index f2e7c0ed1..f12d8d470 100644 --- a/lua/opencode/state/renderer.lua +++ b/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 diff --git a/lua/opencode/state/session.lua b/lua/opencode/state/session.lua index b06ad999d..5aa5ae9d5 100644 --- a/lua/opencode/state/session.lua +++ b/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() diff --git a/lua/opencode/types.lua b/lua/opencode/types.lua index 9ede16c9a..c1b96324f 100644 --- a/lua/opencode/types.lua +++ b/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 diff --git a/lua/opencode/ui/renderer/events.lua b/lua/opencode/ui/renderer/events.lua index f9aab59bb..03fc1a433 100644 --- a/lua/opencode/ui/renderer/events.lua +++ b/lua/opencode/ui/renderer/events.lua @@ -111,7 +111,7 @@ local function replay_orphan_parts(message_id, revert_index) end end ----Update token/cost stats in state from a message +---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 @@ -119,12 +119,8 @@ local function update_stats(message) 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 + if 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) end end @@ -413,13 +409,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.input > 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(tokens.input + tokens.output + tokens.cache.read + tokens.cache.write) end return end @@ -550,6 +542,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_')) From 521d696b7ddf30d43443377ace4fe97d1ace007e Mon Sep 17 00:00:00 2001 From: disrupted Date: Sat, 26 Sep 2026 13:51:22 +0200 Subject: [PATCH 2/2] fix(session): align token count with TUI --- lua/opencode/ui/renderer/events.lua | 14 +++-- tests/unit/token_count_spec.lua | 90 +++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 4 deletions(-) create mode 100644 tests/unit/token_count_spec.lua diff --git a/lua/opencode/ui/renderer/events.lua b/lua/opencode/ui/renderer/events.lua index 03fc1a433..140ee656b 100644 --- a/lua/opencode/ui/renderer/events.lua +++ b/lua/opencode/ui/renderer/events.lua @@ -111,6 +111,12 @@ local function replay_orphan_parts(message_id, revert_index) end end +---@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) @@ -119,8 +125,8 @@ local function update_stats(message) end local tokens = message.info.tokens - if tokens and tokens.input > 0 then - state.renderer.set_tokens_count(tokens.input + tokens.output + tokens.cache.read + tokens.cache.write) + if tokens and tokens.output > 0 then + state.renderer.set_tokens_count(total_tokens(tokens)) end end @@ -409,9 +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 and part.tokens.input > 0 then + if part.type == 'step-finish' and part.tokens and part.tokens.output > 0 then local tokens = part.tokens - state.renderer.set_tokens_count(tokens.input + tokens.output + tokens.cache.read + tokens.cache.write) + state.renderer.set_tokens_count(total_tokens(tokens)) end return end diff --git a/tests/unit/token_count_spec.lua b/tests/unit/token_count_spec.lua new file mode 100644 index 000000000..c131743ee --- /dev/null +++ b/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)