fix: tighten tool markup parse and harden custom-model safety#9
Conversation
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
|
Reviewer's GuideTightens tool-call parsing to avoid executing bare call: markup, narrows ambiguous tool name aliases, and ports generation safety (stop/truncation/sanitize) into the custom-model streaming path. Sequence diagram for custom-model streaming with GenerationSafety and tool parsingsequenceDiagram
participant ModelOrchestrator
participant GenerationSafety
participant ToolCallParser
ModelOrchestrator->>GenerationSafety: maxOutputCharsForCustom()
GenerationSafety-->>ModelOrchestrator: maxOutputChars
loop onTokenReceived
ModelOrchestrator->>GenerationSafety: safetyStopMessage(fullResponse, maxOutputChars)
alt safetyStopMessage returns safetyMsg
GenerationSafety-->>ModelOrchestrator: safetyMsg
ModelOrchestrator->>ModelOrchestrator: _sanitizeAssistantText(fullResponse + safetyMsg)
ModelOrchestrator->>ModelOrchestrator: stopGeneration()
else safetyStopMessage returns null
GenerationSafety-->>ModelOrchestrator: null
ModelOrchestrator->>ModelOrchestrator: _tryParseFunctionCalls(textBuffer.toString())
ModelOrchestrator-->>ToolCallParser: _tryParseFunctionCalls(textBuffer.toString())
end
end
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Co-authored-by: Cursor <cursoragent@cursor.com>
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've left some high level feedback:
- The safety stop check runs on every token against the full
fullResponsestring, which may become expensive for long generations; consider basing the check on the buffer length or an incremental counter to avoid repeated work on growing strings. - After appending the safety stop message and breaking,
fullResponseis sanitized before yielding buttextBufferand any subsequent processing still use the unsanitized buffer; if downstream consumers rely ontextBuffer, consider aligning its contents with the sanitized output or documenting the expected divergence.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The safety stop check runs on every token against the full `fullResponse` string, which may become expensive for long generations; consider basing the check on the buffer length or an incremental counter to avoid repeated work on growing strings.
- After appending the safety stop message and breaking, `fullResponse` is sanitized before yielding but `textBuffer` and any subsequent processing still use the unsanitized buffer; if downstream consumers rely on `textBuffer`, consider aligning its contents with the sanitized output or documenting the expected divergence.Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces safety-stop checks during custom model generation in ModelOrchestrator and refactors ToolCallParser to prevent false positives by requiring explicit delimiters and removing short aliases. Feedback highlights three key issues: appending the safety message before sanitizing the text can cause the message to be stripped, the custom model stream generation loop lacks error handling for raw exceptions, and the tool call regex parser fails to handle nested JSON arguments due to lazy matching.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| final safetyMsg = GenerationSafety.safetyStopMessage( | ||
| fullResponse, | ||
| maxOutputChars, | ||
| ); | ||
| if (safetyMsg != null) { | ||
| safetyStopped = true; | ||
| fullResponse = '$fullResponse$safetyMsg'; | ||
| yield InferenceResult( | ||
| text: _sanitizeAssistantText(fullResponse), | ||
| model: selector.primaryHeavy, | ||
| isStreaming: true, | ||
| thinking: currentThinking, | ||
| toolCalls: allToolCalls.isNotEmpty ? allToolCalls : null, | ||
| inferenceTimeMs: inferenceStopwatch.elapsedMilliseconds, | ||
| ); | ||
| unawaited(stopGeneration()); | ||
| break; | ||
| } |
There was a problem hiding this comment.
If the model is safety-stopped (due to length limit or repetition) while in the middle of generating a tool call, fullResponse will end with an incomplete tool call markup (e.g., <|tool_call>call:open_app{). Appending safetyMsg directly to fullResponse and then calling _sanitizeAssistantText(fullResponse) will cause the safety message to be matched by the lazy stripMarkup regex (which strips everything from <|tool_call> to the end of the string $). As a result, the safety message is completely stripped and never shown to the user.
Sanitizing the text before appending the safety message prevents it from being stripped. Note that the same issue exists in the standard processMessage path.
| final safetyMsg = GenerationSafety.safetyStopMessage( | |
| fullResponse, | |
| maxOutputChars, | |
| ); | |
| if (safetyMsg != null) { | |
| safetyStopped = true; | |
| fullResponse = '$fullResponse$safetyMsg'; | |
| yield InferenceResult( | |
| text: _sanitizeAssistantText(fullResponse), | |
| model: selector.primaryHeavy, | |
| isStreaming: true, | |
| thinking: currentThinking, | |
| toolCalls: allToolCalls.isNotEmpty ? allToolCalls : null, | |
| inferenceTimeMs: inferenceStopwatch.elapsedMilliseconds, | |
| ); | |
| unawaited(stopGeneration()); | |
| break; | |
| } | |
| final safetyMsg = GenerationSafety.safetyStopMessage( | |
| fullResponse, | |
| maxOutputChars, | |
| ); | |
| if (safetyMsg != null) { | |
| safetyStopped = true; | |
| final sanitized = _sanitizeAssistantText(fullResponse); | |
| fullResponse = '$sanitized$safetyMsg'; | |
| yield InferenceResult( | |
| text: fullResponse, | |
| model: selector.primaryHeavy, | |
| isStreaming: true, | |
| thinking: currentThinking, | |
| toolCalls: allToolCalls.isNotEmpty ? allToolCalls : null, | |
| inferenceTimeMs: inferenceStopwatch.elapsedMilliseconds, | |
| ); | |
| unawaited(stopGeneration()); | |
| break; | |
| } |
| @@ -1623,6 +1648,7 @@ class ModelOrchestrator { | |||
| } | |||
There was a problem hiding this comment.
Unlike the standard processMessage stream generation path, _processWithCustomModel does not have any catch blocks around its stream generation loop. If chat.generateChatResponseAsync() or any tool execution throws a PlatformException or Exception, the stream will terminate abruptly and propagate the raw exception to the UI instead of yielding a user-friendly error message (e.g., ⚠️ Inference failed...). Consider wrapping the loop in a try-catch block to handle exceptions gracefully, similar to how it is done in processMessage.
| final pattern = RegExp( | ||
| r'(?:<\|?tool_call\|?>)?\s*call:\s*([a-zA-Z0-9_]+)\s*(\{[\s\S]*?\})\s*' | ||
| r'<\|?tool_call\|?>\s*call:\s*([a-zA-Z0-9_]+)\s*(\{[\s\S]*?\})\s*' | ||
| r'(?:</?\|?tool_call\|?>|<tool_call\|>)?', | ||
| caseSensitive: false, | ||
| ); |
There was a problem hiding this comment.
The lazy match (\\{[\\s\\S]*?\\}) in the regular expression will stop at the first closing brace }. If any tool call arguments contain nested JSON objects (e.g., {"outer": {"inner": "value"}}), this pattern will truncate the arguments and fail to parse them correctly. While current tools may only use flat arguments, consider using a more robust parser or noting this limitation if nested arguments are expected in the future.
Code Review SummaryStatus: No Issues Found | Recommendation: Merge Files Reviewed (3 files)
Reviewed by step-3.7-flash · Input: 215.4K · Output: 60.4K · Cached: 3.5M |
1 participant
Summary
Follow-up to review feedback on #8 (already released as v0.4.1).
Test plan
Made with Cursor
Summary by Sourcery
Harden tool-call parsing and custom-model streaming safety.
New Features:
Bug Fixes:
Enhancements:
Tests: