perf(renderer): index task parts by child session

Avoid scanning every message part when child session updates arrive.
Track the owning task part in render state so task tool rerenders stay constant time.

Author: sudo-tee

lua/opencode/ui/render_state.lua

@@ -39,9 +39,52 @@ function RenderState:reset()
     line_to_message = {},
   }
   self._child_session_parts = {}
+  self._child_session_task_parts = {}
+  self._task_part_child_sessions = {}
   self._line_index_valid = false
 end
 
+---@param part OpencodeMessagePart?
+---@return string?
+local function get_child_session_id_for_task_part(part)
+  if not part or part.tool ~= 'task' then
+    return nil
+  end
+
+  local part_state = part.state
+  local metadata = part_state and part_state.metadata
+
+  return metadata and metadata.sessionId or nil
+end
+
+---@param part_id string
+function RenderState:_clear_task_part_child_session(part_id)
+  local child_session_id = self._task_part_child_sessions[part_id]
+  if not child_session_id then
+    return
+  end
+
+  if self._child_session_task_parts[child_session_id] == part_id then
+    self._child_session_task_parts[child_session_id] = nil
+  end
+
+  self._task_part_child_sessions[part_id] = nil
+end
+
+---@param part_id string
+---@param part OpencodeMessagePart
+function RenderState:_index_task_part_child_session(part_id, part)
+  self:_clear_task_part_child_session(part_id)
+
+  local child_session_id = get_child_session_id_for_task_part(part)
+  if not child_session_id then
+    return
+  end
+
+  self._child_session_task_parts[child_session_id] = part_id
+  self._task_part_child_sessions[part_id] = child_session_id
+end
+
 ---Get parts for a child session
 ---@param session_id string
 ---@return OpencodeMessagePart[]?|nil
@@ -52,6 +95,17 @@ function RenderState:get_child_session_parts(session_id)
   return self._child_session_parts and self._child_session_parts[session_id]
 end
 
+---Get the owning task part for a child session
+---@param session_id string
+---@return string?
+function RenderState:get_task_part_by_child_session(session_id)
+  if not session_id then
+    return nil
+  end
+
+  return self._child_session_task_parts and self._child_session_task_parts[session_id]
+end
+
 ---Upsert a part associated with a child session
 ---@param session_id string
 ---@param part OpencodeMessagePart
@@ -246,6 +300,8 @@ function RenderState:set_part(part, line_start, line_end)
   if line_start and line_end then
     self._line_index_valid = false
   end
+
+  self:_index_task_part_child_session(part_id, part)
 end
 
 ---Update part line positions and shift subsequent content
@@ -290,6 +346,7 @@ function RenderState:update_part_data(part_ref)
   end
 
   rendered_part.part = part_ref
+  self:_index_task_part_child_session(part_ref.id, part_ref)
   return rendered_part
 end
 
@@ -356,10 +413,17 @@ end
 ---@return boolean success
 function RenderState:remove_part(part_id)
   local part_data = self._parts[part_id]
-  if not part_data or not part_data.line_start or not part_data.line_end then
+  if not part_data then
     return false
   end
 
+  self:_clear_task_part_child_session(part_id)
+
+  if not part_data.line_start or not part_data.line_end then
+    self._parts[part_id] = nil
+    return true
+  end
+
   local line_count = part_data.line_end - part_data.line_start + 1
   local shift_from = part_data.line_end + 1
 

lua/opencode/ui/renderer.lua

@@ -1136,23 +1136,12 @@ end
 ---Find and re-render the task tool part in the active session that owns a given child session
 ---@param child_session_id string The child session ID to look up
 function M._rerender_task_tool_for_child_session(child_session_id)
-  if not state.messages then
+  local part_id = M._render_state:get_task_part_by_child_session(child_session_id)
+  if not part_id then
     return
   end
 
-  for _, message in ipairs(state.messages) do
-    for _, part in ipairs(message.parts or {}) do
-      if
-        part.tool == 'task'
-        and part.state
-        and part.state.metadata
-        and part.state.metadata.sessionId == child_session_id
-      then
-        M._rerender_part(part.id)
-        return
-      end
-    end
-  end
+  M._rerender_part(part_id)
 end
 
 ---Re-render existing part with current state

tests/unit/render_state_spec.lua

@@ -103,6 +103,23 @@ describe('RenderState', function()
       assert.is_table(result.actions)
       assert.equals(0, #result.actions)
     end)
+
+    it('indexes task parts by child session ID', function()
+      local part = {
+        id = 'part1',
+        messageID = 'msg1',
+        tool = 'task',
+        state = {
+          metadata = {
+            sessionId = 'child-1',
+          },
+        },
+      }
+
+      render_state:set_part(part, 1, 2)
+
+      assert.equals('part1', render_state:get_task_part_by_child_session('child-1'))
+    end)
   end)
 
   describe('get_part_at_line', function()
@@ -332,6 +349,25 @@ describe('RenderState', function()
       local success = render_state:remove_part('nonexistent')
       assert.is_false(success)
     end)
+
+    it('clears child session index when removing unrendered task parts', function()
+      local part = {
+        id = 'part1',
+        messageID = 'msg1',
+        tool = 'task',
+        state = {
+          metadata = {
+            sessionId = 'child-1',
+          },
+        },
+      }
+
+      render_state:set_part(part)
+      local success = render_state:remove_part('part1')
+
+      assert.is_true(success)
+      assert.is_nil(render_state:get_task_part_by_child_session('child-1'))
+    end)
   end)
 
   describe('remove_message', function()
@@ -487,5 +523,36 @@ describe('RenderState', function()
     it('does nothing for non-existent part', function()
       render_state:update_part_data({ id = 'nonexistent' })
     end)
+
+    it('updates child session index when task metadata changes', function()
+      local original = {
+        id = 'part1',
+        content = 'original',
+        messageID = 'msg1',
+        tool = 'task',
+        state = {
+          metadata = {
+            sessionId = 'child-1',
+          },
+        },
+      }
+      local updated = {
+        id = 'part1',
+        content = 'updated',
+        messageID = 'msg1',
+        tool = 'task',
+        state = {
+          metadata = {
+            sessionId = 'child-2',
+          },
+        },
+      }
+
+      render_state:set_part(original, 10, 15)
+      render_state:update_part_data(updated)
+
+      assert.is_nil(render_state:get_task_part_by_child_session('child-1'))
+      assert.equals('part1', render_state:get_task_part_by_child_session('child-2'))
+    end)
   end)
 end)