Feature/realtime tool display#2
Conversation
- Add tool callback system for real-time notifications - Display tool execution start, completion, and error messages - Add processing status message during AI processing - Remove old tool_used message system - Add tool-specific message styling with color coding - Improve UI state management during tool execution - All messages and comments converted to English
WalkthroughIntroduces a real-time tool callback propagated through MCP client/server and tool execution, emitting tool_start/tool_complete/tool_error events. WebSocket flow in app.py sends processing and completion messages and relays tool events. Frontend JS handles new message types and updates typing/streaming states. CSS adds styles for processing and tool event messages. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant UI as Browser UI
participant WS as WebSocket Server (app.py)
participant Agent as Agent
participant MCP as MCP Client/Server
participant Tool as Tool
UI->>WS: User message
WS-->>UI: processing
WS->>MCP: start(tool_callback)
Note over WS,MCP: tool_callback registered for real-time events
WS->>Agent: run(message, tools)
rect rgba(200,230,255,0.25)
Note right of Agent: Tool invocation path
Agent->>MCP: execute(tool, args)
MCP-->>WS: tool_callback(tool_start, name, args)
WS-->>UI: tool_start (name, args)
MCP->>Tool: run(args)
alt success
Tool-->>MCP: result
MCP-->>WS: tool_callback(tool_complete, name, result)
WS-->>UI: tool_complete (name, result)
else error
Tool--xMCP: error
MCP-->>WS: tool_callback(tool_error, name, err)
WS-->>UI: tool_error (name, err)
end
end
Agent-->>WS: final_text_response or none
alt has final text
WS-->>UI: complete (text)
else tools used, no text
WS-->>UI: complete (tools used)
else no text, no tools
WS-->>UI: complete (no output)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ 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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
mcp_client.py (2)
648-675: Keep tool execution resilient when the callback fails.Every
await tool_callback(...)sits inside the main try/except. If the callback raises (e.g., websocket send_text fails), the bareraisenever runs and the callback exception propagates instead of the original tool result/error. That means a notification glitch can abort or misreport the tool execution. Please wrap the callback invocation so its failures are logged but never mask the tool outcome. One way:async def execute_tool(**kwargs) -> Any: """Execute the tool with the given arguments.""" if not self._initialized: await self.initialize() + async def notify(event: str, payload: Any) -> None: + if not tool_callback: + return + try: + await tool_callback(event, tool.name, payload) + except Exception as callback_error: + logger.error( + "Tool callback %s for '%s' on server '%s' failed: %s", + event, + tool.name, + self.name, + callback_error, + ) + # Apply rate limiting if not self._rate_limiter.is_allowed(): logger.warning(f"Rate limit exceeded for server '{self.name}'") raise RuntimeError(f"Rate limit exceeded. Please try again later.") # Limit concurrent calls @@ try: - # Notify tool execution start - if tool_callback: - await tool_callback("tool_start", tool.name, kwargs) + # Notify tool execution start + await notify("tool_start", kwargs) @@ - # Notify tool execution completion - if tool_callback: - await tool_callback("tool_complete", tool.name, result) + # Notify tool execution completion + await notify("tool_complete", result) return result except Exception as e: logger.error(f"Error executing tool '{tool.name}' on server '{self.name}': {e}") - # Notify tool execution error - if tool_callback: - await tool_callback("tool_error", tool.name, str(e)) + # Notify tool execution error + await notify("tool_error", str(e)) raise
628-691: Remove the duplicatecreate_tool_instancedefinition.The class already defines
create_tool_instanceearlier (Line 530). Keeping both triggers Ruff F811 redefinition errors and leaves the old implementation (with global rate limiting) dead but still in the source. Please merge the callback changes into the original method and drop the duplicate block—otherwise lint will fail and the divergence is confusing. For example, delete the old method and keep this enhanced version (adjusting imports/rate-limiter usage as needed). As per static analysis hints.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
app.py(3 hunks)mcp_client.py(7 hunks)static/chat.js(4 hunks)static/styles.css(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
mcp_client.py (1)
app.py (1)
tool_callback(442-464)
app.py (1)
mcp_client.py (2)
MCPClient(84-308)start(198-241)
🪛 Ruff (0.13.3)
mcp_client.py
628-628: Redefinition of unused create_tool_instance from line 530
(F811)
669-669: Consider moving this statement to an else block
(TRY300)
671-671: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
app.py
461-461: Use explicit conversion flag
Replace with conversion flag
(RUF010)
463-463: Do not catch blind exception: Exception
(BLE001)
464-464: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
| if (messageType.startsWith('tool-')) { | ||
| const contentDiv = messageDiv.querySelector('.message-content'); | ||
| contentDiv.innerHTML = content; // Allow HTML for emojis and styling | ||
| } else { | ||
| messageDiv.querySelector('.message-content').textContent = content; | ||
| } |
There was a problem hiding this comment.
Block HTML injection in tool event messages.
contentDiv.innerHTML = content renders tool names and error text as HTML whenever the server sends a tool_* event. Because those strings include tool_name, a malicious or compromised MCP server could craft a tool name like <img src=x onerror=...> and achieve XSS in every connected browser. Use textContent (or sanitize explicitly) so tool state updates stay safe.
Fix suggestion:
- const contentDiv = messageDiv.querySelector('.message-content');
- contentDiv.innerHTML = content; // Allow HTML for emojis and styling
+ const contentDiv = messageDiv.querySelector('.message-content');
+ contentDiv.textContent = content;📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if (messageType.startsWith('tool-')) { | |
| const contentDiv = messageDiv.querySelector('.message-content'); | |
| contentDiv.innerHTML = content; // Allow HTML for emojis and styling | |
| } else { | |
| messageDiv.querySelector('.message-content').textContent = content; | |
| } | |
| if (messageType.startsWith('tool-')) { | |
| const contentDiv = messageDiv.querySelector('.message-content'); | |
| contentDiv.textContent = content; | |
| } else { | |
| messageDiv.querySelector('.message-content').textContent = content; | |
| } |
🤖 Prompt for AI Agents
In static/chat.js around lines 322 to 327, tool-event messages are using
contentDiv.innerHTML = content which allows untrusted HTML from tool
names/errors to execute (XSS); change this to set textContent (or otherwise
sanitize the content) so tool state updates are rendered as plain text. Replace
the innerHTML assignment with contentDiv.textContent = content and ensure any
legitimate emoji/styling is handled by a separate, explicit sanitizer or
rendering path that only allows safe tokens.
Add Real-time Tool Execution Display
What Changed
Files Modified
app.py: Added tool callback systemmcp_client.py: Enhanced with real-time notificationsstatic/chat.js: Added new message handlersstatic/styles.css: Added tool message stylingBenefits
Summary by CodeRabbit