Skip to content

fix: tighten tool markup parse and harden custom-model safety #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
involvex merged 1 commit into from
Jul 19, 2026
Merged

fix: tighten tool markup parse and harden custom-model safety #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 27 additions & 1 deletion lib/services/model_orchestrator.dart
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1511,6 +1511,8 @@ class ModelOrchestrator {
bool hasPendingToolCalls = true;
int toolRounds = 0;
final List<Map<String, dynamic>> allToolCalls = [];
var safetyStopped = false;
final maxOutputChars = GenerationSafety.maxOutputCharsForCustom();

_isStreaming = true;
_streamingCompleter = Completer<void>();
Expand All @@ -1529,6 +1531,25 @@ class ModelOrchestrator {
fullResponse += token;
textBuffer.write(token);

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;
}
Comment on lines +1534 to +1551

@gemini-code-assist gemini-code-assist Bot Jul 19, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

high

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.

Suggested change
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;
}


final parsedCalls = _tryParseFunctionCalls(textBuffer.toString());
if (parsedCalls != null && parsedCalls.isNotEmpty) {
fullResponse = _sanitizeAssistantText(fullResponse);
Expand Down Expand Up @@ -1587,7 +1608,7 @@ class ModelOrchestrator {
} else if (event is ThinkingResponse) {
currentThinking = event.content;
yield InferenceResult(
text: fullResponse,
text: _sanitizeAssistantText(fullResponse),
model: selector.primaryHeavy,
isStreaming: true,
thinking: currentThinking,
Expand All @@ -1597,6 +1618,10 @@ class ModelOrchestrator {
}
}

if (safetyStopped || (_streamingCompleter?.isCompleted ?? false)) {
break;
}

if (hasPendingToolCalls) {
final toolResultMessages = <Message>[];
for (final tc in allToolCalls.where(
Expand All @@ -1623,6 +1648,7 @@ class ModelOrchestrator {
}

@gemini-code-assist gemini-code-assist Bot Jul 19, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

high

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.


inferenceStopwatch.stop();
fullResponse = _sanitizeAssistantText(fullResponse);

yield InferenceResult(
text: fullResponse,
Expand Down
20 changes: 10 additions & 10 deletions lib/utils/tool_call_parser.dart
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,26 +5,27 @@ class ToolCallParser {
const ToolCallParser._();

/// Common hallucinated names → Nova built-in tool names.
///
/// Keep aliases specific — short names like `search` / `time` cause false
/// positives when the model narrates tool use in prose.
static const aliases = <String, String>{
'google_search': 'search_web',
'web_search': 'search_web',
'search': 'search_web',
'bing_search': 'search_web',
'duckduckgo': 'search_web',
'duckduckgo_search': 'search_web',
'open_browser': 'search_web',
'browser_search': 'search_web',
'get_current_time': 'get_time',
'current_time': 'get_time',
'time': 'get_time',
'set_an_alarm': 'set_alarm',
'create_alarm': 'set_alarm',
'launch_app': 'open_app',
'start_app': 'open_app',
'open_application': 'open_app',
'screenshot': 'take_screenshot',
'capture_screen': 'take_screenshot',
'weather': 'get_weather',
'take_a_screenshot': 'take_screenshot',
'check_weather': 'get_weather',
'get_current_weather': 'get_weather',
};

/// Returns normalized tool calls, or null if none found.
Expand Down Expand Up @@ -100,10 +101,6 @@ class ToolCallParser {
),
'',
);
out = out.replaceAll(
RegExp(r'call:\s*[a-zA-Z0-9_]+\s*\{[^{}]*(?:\}|$)', caseSensitive: false),
'',
);
out = out.replaceAll(RegExp(r'<\|?"?\|>'), '');
out = out.replaceAll(RegExp(r'[ \t]+\n'), '\n');
out = out.replaceAll(RegExp(r'\n{3,}'), '\n\n');
Expand Down Expand Up @@ -145,10 +142,13 @@ class ToolCallParser {

/// Handles forms like:
/// `<|tool_call>call:google_search{queries:[<|"|>Missypwns twitch<|"|>]}`
///
/// Requires a tool-call delimiter so explanatory prose such as
/// `call:open_app{package:...}` is not executed as a real tool.
static List<Map<String, dynamic>> _parseMarkupCalls(String text) {
final results = <Map<String, dynamic>>[];
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,
);
Comment on lines 150 to 154

@gemini-code-assist gemini-code-assist Bot Jul 19, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

medium

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.

Expand Down
23 changes: 23 additions & 0 deletions test/utils/tool_call_parser_test.dart
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,5 +41,28 @@ void main() {
test('returns null for ordinary prose', () {
expect(ToolCallParser.parse('Hello, how can I help?'), isNull);
});

test('ignores bare call: markup without tool_call delimiter', () {
const prose =
'I would call:open_app{package:com.android.chrome} if tools worked.';
expect(ToolCallParser.parse(prose), isNull);
});

test('does not alias short words like search or time', () {
expect(
ToolCallParser.normalizeCall({
'name': 'search',
'args': {'query': 'x'},
})['name'],
'search',
);
expect(
ToolCallParser.normalizeCall({
'name': 'time',
'args': <String, dynamic>{},
})['name'],
'time',
);
});
});
}
Loading