Skip to content

fix(runtime-work): stabilize local task flow - #1544

Merged
qdaxb merged 1 commit into
wecode-ai:mainfrom
FicoHu:codex/runtime-local-task-flow
Jun 25, 2026
Merged

qdaxb merged 1 commit into
wecode-ai:mainfrom
FicoHu:codex/runtime-local-task-flow

Conversation

@FicoHu

@FicoHu FicoHu commented Jun 25, 2026 •

Copy link
Copy Markdown
Contributor

Summary by CodeRabbit

  • New Features

    • Runtime tasks can now open immediately with a client-generated task ID, improving responsiveness while the task is being created.
    • Local image references in messages are now shown with inline preview attachments when available.
    • Tool details stay expanded during active streaming, making in-progress output easier to follow.
  • Bug Fixes

    • Improved waiting states so the app shows a clearer “thinking” indicator while assistant responses are starting.
    • Better handling of queued follow-up messages to keep their order consistent.

@coderabbitai

coderabbitai Bot commented Jun 25, 2026 •

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

The PR forwards optional localTaskId through runtime task creation, adds assistant-start waiting state and routing in the workbench, updates tool block expansion during streaming, and adds Codex transcript handling for local image references with preview attachments.

Changes

Runtime task start and streaming flow

Layer / File(s) Summary
Request contract and backend forwarding
backend/app/schemas/runtime_work.py, backend/app/services/runtime_work_service.py, backend/tests/services/test_runtime_work_service.py, wework/src/types/api.ts, docs/en/developer-guide/runtime-local-work.md, docs/zh/developer-guide/runtime-local-work.md
RuntimeTaskCreateRequest and the runtime create RPC now carry optional localTaskId, and the developer guide describes the same field in the client-back-end flow.
Assistant-start waiting state
wework/src/features/workbench/WorkbenchProvider.tsx
WorkbenchProvider generates local runtime task IDs, tracks isAwaitingAssistantStart, opens runtime tasks from generated local addresses, and gates sending and queue handling on the new busy state.
Layout propagation and workbench tests
wework/src/pages/WorkbenchPage.tsx, wework/src/components/layout/DesktopWorkbenchLayout.tsx, wework/src/components/layout/DesktopWorkbenchMain.tsx, wework/src/components/layout/MobileWorkbenchLayout.tsx, wework/src/features/workbench/WorkbenchProvider.test.tsx
WorkbenchPage and the layout components pass isAwaitingAssistantStart through to ScrollableMessageArea, and the workbench tests cover pending and queued runtime sends.
Streaming tool block expansion
wework/src/components/chat/blocks/toolBlockActivity.ts, wework/src/components/chat/blocks/ToolBlockItem.tsx, wework/src/components/chat/blocks/ToolBlocksDisplay.tsx, wework/src/components/chat/blocks/ToolBlocksDisplay.test.tsx
buildProcessingDisplayRows can skip grouping completed tool blocks during streaming, ToolBlockItem can stay expanded, and the display test covers the streaming state.

Codex local image previews

Layer / File(s) Summary
Local image transcript rendering
executor/runtime_work/codex_discovery.py
CodexSessionDiscovery extracts local image references from transcript payloads, resolves readable file:// paths, builds preview attachments, and adds them to transcript pages.
Local image tests and docs
executor/tests/runtime_work/test_local_task_store.py, docs/en/developer-guide/runtime-local-work.md, docs/zh/developer-guide/runtime-local-work.md
read_transcript_page tests cover local image content and preview metadata, and the developer guide documents the same behavior in English and Chinese.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

  • wecode-ai/Wegent#1513 — Also changes RuntimeTaskCreateRequest and backend runtime.tasks.create request forwarding.
  • wecode-ai/Wegent#1537 — Also changes executor/runtime_work/codex_discovery.py transcript reading/windowing logic.

Suggested reviewers

  • qdaxb
  • Micro66

Poem

🐰 I hop by localTaskId light,
and preview leaves now gleam so bright.
With file paths tucked in transcripts deep,
I nibble tool blocks while they sleep.
Then “正在思考” twinkles in sight.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 3.13% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title is concise and matches the main change: stabilizing the runtime local task flow.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
executor/runtime_work/codex_discovery.py (1)

902-930: 🚀 Performance & Scalability | 🔵 Trivial

Preview data URLs are embedded into cached transcript messages.

_local_image_attachment inlines up to 5 MiB per image as a base64 local_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

📥 Commits

Reviewing files that changed from the base of the PR and between 6091386 and 101c44f.

📒 Files selected for processing (18)
  • backend/app/schemas/runtime_work.py
  • backend/app/services/runtime_work_service.py
  • backend/tests/services/test_runtime_work_service.py
  • docs/en/developer-guide/runtime-local-work.md
  • docs/zh/developer-guide/runtime-local-work.md
  • executor/runtime_work/codex_discovery.py
  • executor/tests/runtime_work/test_local_task_store.py
  • wework/src/components/chat/blocks/ToolBlockItem.tsx
  • wework/src/components/chat/blocks/ToolBlocksDisplay.test.tsx
  • wework/src/components/chat/blocks/ToolBlocksDisplay.tsx
  • wework/src/components/chat/blocks/toolBlockActivity.ts
  • wework/src/components/layout/DesktopWorkbenchLayout.tsx
  • wework/src/components/layout/DesktopWorkbenchMain.tsx
  • wework/src/components/layout/MobileWorkbenchLayout.tsx
  • wework/src/features/workbench/WorkbenchProvider.test.tsx
  • wework/src/features/workbench/WorkbenchProvider.tsx
  • wework/src/pages/WorkbenchPage.tsx
  • wework/src/types/api.ts

Comment on lines +2160 to +2161
const currentRuntimeTaskBusy =
currentRuntimeTaskRunning || hasActiveTurn || isAwaitingAssistantStart

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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/src

Repository: 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.tsx

Repository: 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants