Repository navigation
fix(runtime-work): stabilize local task flow - #1544
Conversation
📝 WalkthroughWalkthroughThe PR forwards optional ChangesRuntime task start and streaming flow
Codex local image previews
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
executor/runtime_work/codex_discovery.py (1)
902-930: 🚀 Performance & Scalability | 🔵 TrivialPreview data URLs are embedded into cached transcript messages.
_local_image_attachmentinlines up to 5 MiB per image as a base64local_preview_url, and these messages are retained in_transcript_cache. A long Codex session referencing many local images (across many cached threads) can hold a large amount of base64 in memory. Consider bounding the number/aggregate size of inlined previews per transcript, or storing paths and encoding lazily on demand, if memory pressure becomes a concern.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@executor/runtime_work/codex_discovery.py` around lines 902 - 930, The `_local_image_attachment` helper is embedding full base64 preview data into transcript-cached message payloads, which can grow memory usage without bound. Update the attachment flow so `local_preview_url` is not eagerly stored for every image in `_transcript_cache`; instead, either cap the number or total bytes of inlined previews per transcript, or defer encoding until a consumer explicitly requests the preview. Keep the fix localized around `_local_image_attachment` and the transcript caching path that stores its returned dict.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@wework/src/features/workbench/WorkbenchProvider.tsx`:
- Around line 2160-2161: `isAwaitingAssistantStart` can get stuck true in
`WorkbenchProvider` because it is only cleared by WebSocket callbacks that pass
`isCurrentLocalTaskEvent`. Add a fallback timeout-based reset (for example in a
`useEffect` watching `isAwaitingAssistantStart`) so the flag is cleared after a
short delay if no start/done/error event arrives, and also ensure any failure
paths in `createRuntimeTask` and `sendRuntimeMessage` clear it when a task opens
but never streams.
---
Nitpick comments:
In `@executor/runtime_work/codex_discovery.py`:
- Around line 902-930: The `_local_image_attachment` helper is embedding full
base64 preview data into transcript-cached message payloads, which can grow
memory usage without bound. Update the attachment flow so `local_preview_url` is
not eagerly stored for every image in `_transcript_cache`; instead, either cap
the number or total bytes of inlined previews per transcript, or defer encoding
until a consumer explicitly requests the preview. Keep the fix localized around
`_local_image_attachment` and the transcript caching path that stores its
returned dict.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 647b621b-3881-4579-bceb-b66a87f8f69e
📒 Files selected for processing (18)
backend/app/schemas/runtime_work.pybackend/app/services/runtime_work_service.pybackend/tests/services/test_runtime_work_service.pydocs/en/developer-guide/runtime-local-work.mddocs/zh/developer-guide/runtime-local-work.mdexecutor/runtime_work/codex_discovery.pyexecutor/tests/runtime_work/test_local_task_store.pywework/src/components/chat/blocks/ToolBlockItem.tsxwework/src/components/chat/blocks/ToolBlocksDisplay.test.tsxwework/src/components/chat/blocks/ToolBlocksDisplay.tsxwework/src/components/chat/blocks/toolBlockActivity.tswework/src/components/layout/DesktopWorkbenchLayout.tsxwework/src/components/layout/DesktopWorkbenchMain.tsxwework/src/components/layout/MobileWorkbenchLayout.tsxwework/src/features/workbench/WorkbenchProvider.test.tsxwework/src/features/workbench/WorkbenchProvider.tsxwework/src/pages/WorkbenchPage.tsxwework/src/types/api.ts
| const currentRuntimeTaskBusy = | ||
| currentRuntimeTaskRunning || hasActiveTurn || isAwaitingAssistantStart |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find all sites that clear isAwaitingAssistantStart and the WS event-to-address matching logic.
rg -nP 'setIsAwaitingAssistantStart\(' wework/src
rg -nP 'isCurrentLocalTaskEvent' wework/srcRepository: wecode-ai/Wegent
Length of output: 2626
🏁 Script executed:
#!/bin/bash
# 1. Check isCurrentLocalTaskEvent implementation
sed -n '470,510p' wework/src/features/workbench/WorkbenchProvider.tsx
# 2. Check the 'set to true' context around line 2340 (likely create/send flow)
sed -n '2330,2400p' wework/src/features/workbench/WorkbenchProvider.tsx
# 3. Check the logic around line 2477 (likely await assistant start flow)
sed -n '2470,2500p' wework/src/features/workbench/WorkbenchProvider.tsx
# 4. Check startNewChat logic which clears the flag (around line 1655)
sed -n '1640,1730p' wework/src/features/workbench/WorkbenchProvider.tsxRepository: wecode-ai/Wegent
Length of output: 8518
Add a safety timeout to clear isAwaitingAssistantStart if stuck.
The flag isAwaitingAssistantStart is set to true in createRuntimeTask and sendRuntimeMessage but relies on specific WebSocket events (onChatStart, onChatDone, onChatError) passing the isCurrentLocalTaskEvent guard to clear it. If the socket drops, disconnects, or the task ID changes before those events are received, the "thinking" indicator will remain stuck indefinitely, blocking queued sends.
Add a useEffect that resets the flag after a short timeout (e.g., 5 seconds) whenever it is set, or explicitly in catch blocks for tasks that open to a new ID but fail to stream.
Current implementation details
// wework/src/features/workbench/WorkbenchProvider.tsx
// Flag is set at lines 2340 and 2477
const currentRuntimeTaskBusy =
currentRuntimeTaskRunning || hasActiveTurn || isAwaitingAssistantStart
// Cleared only if specific WS events match currentTaskRef
// (See lines 1344-1384 where isCurrentLocalTaskEvent guard protects setIsAwaitingAssistantStart(false))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@wework/src/features/workbench/WorkbenchProvider.tsx` around lines 2160 -
2161, `isAwaitingAssistantStart` can get stuck true in `WorkbenchProvider`
because it is only cleared by WebSocket callbacks that pass
`isCurrentLocalTaskEvent`. Add a fallback timeout-based reset (for example in a
`useEffect` watching `isAwaitingAssistantStart`) so the flag is cleared after a
short delay if no start/done/error event arrives, and also ensure any failure
paths in `createRuntimeTask` and `sendRuntimeMessage` clear it when a task opens
but never streams.
Summary by CodeRabbit
New Features
Bug Fixes