Skip to content

Commit 6efcb2e

Browse files
committed
feat(ui/renderer): support orphan message parts and replay on message arrival
Add orphan-parts storage and indexed operations to RenderState: - upsert_orphan_part, consume_orphan_parts, remove_orphan_part, clear_orphan_parts Replay orphan parts when a message is set/updated so parts that arrive before their parent message are applied as soon as the message exists. Clear orphan parts on message removal and handle orphan removals during part deletion. Add test fixtures for the part-before-message scenario and update replay tests to include the new case.
1 parent 43c5fa7 commit 6efcb2e

5 files changed

Lines changed: 318 additions & 1 deletion

File tree

lua/opencode/ui/render_state.lua

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ end
3232
function RenderState:reset()
3333
self._messages = {}
3434
self._parts = {}
35+
self._orphan_parts = {}
36+
self._orphan_parts_index = {}
3537
self._part_ranges = {}
3638
self._message_ranges = {}
3739
self._ranges_valid = false
@@ -205,6 +207,82 @@ function RenderState:get_message(message_id)
205207
return self._messages[message_id]
206208
end
207209

210+
---@param message_id string
211+
---@param part OpencodeMessagePart
212+
function RenderState:upsert_orphan_part(message_id, part)
213+
if not message_id or not part or not part.id then
214+
return
215+
end
216+
217+
local orphan_parts = self._orphan_parts[message_id]
218+
if not orphan_parts then
219+
orphan_parts = {}
220+
self._orphan_parts[message_id] = orphan_parts
221+
self._orphan_parts_index[message_id] = {}
222+
end
223+
224+
local orphan_index = self._orphan_parts_index[message_id]
225+
local idx = orphan_index[part.id]
226+
if idx then
227+
orphan_parts[idx] = part
228+
else
229+
orphan_parts[#orphan_parts + 1] = part
230+
orphan_index[part.id] = #orphan_parts
231+
end
232+
end
233+
234+
---@param message_id string
235+
---@return OpencodeMessagePart[]
236+
function RenderState:consume_orphan_parts(message_id)
237+
if not message_id then
238+
return {}
239+
end
240+
241+
local orphan_parts = self._orphan_parts[message_id] or {}
242+
self._orphan_parts[message_id] = nil
243+
self._orphan_parts_index[message_id] = nil
244+
return orphan_parts
245+
end
246+
247+
---@param message_id string
248+
---@param part_id string
249+
---@return boolean
250+
function RenderState:remove_orphan_part(message_id, part_id)
251+
local orphan_parts = message_id and self._orphan_parts[message_id]
252+
local orphan_index = message_id and self._orphan_parts_index[message_id]
253+
local idx = orphan_index and orphan_index[part_id]
254+
if not idx then
255+
return false
256+
end
257+
258+
table.remove(orphan_parts, idx)
259+
orphan_index[part_id] = nil
260+
261+
for i = idx, #orphan_parts do
262+
local part = orphan_parts[i]
263+
if part and part.id then
264+
orphan_index[part.id] = i
265+
end
266+
end
267+
268+
if #orphan_parts == 0 then
269+
self._orphan_parts[message_id] = nil
270+
self._orphan_parts_index[message_id] = nil
271+
end
272+
273+
return true
274+
end
275+
276+
---@param message_id string
277+
function RenderState:clear_orphan_parts(message_id)
278+
if not message_id then
279+
return
280+
end
281+
282+
self._orphan_parts[message_id] = nil
283+
self._orphan_parts_index[message_id] = nil
284+
end
285+
208286
---@param line integer 1-indexed
209287
---@return RenderedMessage?
210288
function RenderState:get_message_at_line(line)

lua/opencode/ui/renderer/events.lua

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,15 @@ end
4141

4242
local M = {}
4343

44+
---@param message_id string
45+
---@param revert_index? integer
46+
local function replay_orphan_parts(message_id, revert_index)
47+
local orphan_parts = ctx.render_state:consume_orphan_parts(message_id)
48+
for _, orphan_part in ipairs(orphan_parts) do
49+
M.on_part_updated({ part = orphan_part }, revert_index)
50+
end
51+
end
52+
4453
---Update token/cost stats in state from a message
4554
---@param message OpencodeMessage
4655
local function update_stats(message)
@@ -159,6 +168,7 @@ function M.on_message_updated(message, revert_index)
159168
table.insert(state.messages, msg)
160169
end
161170
ctx.render_state:set_message(msg, 0, 0)
171+
replay_orphan_parts(msg.info.id, revert_index)
162172
return
163173
end
164174

@@ -180,6 +190,7 @@ function M.on_message_updated(message, revert_index)
180190
else
181191
table.insert(state.messages, msg)
182192
ctx.render_state:set_message(msg)
193+
replay_orphan_parts(msg.info.id)
183194
flush.mark_message_dirty(msg.info.id)
184195
state.renderer.set_current_message(msg)
185196
if message.info.role == 'user' then
@@ -204,6 +215,7 @@ function M.on_message_removed(properties)
204215
end
205216

206217
local rendered_message = ctx.render_state:get_message(message_id)
218+
ctx.render_state:clear_orphan_parts(message_id)
207219
if not rendered_message or not rendered_message.message then
208220
return
209221
end
@@ -251,7 +263,7 @@ function M.on_part_updated(properties, revert_index)
251263

252264
local rendered_message = ctx.render_state:get_message(part.messageID)
253265
if not rendered_message or not rendered_message.message then
254-
vim.notify('Could not find message for part: ' .. vim.inspect(part), vim.log.levels.WARN)
266+
ctx.render_state:upsert_orphan_part(part.messageID, part)
255267
return
256268
end
257269

@@ -344,6 +356,10 @@ function M.on_part_removed(properties)
344356
return
345357
end
346358

359+
if properties.messageID and ctx.render_state:remove_orphan_part(properties.messageID, part_id) then
360+
return
361+
end
362+
347363
-- Remove the part from the in-memory message too
348364
local cached = ctx.render_state:get_part(part_id)
349365
local message_id = cached and cached.message_id
Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
{
2+
"actions": [],
3+
"extmarks": [
4+
[
5+
1,
6+
1,
7+
0,
8+
{
9+
"ns_id": 3,
10+
"priority": 10,
11+
"right_gravity": true,
12+
"virt_text": [
13+
[
14+
"▌󰭻 ",
15+
"OpencodeMessageRoleUser"
16+
],
17+
[
18+
" "
19+
],
20+
[
21+
"USER",
22+
"OpencodeMessageRoleUser"
23+
],
24+
[
25+
"",
26+
"OpencodeHint"
27+
],
28+
[
29+
" [msg_0000000000001]",
30+
"OpencodeHint"
31+
]
32+
],
33+
"virt_text_hide": false,
34+
"virt_text_pos": "win_col",
35+
"virt_text_repeat_linebreak": false,
36+
"virt_text_win_col": -3
37+
}
38+
],
39+
[
40+
2,
41+
1,
42+
0,
43+
{
44+
"ns_id": 3,
45+
"priority": 9,
46+
"right_gravity": true,
47+
"virt_text": [
48+
[
49+
" 2023-11-14 22:13:20",
50+
"OpencodeHint"
51+
]
52+
],
53+
"virt_text_hide": false,
54+
"virt_text_pos": "right_align",
55+
"virt_text_repeat_linebreak": false
56+
}
57+
],
58+
[
59+
3,
60+
2,
61+
0,
62+
{
63+
"ns_id": 3,
64+
"priority": 4096,
65+
"right_gravity": true,
66+
"virt_text": [
67+
[
68+
"",
69+
"OpencodeMessageRoleUser"
70+
]
71+
],
72+
"virt_text_hide": false,
73+
"virt_text_pos": "win_col",
74+
"virt_text_repeat_linebreak": true,
75+
"virt_text_win_col": -3
76+
}
77+
],
78+
[
79+
4,
80+
3,
81+
0,
82+
{
83+
"ns_id": 3,
84+
"priority": 4096,
85+
"right_gravity": true,
86+
"virt_text": [
87+
[
88+
"",
89+
"OpencodeMessageRoleUser"
90+
]
91+
],
92+
"virt_text_hide": false,
93+
"virt_text_pos": "win_col",
94+
"virt_text_repeat_linebreak": true,
95+
"virt_text_win_col": -3
96+
}
97+
],
98+
[
99+
5,
100+
6,
101+
0,
102+
{
103+
"ns_id": 3,
104+
"priority": 10,
105+
"right_gravity": true,
106+
"virt_text": [
107+
[
108+
"",
109+
"OpencodeMessageRoleAssistant"
110+
],
111+
[
112+
" "
113+
],
114+
[
115+
"BUILD",
116+
"OpencodeMessageRoleAssistant"
117+
],
118+
[
119+
" claude-sonnet-4-5",
120+
"OpencodeHint"
121+
],
122+
[
123+
" [msg_0000000000002]",
124+
"OpencodeHint"
125+
]
126+
],
127+
"virt_text_hide": false,
128+
"virt_text_pos": "win_col",
129+
"virt_text_repeat_linebreak": false,
130+
"virt_text_win_col": -3
131+
}
132+
],
133+
[
134+
6,
135+
6,
136+
0,
137+
{
138+
"ns_id": 3,
139+
"priority": 9,
140+
"right_gravity": true,
141+
"virt_text": [
142+
[
143+
" 2023-11-14 22:13:21",
144+
"OpencodeHint"
145+
]
146+
],
147+
"virt_text_hide": false,
148+
"virt_text_pos": "right_align",
149+
"virt_text_repeat_linebreak": false
150+
}
151+
]
152+
],
153+
"lines": [
154+
"----",
155+
"",
156+
"",
157+
"Can you help me fix this bug?",
158+
"",
159+
"----",
160+
"",
161+
"",
162+
"Sure, I can help with that.",
163+
"",
164+
""
165+
],
166+
"timestamp": 1773091221
167+
}
Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
[
2+
{
3+
"type": "message.updated",
4+
"properties": {
5+
"info": {
6+
"id": "msg_0000000000001",
7+
"sessionID": "ses_0000000000001",
8+
"role": "user",
9+
"agent": "build",
10+
"time": { "created": 1700000000000 },
11+
"model": { "providerID": "anthropic", "modelID": "claude-sonnet-4-5" }
12+
}
13+
}
14+
},
15+
{
16+
"type": "message.part.updated",
17+
"properties": {
18+
"part": {
19+
"id": "prt_0000000000001",
20+
"messageID": "msg_0000000000001",
21+
"sessionID": "ses_0000000000001",
22+
"type": "text",
23+
"text": "Can you help me fix this bug?"
24+
}
25+
}
26+
},
27+
{
28+
"type": "message.part.delta",
29+
"properties": {
30+
"partID": "prt_0000000000002",
31+
"messageID": "msg_0000000000002",
32+
"sessionID": "ses_0000000000001",
33+
"field": "text",
34+
"delta": "Sure, I can help with that."
35+
}
36+
},
37+
{
38+
"type": "message.updated",
39+
"properties": {
40+
"info": {
41+
"id": "msg_0000000000002",
42+
"sessionID": "ses_0000000000001",
43+
"role": "assistant",
44+
"agent": "build",
45+
"mode": "build",
46+
"providerID": "anthropic",
47+
"modelID": "claude-sonnet-4-5",
48+
"parentID": "msg_0000000000001",
49+
"time": { "created": 1700000001000, "completed": 1700000050000 },
50+
"cost": 0,
51+
"tokens": { "input": 100, "output": 20, "reasoning": 0, "cache": { "read": 0, "write": 0 } }
52+
}
53+
}
54+
}
55+
]

tests/replay/renderer_spec.lua

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,7 @@ describe('renderer functional tests', function()
269269
local skip_full_session = {
270270
'permission-prompt',
271271
'permission-ask-new',
272+
'part-before-message-delta',
272273
'question-ask',
273274
'question-ask-other',
274275
'multiple-question-ask',

0 commit comments

Comments
 (0)