Skip to content

fix: preserve tool-call history across thread hydration to prevent model re-attempts (#568) #670

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

Open
reidliu41 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: preserve tool-call history across thread hydration to prevent model re-attempts (#568) #670

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
32 changes: 15 additions & 17 deletions src/agent/dispatcher.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -687,23 +687,6 @@ impl Agent {
.await;
}

// Record result in thread
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
match &tool_result {
Ok(output) => {
turn.record_tool_result(serde_json::json!(output));
}
Err(e) => {
turn.record_tool_error(e.to_string());
}
}
}
}

// Check for auth awaiting — defer the return
// until all results are recorded.
if deferred_auth.is_none()
Expand Down Expand Up @@ -743,6 +726,7 @@ impl Agent {
}

// Sanitize and add tool result to context
let is_tool_error = tool_result.is_err();
let result_content = match tool_result {
Ok(output) => {
let sanitized =
Expand All @@ -756,6 +740,20 @@ impl Agent {
Err(e) => format!("Tool '{}' failed: {}", tc.name, e),
};

// Record sanitized result in thread so messages()
// and persist_tool_calls() use cleaned content.
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
if is_tool_error {
turn.record_tool_error(result_content.clone());
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Mar 7, 2026

Choose a reason for hiding this comment

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

security-medium medium

Improper Output Handling (Unsanitized Tool Errors)

Tool execution error messages are recorded in the thread without being sanitized. While successful tool outputs are sanitized using the SafetyLayer (line 733), error messages are formatted as simple strings (line 740) and pushed directly to the context. If a tool (e.g., an HTTP tool) includes untrusted data from an external service in its error message, an attacker could perform a prompt injection attack when this history is fed back to the LLM.

Remediation: Ensure that all tool outputs, including error messages, are passed through the SafetyLayer sanitization process before being recorded in the thread or added to the LLM context.

References
  1. Sanitization should only be applied to data paths sent to external services, such as an LLM, to prevent issues like prompt injection.

}
turn.record_tool_result(serde_json::json!(result_content));
}
}

context_messages.push(ChatMessage::tool_result(
&tc.id,
&tc.name,
Expand Down
Loading