b3nw/zerogravity
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: multi-round tool history rewrite and finishReason handling

Browse Source 
- Add ToolRound struct to pair function calls with results per-round
- Replace single-match history rewrite (broke after first round) with
  multi-round loop that rewrites ALL placeholder model turns
- Fix tool result name fallback: use positional index instead of always
  picking the first call
- Set is_complete for any finishReason (FUNCTION_CALL, MAX_TOKENS, etc.)
  not just STOP — prevents response_complete flag from never being set
- Legacy fallback: responses.rs path (single-round via last_calls +
  pending_results) still works when tool_rounds is empty
- Add tests: multi-round rewrite, single-round legacy, no-op, and
  FUNCTION_CALL/MAX_TOKENS finishReason handling
This commit is contained in:
Nikketryhard
2026-02-16 19:05:37 -06:00
parent 6bda2ecafa
commit 39381a4dfe
5 changed files with 410 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

93
src/api/completions.rs
View File

@@ -16,7 +16,7 @@ use super::polling::{
use super::types::*;
use super::util::{err_response, now_unix, upstream_err_response};
use super::AppState;
use crate::mitm::store::{CapturedFunctionCall, PendingToolResult};
use crate::mitm::store::{CapturedFunctionCall, PendingToolResult, ToolRound};
/// Extract a conversation/session ID from a flexible JSON value.
/// Accepts a plain string or an object with an "id" field.
@@ -229,18 +229,25 @@ pub(crate) async fn handle_completions(
// When OpenCode sends back tool results, the messages array contains:
// 1. assistant message with tool_calls (the model's previous function calls)
// 2. tool messages with results (the executed tool outputs)
// We need to store these so modify_request can rewrite the LS's
// conversation history with proper functionCall/functionResponse parts
// instead of the placeholder "Tool call completed" text.
// We build ToolRounds: each round pairs one assistant's tool_calls with
// the subsequent tool result messages. This enables correct per-turn
// history rewriting for multi-step tool use.
{
let mut last_calls: Vec<CapturedFunctionCall> = Vec::new();
let mut pending_results: Vec<PendingToolResult> = Vec::new();
let mut rounds: Vec<ToolRound> = Vec::new();
let mut current_round: Option<ToolRound> = None;
for msg in &body.messages {
match msg.role.as_str() {
"assistant" => {
// Extract function calls from assistant's tool_calls
// Finalize any open round
if let Some(round) = current_round.take() {
if !round.calls.is_empty() {
rounds.push(round);
}
}
// Start new round if this assistant has tool_calls
if let Some(ref tool_calls) = msg.tool_calls {
let mut calls = Vec::new();
for tc in tool_calls {
if let Some(func) = tc.get("function") {
let name = func["name"].as_str().unwrap_or("unknown").to_string();
@@ -254,7 +261,7 @@ pub(crate) async fn handle_completions(
state.mitm_store.register_call_id(call_id, name.clone()).await;
}
last_calls.push(CapturedFunctionCall {
calls.push(CapturedFunctionCall {
name,
args,
captured_at: std::time::SystemTime::now()
@@ -264,21 +271,31 @@ pub(crate) async fn handle_completions(
});
}
}
if !calls.is_empty() {
current_round = Some(ToolRound {
calls,
results: Vec::new(),
});
}
}
}
"tool" => {
// Extract tool results
let text = extract_message_text(&msg.content);
if let Some(ref call_id) = msg.tool_call_id {
// Look up function name from call_id
// Look up function name from call_id, fall back to
// positional index within the current round's calls
let result_index = current_round
.as_ref()
.map(|r| r.results.len())
.unwrap_or(0);
let name = state
.mitm_store
.lookup_call_id(call_id)
.await
.unwrap_or_else(|| {
// Fallback: try to find the name from last_calls by position
last_calls
.first()
current_round
.as_ref()
.and_then(|r| r.calls.get(result_index))
.map(|fc| fc.name.clone())
.unwrap_or_else(|| "unknown_function".to_string())
});
@@ -286,35 +303,43 @@ pub(crate) async fn handle_completions(
let result_value = serde_json::from_str::<serde_json::Value>(&text)
.unwrap_or_else(|_| serde_json::json!({"result": text}));
pending_results.push(PendingToolResult {
name,
result: result_value,
});
if let Some(ref mut round) = current_round {
round.results.push(PendingToolResult {
name,
result: result_value,
});
}
}
}
_ => {}
_ => {
// Any other role (user, system) finalizes the current round
if let Some(round) = current_round.take() {
if !round.calls.is_empty() {
rounds.push(round);
}
}
}
}
}
// Finalize last round
if let Some(round) = current_round.take() {
if !round.calls.is_empty() {
rounds.push(round);
}
}
if !last_calls.is_empty() {
if !rounds.is_empty() {
info!(
count = last_calls.len(),
tools = ?last_calls.iter().map(|c| &c.name).collect::<Vec<_>>(),
"Completions: stored last function calls for MITM history rewrite"
round_count = rounds.len(),
calls = ?rounds.iter().map(|r| r.calls.iter().map(|c| &c.name).collect::<Vec<_>>()).collect::<Vec<_>>(),
"Completions: stored {} tool round(s) for MITM history rewrite",
rounds.len(),
);
state.mitm_store.set_last_function_calls(last_calls).await;
}
if !pending_results.is_empty() {
info!(
count = pending_results.len(),
tools = ?pending_results.iter().map(|r| &r.name).collect::<Vec<_>>(),
"Completions: stored tool results for MITM injection"
);
for result in pending_results {
state.mitm_store.add_tool_result(result).await;
// Also set last_function_calls from the latest round for proxy.rs recording compat
if let Some(last_round) = rounds.last() {
state.mitm_store.set_last_function_calls(last_round.calls.clone()).await;
}
// Clear awaiting flag — we have the results now
state.mitm_store.set_tool_rounds(rounds).await;
state.mitm_store.clear_awaiting_tool_result();
}
}