fix: normalize each Gemini response call before Tool dispatch
This commit is contained in:
@@ -66,11 +66,15 @@ pub(crate) struct CandidatePart {
|
||||
|
||||
/// Function call.
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct FunctionCall {
|
||||
/// Function name.
|
||||
pub name: String,
|
||||
/// Arguments.
|
||||
pub args: Option<serde_json::Value>,
|
||||
/// Partial-argument protocol is not a complete args object.
|
||||
pub partial_args: Option<serde_json::Value>,
|
||||
pub will_continue: Option<bool>,
|
||||
}
|
||||
|
||||
/// Usage metadata.
|
||||
@@ -110,6 +114,14 @@ pub(crate) struct SafetyRating {
|
||||
// Event conversion
|
||||
// ============================================================================
|
||||
|
||||
/// Response-local Gemini block allocation. Parts in separate SSE frames may
|
||||
/// restart their array index, and function names are not call identities.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct GeminiStreamState {
|
||||
next_index: usize,
|
||||
last_text_index: Option<usize>,
|
||||
}
|
||||
|
||||
impl GeminiScheme {
|
||||
/// Convert SSE data to Event.
|
||||
///
|
||||
@@ -120,7 +132,11 @@ impl GeminiScheme {
|
||||
/// * `Ok(Some(Vec<Event>))` - Conversion succeeded.
|
||||
/// * `Ok(None)` - Event ignored.
|
||||
/// * `Err(ClientError)` - Parse error.
|
||||
pub(crate) fn parse_event(&self, data: &str) -> Result<Option<Vec<Event>>, ClientError> {
|
||||
pub(crate) fn parse_event(
|
||||
&self,
|
||||
data: &str,
|
||||
state: &mut GeminiStreamState,
|
||||
) -> Result<Option<Vec<Event>>, ClientError> {
|
||||
// Skip empty or invalid data.
|
||||
if data.is_empty() || data == "[DONE]" {
|
||||
return Ok(None);
|
||||
@@ -144,31 +160,39 @@ impl GeminiScheme {
|
||||
// Process candidates.
|
||||
if let Some(candidates) = response.candidates {
|
||||
for candidate in candidates {
|
||||
let candidate_index = candidate.index.unwrap_or(0);
|
||||
|
||||
if let Some(content) = candidate.content {
|
||||
if let Some(parts) = content.parts {
|
||||
for (part_index, part) in parts.iter().enumerate() {
|
||||
for part in &parts {
|
||||
let index = state.next_index;
|
||||
state.next_index += 1;
|
||||
// Text delta.
|
||||
if let Some(text) = &part.text {
|
||||
if !text.is_empty() {
|
||||
// Gemini sends no explicit BlockStart, so emit TextDelta
|
||||
// directly (Timeline handles the implicit start).
|
||||
events.push(Event::text_delta(part_index, text.clone()));
|
||||
state.last_text_index = Some(index);
|
||||
events.push(Event::text_delta(index, text.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
// Function call.
|
||||
if let Some(function_call) = &part.function_call {
|
||||
// Function call start.
|
||||
// Gemini usually sends function calls all at once, but may
|
||||
// send partial calls when streaming arguments are enabled.
|
||||
|
||||
// Gemini has no function call IDs, so use the name as the ID.
|
||||
let function_id = format!("call_{}", function_call.name);
|
||||
if function_call.partial_args.is_some()
|
||||
|| function_call.will_continue.is_some()
|
||||
{
|
||||
return Err(ClientError::Api {
|
||||
status: None,
|
||||
code: Some("unsupported_partial_function_call".into()),
|
||||
message: "Gemini partial function arguments require an unsupported streaming protocol".into(),
|
||||
retry_after: None,
|
||||
});
|
||||
}
|
||||
// These synthetic IDs correlate calls within a response;
|
||||
// durable identity still comes from Host history entries.
|
||||
let function_id = format!("call_{index}");
|
||||
|
||||
events.push(Event::BlockStart(BlockStart {
|
||||
index: candidate_index * 10 + part_index, // Composite index.
|
||||
index,
|
||||
block_type: BlockType::ToolUse,
|
||||
metadata: BlockMetadata::ToolUse {
|
||||
id: function_id,
|
||||
@@ -180,12 +204,13 @@ impl GeminiScheme {
|
||||
if let Some(args) = &function_call.args {
|
||||
let args_str = serde_json::to_string(args).unwrap_or_default();
|
||||
if !args_str.is_empty() && args_str != "null" {
|
||||
events.push(Event::tool_input_delta(
|
||||
candidate_index * 10 + part_index,
|
||||
args_str,
|
||||
));
|
||||
events.push(Event::tool_input_delta(index, args_str));
|
||||
}
|
||||
}
|
||||
// Each args object is complete in this wire part. Close
|
||||
// its collector before the next function call starts;
|
||||
// dispatch still waits for the response-complete gate.
|
||||
events.push(Event::tool_use_stop(index));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -200,13 +225,14 @@ impl GeminiScheme {
|
||||
_ => None,
|
||||
};
|
||||
|
||||
// Text block stop.
|
||||
events.push(Event::BlockStop(BlockStop {
|
||||
index: candidate_index,
|
||||
block_type: BlockType::Text,
|
||||
stop_reason,
|
||||
reasoning: None,
|
||||
}));
|
||||
if let Some(index) = state.last_text_index.take() {
|
||||
events.push(Event::BlockStop(BlockStop {
|
||||
index,
|
||||
block_type: BlockType::Text,
|
||||
stop_reason,
|
||||
reasoning: None,
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -240,7 +266,10 @@ mod tests {
|
||||
let data =
|
||||
r#"{"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"},"index":0}]}"#;
|
||||
|
||||
let events = scheme.parse_event(data).unwrap().unwrap();
|
||||
let events = scheme
|
||||
.parse_event(data, &mut GeminiStreamState::default())
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(events.len(), 1);
|
||||
|
||||
if let Event::BlockDelta(delta) = &events[0] {
|
||||
@@ -260,7 +289,10 @@ mod tests {
|
||||
let scheme = GeminiScheme::new();
|
||||
let data = r#"{"candidates":[{"content":{"parts":[{"text":"Hi"}],"role":"model"},"index":0}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}"#;
|
||||
|
||||
let events = scheme.parse_event(data).unwrap().unwrap();
|
||||
let events = scheme
|
||||
.parse_event(data, &mut GeminiStreamState::default())
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
// Should contain a Usage event.
|
||||
let usage_event = events.iter().find(|e| matches!(e, Event::Usage(_)));
|
||||
@@ -278,7 +310,10 @@ mod tests {
|
||||
let scheme = GeminiScheme::new();
|
||||
let data = r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"get_weather","args":{"location":"Tokyo"}}}],"role":"model"},"index":0}]}"#;
|
||||
|
||||
let events = scheme.parse_event(data).unwrap().unwrap();
|
||||
let events = scheme
|
||||
.parse_event(data, &mut GeminiStreamState::default())
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
// Should contain a BlockStart event.
|
||||
let start_event = events.iter().find(|e| matches!(e, Event::BlockStart(_)));
|
||||
@@ -309,7 +344,10 @@ mod tests {
|
||||
let scheme = GeminiScheme::new();
|
||||
let data = r#"{"candidates":[{"content":{"parts":[{"text":"Done"}],"role":"model"},"finishReason":"STOP","index":0}]}"#;
|
||||
|
||||
let events = scheme.parse_event(data).unwrap().unwrap();
|
||||
let events = scheme
|
||||
.parse_event(data, &mut GeminiStreamState::default())
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
// Should contain a BlockStop event.
|
||||
let stop_event = events.iter().find(|e| matches!(e, Event::BlockStop(_)));
|
||||
@@ -320,10 +358,85 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_name_calls_have_distinct_ids_and_block_indexes_across_parts_and_frames() {
|
||||
use crate::llm_client::scheme::Scheme;
|
||||
let scheme = GeminiScheme::new();
|
||||
for frames in [
|
||||
vec![
|
||||
r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"echo","args":{"text":"first"}}},{"functionCall":{"name":"echo","args":{"text":"second"}}}]}}]}"#,
|
||||
],
|
||||
vec![
|
||||
r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"echo","args":{"text":"first"}}}]}}]}"#,
|
||||
r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"echo","args":{"text":"second"}}}]}}]}"#,
|
||||
],
|
||||
] {
|
||||
let mut state = Default::default();
|
||||
let mut events = Vec::new();
|
||||
for frame in frames {
|
||||
events.extend(scheme.parse_sse("", frame, &mut state).unwrap());
|
||||
}
|
||||
let calls: Vec<_> = events
|
||||
.iter()
|
||||
.filter_map(|event| match event {
|
||||
Event::BlockStart(BlockStart {
|
||||
index,
|
||||
metadata: BlockMetadata::ToolUse { id, name },
|
||||
..
|
||||
}) => Some((*index, id.as_str(), name.as_str())),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(calls, [(0, "call_0", "echo"), (1, "call_1", "echo")]);
|
||||
let inputs: Vec<_> = events
|
||||
.iter()
|
||||
.filter_map(|event| match event {
|
||||
Event::BlockDelta(delta) => match &delta.delta {
|
||||
DeltaContent::InputJson(json) => Some((delta.index, json.as_str())),
|
||||
_ => None,
|
||||
},
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(
|
||||
inputs,
|
||||
[(0, r#"{"text":"first"}"#), (1, r#"{"text":"second"}"#)]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_argument_protocol_is_not_misread_as_an_empty_complete_call() {
|
||||
use crate::llm_client::scheme::Scheme;
|
||||
let scheme = GeminiScheme::new();
|
||||
for call in [
|
||||
r#"{"name":"echo","partialArgs":[{"jsonPath":"$.text","stringValue":"part"}]}"#,
|
||||
r#"{"name":"echo","willContinue":true,"args":{"text":"part"}}"#,
|
||||
r#"{"name":"echo","willContinue":false}"#,
|
||||
] {
|
||||
let data = format!(
|
||||
r#"{{"candidates":[{{"content":{{"parts":[{{"functionCall":{call}}}]}}}}]}}"#
|
||||
);
|
||||
assert!(
|
||||
matches!(scheme.parse_sse("", &data, &mut Default::default()), Err(ClientError::Api { code: Some(code), .. }) if code == "unsupported_partial_function_call")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_empty_data() {
|
||||
let scheme = GeminiScheme::new();
|
||||
assert!(scheme.parse_event("").unwrap().is_none());
|
||||
assert!(scheme.parse_event("[DONE]").unwrap().is_none());
|
||||
assert!(
|
||||
scheme
|
||||
.parse_event("", &mut GeminiStreamState::default())
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
scheme
|
||||
.parse_event("[DONE]", &mut GeminiStreamState::default())
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -239,6 +239,8 @@ impl GeminiScheme {
|
||||
let mut contents = Vec::new();
|
||||
let mut pending_model_parts: Vec<GeminiPart> = Vec::new();
|
||||
let mut pending_user_parts: Vec<GeminiPart> = Vec::new();
|
||||
let mut call_names: std::collections::HashMap<&str, std::collections::VecDeque<&str>> =
|
||||
std::collections::HashMap::new();
|
||||
|
||||
for item in items {
|
||||
match item {
|
||||
@@ -270,8 +272,12 @@ impl GeminiScheme {
|
||||
}
|
||||
|
||||
Item::ToolCall {
|
||||
name, arguments, ..
|
||||
call_id,
|
||||
name,
|
||||
arguments,
|
||||
..
|
||||
} => {
|
||||
call_names.entry(call_id).or_default().push_back(name);
|
||||
// Flush pending user parts first
|
||||
if !pending_user_parts.is_empty() {
|
||||
contents.push(GeminiContent {
|
||||
@@ -309,11 +315,17 @@ impl GeminiScheme {
|
||||
Some(c) => format!("{summary}\n{c}"),
|
||||
None => summary.clone(),
|
||||
};
|
||||
// Gemini responses identify the function by its declaration
|
||||
// name, never by the adapter's response-local correlation ID.
|
||||
let name = call_names
|
||||
.get_mut(call_id.as_str())
|
||||
.and_then(std::collections::VecDeque::pop_front)
|
||||
.unwrap_or(call_id.as_str());
|
||||
pending_user_parts.push(GeminiPart::FunctionResponse {
|
||||
function_response: GeminiFunctionResponse {
|
||||
name: call_id.clone(),
|
||||
name: name.to_string(),
|
||||
response: GeminiFunctionResponseContent {
|
||||
name: call_id.clone(),
|
||||
name: name.to_string(),
|
||||
content: Value::String(text),
|
||||
},
|
||||
},
|
||||
@@ -468,6 +480,36 @@ mod tests {
|
||||
assert_eq!(gemini_req.contents[0].role, "user");
|
||||
assert_eq!(gemini_req.contents[1].role, "model");
|
||||
assert_eq!(gemini_req.contents[2].role, "user");
|
||||
let body = serde_json::to_value(gemini_req).unwrap();
|
||||
assert_eq!(
|
||||
body["contents"][2]["parts"][0]["functionResponse"]["name"],
|
||||
"get_weather"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parallel_results_resolve_function_names_not_synthetic_call_ids() {
|
||||
let scheme = GeminiScheme::new();
|
||||
let request = Request::new()
|
||||
.item(Item::tool_call("call_0", "echo", r#"{"text":"first"}"#))
|
||||
.item(Item::tool_call("call_1", "echo", r#"{"text":"second"}"#))
|
||||
.item(Item::tool_result("call_1", "second"))
|
||||
.item(Item::tool_result("call_0", "first"))
|
||||
.item(Item::tool_call("call_0", "other", "{}"))
|
||||
.item(Item::tool_result("call_0", "later"));
|
||||
let body = serde_json::to_value(scheme.build_request(&request, &cap())).unwrap();
|
||||
let responses: Vec<_> = body["contents"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.flat_map(|content| content["parts"].as_array().unwrap())
|
||||
.filter_map(|part| part.get("functionResponse"))
|
||||
.collect();
|
||||
assert_eq!(responses.len(), 3);
|
||||
for (response, name) in responses.iter().zip(["echo", "echo", "other"]) {
|
||||
assert_eq!(response["name"], name);
|
||||
assert_eq!(response["response"]["name"], name);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -10,7 +10,7 @@ use crate::llm_client::{
|
||||
use super::GeminiScheme;
|
||||
|
||||
impl Scheme for GeminiScheme {
|
||||
type State = ();
|
||||
type State = super::events::GeminiStreamState;
|
||||
|
||||
fn default_base_url(&self) -> &'static str {
|
||||
"https://generativelanguage.googleapis.com"
|
||||
@@ -38,9 +38,9 @@ impl Scheme for GeminiScheme {
|
||||
&self,
|
||||
_event_type: &str,
|
||||
data: &str,
|
||||
_state: &mut Self::State,
|
||||
state: &mut Self::State,
|
||||
) -> Result<Vec<Event>, ClientError> {
|
||||
Ok(self.parse_event(data)?.unwrap_or_default())
|
||||
Ok(self.parse_event(data, state)?.unwrap_or_default())
|
||||
}
|
||||
|
||||
fn default_capability(&self) -> ModelCapability {
|
||||
|
||||
@@ -553,6 +553,620 @@ pub(super) mod plugin_runtime_tests {
|
||||
f.shutdown().await;
|
||||
}
|
||||
|
||||
// Same-name calls in one Gemini response must remain distinct through
|
||||
// adapter normalization, committed Host history and real Component dispatch.
|
||||
#[derive(Clone)]
|
||||
struct GeminiTwoCallModel {
|
||||
tool_name: String,
|
||||
https: bool,
|
||||
separate_frames: bool,
|
||||
requests: Arc<Mutex<Vec<ModelRequest>>>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl LlmClient for GeminiTwoCallModel {
|
||||
fn clone_boxed(&self) -> Box<dyn LlmClient> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
|
||||
fn tool_call_completion_support(&self) -> agen::llm_client::ToolCallCompletionSupport {
|
||||
use agen::llm_client::scheme::{Scheme, gemini::GeminiScheme};
|
||||
// Match HttpTransport: Gemini does not have authoritative per-call
|
||||
// done events. Do not silently test the Engine's PerBlock path.
|
||||
let support = GeminiScheme::new().tool_call_completion_support();
|
||||
assert_eq!(
|
||||
support,
|
||||
agen::llm_client::ToolCallCompletionSupport::ResponseComplete
|
||||
);
|
||||
support
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
request: ModelRequest,
|
||||
) -> std::result::Result<agen::llm_client::client::ResponseStream, ClientError> {
|
||||
use agen::llm_client::scheme::{Scheme, gemini::GeminiScheme};
|
||||
let turn = {
|
||||
let mut requests = self.requests.lock().unwrap();
|
||||
let turn = requests.len();
|
||||
requests.push(request);
|
||||
turn
|
||||
};
|
||||
let scheme = GeminiScheme::new();
|
||||
let mut state = <GeminiScheme as Scheme>::State::default();
|
||||
let frames = match turn {
|
||||
// Exercise both two parts in one frame and part-index restart
|
||||
// across frames. Neither path fabricates Engine call IDs.
|
||||
0 => {
|
||||
let parts = ["1", "2"]
|
||||
.into_iter()
|
||||
.map(|text| {
|
||||
json!({"functionCall":{"name":self.tool_name,"args":
|
||||
if self.https { json!({}) } else { json!({"text":text}) }}})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let chunks = if self.separate_frames {
|
||||
parts.into_iter().map(|part| vec![part]).collect::<Vec<_>>()
|
||||
} else {
|
||||
vec![parts]
|
||||
};
|
||||
chunks.into_iter().map(|parts| {
|
||||
json!({"candidates":[{"index":0,"content":{"role":"model","parts":parts}}]})
|
||||
}).chain([json!({"candidates":[{"index":0,"finishReason":"STOP"}]})]).collect::<Vec<_>>()
|
||||
}
|
||||
1 => vec![json!({"candidates":[{"index":0,"content":{"role":"model",
|
||||
"parts":[{"text":"Observed both tool results."}]},"finishReason":"STOP"}]})],
|
||||
_ => panic!("restore must not start an extra model turn"),
|
||||
};
|
||||
let mut events = Vec::new();
|
||||
for frame in frames {
|
||||
events.extend(scheme.parse_sse("", &frame.to_string(), &mut state)?);
|
||||
}
|
||||
// Transport owns this terminal event; the scheme parses wire frames.
|
||||
events.push(ModelEvent::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}));
|
||||
Ok(Box::pin(futures::stream::iter(events.into_iter().map(Ok))))
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_gemini_two_call_history(items: &[agen::Item], public_name: &str, https: bool) {
|
||||
let calls = items
|
||||
.iter()
|
||||
.filter_map(|item| match item {
|
||||
agen::Item::ToolCall {
|
||||
call_id,
|
||||
name,
|
||||
arguments,
|
||||
..
|
||||
} => Some((call_id, name, arguments)),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(calls.len(), 2, "same-name calls must not alias");
|
||||
assert_ne!(
|
||||
calls[0].0, calls[1].0,
|
||||
"provider correlation IDs must be distinct"
|
||||
);
|
||||
let results = items
|
||||
.iter()
|
||||
.filter_map(|item| match item {
|
||||
agen::Item::ToolResult {
|
||||
call_id,
|
||||
summary,
|
||||
content,
|
||||
disposition,
|
||||
is_error,
|
||||
..
|
||||
} => Some((call_id, summary, content, disposition, is_error)),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(results.len(), 2);
|
||||
for (index, (id, name, arguments)) in calls.iter().enumerate() {
|
||||
assert_eq!(name.as_str(), public_name);
|
||||
let text = (index + 1).to_string();
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(arguments).unwrap(),
|
||||
if https {
|
||||
json!({})
|
||||
} else {
|
||||
json!({"text":text})
|
||||
}
|
||||
);
|
||||
let matching = results
|
||||
.iter()
|
||||
.filter(|result| result.0 == *id)
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
matching.len(),
|
||||
1,
|
||||
"each call has exactly one correlated result"
|
||||
);
|
||||
let (_, summary, content, disposition, is_error) = *matching[0];
|
||||
assert_eq!(
|
||||
*disposition,
|
||||
agen::ToolResultDisposition::Success,
|
||||
"{summary}: {content:?}"
|
||||
);
|
||||
assert!(!*is_error);
|
||||
let output: Value = serde_json::from_str(content.as_deref().unwrap()).unwrap();
|
||||
if https {
|
||||
assert_eq!(summary, "HTTP response observed");
|
||||
assert_eq!(output["structured"], json!({"status":503,"body_bytes":9}));
|
||||
} else {
|
||||
assert_eq!(
|
||||
output["structured"],
|
||||
json!({"tool":"example_echo","text":text})
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn gemini_committed_calls(
|
||||
saved: &session_store::RestoredState,
|
||||
public_name: &str,
|
||||
) -> Vec<session_store::LoggedHistoryEntry> {
|
||||
let calls = saved.annotated_history.iter().filter(|entry| matches!(
|
||||
&entry.item, session_store::LoggedItem::ToolCall { name, .. } if name == public_name
|
||||
)).cloned().collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
calls.len(),
|
||||
2,
|
||||
"both calls must have committed Host entries"
|
||||
);
|
||||
assert_ne!(calls[0].metadata.entry_id, calls[1].metadata.entry_id);
|
||||
for call in &calls {
|
||||
assert!(matches!(
|
||||
&call.metadata.origin,
|
||||
session_store::LoggedSessionHistoryOrigin::ModelOutput { .. }
|
||||
));
|
||||
}
|
||||
calls
|
||||
}
|
||||
|
||||
fn gemini_invocation_rows(f: &Fixture) -> Vec<(String, String)> {
|
||||
f.api
|
||||
.config_store
|
||||
.with_conn(|conn| {
|
||||
let mut statement = conn.prepare(
|
||||
"SELECT invocation_id,consumer_json FROM plugin_invocations \
|
||||
WHERE workspace_id=?1 AND plan_id=?2 AND worker_id=?3 \
|
||||
AND installation_id=?4 ORDER BY invocation_id",
|
||||
)?;
|
||||
let rows = statement.query_map(
|
||||
rusqlite::params![
|
||||
TEST_WORKSPACE_ID,
|
||||
f.plan.plan_id,
|
||||
f.source.worker_id,
|
||||
f.installed.installation.installation_id,
|
||||
],
|
||||
|row| Ok((row.get(0)?, row.get(1)?)),
|
||||
)?;
|
||||
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn gemini_receipt_rows(f: &Fixture) -> Vec<(String, String, String, i64, Vec<u8>)> {
|
||||
f.api
|
||||
.config_store
|
||||
.with_conn(|conn| {
|
||||
let mut statement = conn.prepare(
|
||||
"SELECT r.invocation_id,r.effect_key,r.state,r.response_status,r.response_body \
|
||||
FROM plugin_invocations i JOIN plugin_effect_receipts r \
|
||||
USING(workspace_id,invocation_id) WHERE i.workspace_id=?1 \
|
||||
AND i.plan_id=?2 AND i.worker_id=?3 AND i.installation_id=?4 \
|
||||
ORDER BY r.invocation_id,r.effect_key",
|
||||
)?;
|
||||
let rows = statement.query_map(
|
||||
rusqlite::params![
|
||||
TEST_WORKSPACE_ID,
|
||||
f.plan.plan_id,
|
||||
f.source.worker_id,
|
||||
f.installed.installation.installation_id,
|
||||
],
|
||||
|row| {
|
||||
Ok((
|
||||
row.get(0)?,
|
||||
row.get(1)?,
|
||||
row.get(2)?,
|
||||
row.get(3)?,
|
||||
row.get(4)?,
|
||||
))
|
||||
},
|
||||
)?;
|
||||
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn gemini_two_calls_through_runtime_and_restore(https: bool, restore_pending: bool) {
|
||||
let mut peer = if https {
|
||||
Some(
|
||||
Peer::new(
|
||||
HOST,
|
||||
vec![
|
||||
response(503, &[], b"peer data"),
|
||||
response(503, &[], b"peer data"),
|
||||
],
|
||||
)
|
||||
.await,
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
// Pure Echo accepts distinct text inputs; HTTPSComponent only accepts
|
||||
// {}, but must still perform TWO independent effects/receipts.
|
||||
let mut f = Fixture::new(peer.as_ref().map(peer_binding)).await;
|
||||
f.api.plugin_https_transport = peer
|
||||
.as_ref()
|
||||
.map(|peer| Arc::new(peer.transport_for_requests(2)));
|
||||
allow_artifact_read(&f);
|
||||
let (artifacts, delivery_server) = artifact_source(&f, false).await;
|
||||
assert!(delivery_server.is_none());
|
||||
let worker_ref = WorkerRef::try_from(&f.source).unwrap();
|
||||
let providers = Arc::new(RuntimeProviders {
|
||||
worker: worker_ref.clone(),
|
||||
artifacts,
|
||||
owner: f
|
||||
.owner(peer.as_ref().map(|peer| peer.transport_for_requests(2)))
|
||||
.await,
|
||||
});
|
||||
let bridge = PinnedPluginToolRuntime::new(
|
||||
providers,
|
||||
&f._directory.path().join("gemini-runtime-cache"),
|
||||
)
|
||||
.unwrap();
|
||||
let public_name = format!("{}__{}", f.installed.installation.installation_id, f.tool);
|
||||
let model = GeminiTwoCallModel {
|
||||
tool_name: public_name.clone(),
|
||||
https,
|
||||
separate_frames: restore_pending,
|
||||
requests: Arc::default(),
|
||||
};
|
||||
let mut manifest = worker::WorkerManifest::from_toml(&format!(
|
||||
"[worker]\nname=\"{}\"\n[model]\nscheme=\"gemini\"\nmodel_id=\"scripted\"\n[engine]\n[scope]\nallow=[]\n",
|
||||
f.source.worker_id,
|
||||
)).unwrap();
|
||||
// Needed by the real metadata restore path, not a new Profile grant.
|
||||
manifest
|
||||
.plugins
|
||||
.selections
|
||||
.push(manifest::PluginProfileSelection {
|
||||
installation_id: f.installed.installation.installation_id.clone(),
|
||||
tools: vec![f.tool.into()],
|
||||
logging: false,
|
||||
https_slots: if https { vec!["api".into()] } else { vec![] },
|
||||
config: Default::default(),
|
||||
});
|
||||
let store_root = f._directory.path().join("gemini-tool-history");
|
||||
let metadata_root = f._directory.path().join("gemini-worker-metadata");
|
||||
let mut host = worker::Worker::new(
|
||||
manifest.clone(),
|
||||
agen::Engine::new_annotated(Box::new(model.clone()) as Box<dyn LlmClient>)
|
||||
.system_prompt("Use the selected tool twice."),
|
||||
session_store::CombinedStore::new(
|
||||
session_store::FsStore::new(&store_root).unwrap(),
|
||||
session_store::FsWorkerStore::new(&metadata_root).unwrap(),
|
||||
),
|
||||
worker::WorkerWorkspaceContext::unavailable(
|
||||
Some(worker::WorkspaceId::new(TEST_WORKSPACE_ID).unwrap()),
|
||||
"unrelated Workspace operations unavailable",
|
||||
),
|
||||
worker::WorkerFilesystemAuthority::None,
|
||||
manifest::Scope::empty(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let session_id = host.session_id();
|
||||
let snapshot = runtime_api::PluginExecutionSnapshot {
|
||||
plan_id: f.plan.plan_id.clone(),
|
||||
fingerprint: f.plan.fingerprint.clone(),
|
||||
plan: f.plan.plan.clone(),
|
||||
};
|
||||
let feature = bridge
|
||||
.prepare(&worker_ref, &snapshot, &session_id.to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
let report = host.install_features(FeatureRegistryBuilder::new().with_module(feature));
|
||||
assert!(!report.has_errors(), "{}", report.error_message());
|
||||
host.enable_worker_metadata_write_through().unwrap();
|
||||
if restore_pending {
|
||||
host.add_pre_tool_call_hook(PauseBeforeFirstCall);
|
||||
}
|
||||
host.wire_history_persistence();
|
||||
let result = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(30),
|
||||
host.run(vec![protocol::Segment::Text {
|
||||
content: "Call the selected tool twice in one response.".into(),
|
||||
}]),
|
||||
)
|
||||
.await
|
||||
.expect("two-call runtime deadline")
|
||||
.unwrap();
|
||||
let mut pending_calls = None;
|
||||
if restore_pending {
|
||||
assert!(matches!(result, worker::WorkerRunResult::Paused));
|
||||
assert_eq!(model.requests.lock().unwrap().len(), 1);
|
||||
assert_eq!(f.count("plugin_invocations"), 0);
|
||||
assert_eq!(f.count("plugin_effect_receipts"), 0);
|
||||
if let Some(peer) = &peer {
|
||||
assert_eq!(peer.connection_count(), 0);
|
||||
}
|
||||
let reopened = session_store::FsStore::new(&store_root).unwrap();
|
||||
let segments = reopened.list_segments(session_id).unwrap();
|
||||
assert_eq!(segments.len(), 1);
|
||||
let pending = session_store::restore(&reopened, session_id, segments[0]).unwrap();
|
||||
let calls = gemini_committed_calls(&pending, &public_name);
|
||||
let consumers = persisted_call_consumers(session_id, &pending, &public_name);
|
||||
assert_eq!(consumers.len(), 2);
|
||||
assert_ne!(consumers[0], consumers[1]);
|
||||
assert!(
|
||||
!pending
|
||||
.history
|
||||
.iter()
|
||||
.any(|item| matches!(item, agen::Item::ToolResult { .. }))
|
||||
);
|
||||
// Inspect live Host annotations as well as the reopened Store.
|
||||
let live_ids = host
|
||||
.session_history()
|
||||
.iter()
|
||||
.filter_map(|entry| match &entry.item {
|
||||
agen::Item::ToolCall { .. } => Some(entry.annotation.entry_id.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
live_ids,
|
||||
calls
|
||||
.iter()
|
||||
.map(|entry| entry.metadata.entry_id.clone())
|
||||
.collect::<Vec<_>>()
|
||||
);
|
||||
drop(host);
|
||||
let mut restored = worker::Worker::prepare_restored_from_worker_metadata_with_context_and_model_client(
|
||||
&manifest.worker.name, manifest.clone(), session_store::CombinedStore::new(
|
||||
reopened, session_store::FsWorkerStore::new(&metadata_root).unwrap(),
|
||||
), worker::PromptCatalogSource::builtins_only(),
|
||||
worker::WorkerWorkspaceContext::unavailable(Some(worker::WorkspaceId::new(TEST_WORKSPACE_ID).unwrap()), "unrelated Workspace operations unavailable"),
|
||||
worker::WorkerFilesystemAuthority::None, Some(Box::new(model.clone())),
|
||||
).await.unwrap().activate().await.unwrap();
|
||||
assert_eq!(restored.session_id(), session_id);
|
||||
let reloaded = session_store::restore(
|
||||
&session_store::FsStore::new(&store_root).unwrap(),
|
||||
session_id,
|
||||
segments[0],
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(gemini_committed_calls(&reloaded, &public_name), calls);
|
||||
assert_eq!(
|
||||
persisted_call_consumers(session_id, &reloaded, &public_name),
|
||||
consumers
|
||||
);
|
||||
let restored_ids = restored
|
||||
.session_history()
|
||||
.iter()
|
||||
.filter_map(|entry| match &entry.item {
|
||||
agen::Item::ToolCall { .. } => Some(entry.annotation.entry_id.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
restored_ids, live_ids,
|
||||
"metadata reopen retains both live Host entries"
|
||||
);
|
||||
let feature = bridge
|
||||
.prepare(&worker_ref, &snapshot, &session_id.to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
let report =
|
||||
restored.install_features(FeatureRegistryBuilder::new().with_module(feature));
|
||||
assert!(!report.has_errors(), "{}", report.error_message());
|
||||
restored.wire_history_persistence();
|
||||
let result =
|
||||
tokio::time::timeout(std::time::Duration::from_secs(30), restored.resume())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(matches!(result, worker::WorkerRunResult::Finished));
|
||||
pending_calls = Some((calls, consumers));
|
||||
host = restored;
|
||||
} else {
|
||||
assert!(matches!(result, worker::WorkerRunResult::Finished));
|
||||
}
|
||||
assert_gemini_two_call_history(&host.history(), &public_name, https);
|
||||
{
|
||||
use agen::llm_client::scheme::{Scheme, gemini::GeminiScheme};
|
||||
let requests = model.requests.lock().unwrap();
|
||||
assert_eq!(requests.len(), 2);
|
||||
assert_eq!(
|
||||
requests[0]
|
||||
.tools
|
||||
.iter()
|
||||
.map(|tool| tool.name.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
[public_name.as_str()]
|
||||
);
|
||||
assert_gemini_two_call_history(&requests[1].items, &public_name, https);
|
||||
let scheme = GeminiScheme::new();
|
||||
let body =
|
||||
scheme.build_request_body("scripted", &requests[1], &scheme.default_capability());
|
||||
let responses = body["contents"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.flat_map(|content| content["parts"].as_array().unwrap())
|
||||
.filter_map(|part| part.get("functionResponse"))
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(responses.len(), 2);
|
||||
for response in responses {
|
||||
assert_eq!(
|
||||
response["name"], public_name,
|
||||
"functionResponse uses the real Tool name, not a generated call ID"
|
||||
);
|
||||
assert_eq!(response["response"]["name"], public_name);
|
||||
}
|
||||
}
|
||||
let reopened = session_store::FsStore::new(&store_root).unwrap();
|
||||
let segments = reopened.list_segments(session_id).unwrap();
|
||||
assert_eq!(segments.len(), 1);
|
||||
let saved = session_store::restore(&reopened, session_id, segments[0]).unwrap();
|
||||
assert_gemini_two_call_history(&saved.history, &public_name, https);
|
||||
let calls = gemini_committed_calls(&saved, &public_name);
|
||||
let mut consumers = persisted_call_consumers(session_id, &saved, &public_name);
|
||||
assert_eq!(consumers.len(), 2);
|
||||
assert_ne!(consumers[0], consumers[1]);
|
||||
if let Some((pending, expected_consumers)) = pending_calls {
|
||||
assert_eq!(
|
||||
calls, pending,
|
||||
"resume preserves both committed Host entries"
|
||||
);
|
||||
assert_eq!(consumers, expected_consumers);
|
||||
}
|
||||
let invocations = gemini_invocation_rows(&f);
|
||||
assert_eq!(invocations.len(), 2);
|
||||
assert_ne!(invocations[0].0, invocations[1].0);
|
||||
let mut actual_consumers = invocations
|
||||
.iter()
|
||||
.map(
|
||||
|(_, consumer)| match serde_json::from_str::<PluginConsumer>(consumer).unwrap() {
|
||||
PluginConsumer::ToolCall {
|
||||
tool_call_id,
|
||||
tool_name,
|
||||
} => {
|
||||
assert_eq!(tool_name, f.tool);
|
||||
tool_call_id
|
||||
}
|
||||
other => panic!("unexpected consumer: {other:?}"),
|
||||
},
|
||||
)
|
||||
.collect::<Vec<_>>();
|
||||
consumers.sort();
|
||||
actual_consumers.sort();
|
||||
assert_eq!(
|
||||
actual_consumers, consumers,
|
||||
"SQLite admission uses both durable Host consumers"
|
||||
);
|
||||
let receipts = gemini_receipt_rows(&f);
|
||||
assert_eq!(f.count("plugin_invocations"), 2);
|
||||
assert_eq!(receipts.len(), if https { 2 } else { 0 });
|
||||
assert_eq!(f.count("plugin_effect_receipts"), if https { 2 } else { 0 });
|
||||
if let Some(peer) = &mut peer {
|
||||
assert_eq!(
|
||||
peer.connection_count(),
|
||||
2,
|
||||
"identical HTTPS inputs are separate effects"
|
||||
);
|
||||
for _ in 0..2 {
|
||||
let observed =
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), peer.observed())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
String::from_utf8(observed)
|
||||
.unwrap()
|
||||
.starts_with("GET /status HTTP/1.1\r\n")
|
||||
);
|
||||
}
|
||||
for (index, (invocation, key, state, status, body)) in receipts.iter().enumerate() {
|
||||
assert_eq!(invocation, &invocations[index].0);
|
||||
assert_eq!(key, "status-v1");
|
||||
assert_eq!(state, "completed");
|
||||
assert_eq!(*status, 503);
|
||||
assert_eq!(body, b"peer data");
|
||||
}
|
||||
}
|
||||
// A completed metadata reopen/activation is not a new run. Do not call
|
||||
// resume() (which requires a paused logical run), or fabricate history.
|
||||
drop(host);
|
||||
let mut completed =
|
||||
worker::Worker::prepare_restored_from_worker_metadata_with_context_and_model_client(
|
||||
&manifest.worker.name,
|
||||
manifest.clone(),
|
||||
session_store::CombinedStore::new(
|
||||
reopened,
|
||||
session_store::FsWorkerStore::new(&metadata_root).unwrap(),
|
||||
),
|
||||
worker::PromptCatalogSource::builtins_only(),
|
||||
worker::WorkerWorkspaceContext::unavailable(
|
||||
Some(worker::WorkspaceId::new(TEST_WORKSPACE_ID).unwrap()),
|
||||
"unrelated Workspace operations unavailable",
|
||||
),
|
||||
worker::WorkerFilesystemAuthority::None,
|
||||
Some(Box::new(model.clone())),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.activate()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(completed.session_id(), session_id);
|
||||
let feature = bridge
|
||||
.prepare(&worker_ref, &snapshot, &session_id.to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
let report = completed.install_features(FeatureRegistryBuilder::new().with_module(feature));
|
||||
assert!(!report.has_errors(), "{}", report.error_message());
|
||||
completed.wire_history_persistence();
|
||||
assert!(completed.engine().active_run_turn_count().is_none());
|
||||
assert_gemini_two_call_history(&completed.history(), &public_name, https);
|
||||
let completed_ids = completed
|
||||
.session_history()
|
||||
.iter()
|
||||
.filter_map(|entry| match &entry.item {
|
||||
agen::Item::ToolCall { .. } => Some(entry.annotation.entry_id.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
completed_ids,
|
||||
calls
|
||||
.iter()
|
||||
.map(|entry| entry.metadata.entry_id.clone())
|
||||
.collect::<Vec<_>>(),
|
||||
"completed Worker activation retains both committed Host identities"
|
||||
);
|
||||
let after = session_store::restore(
|
||||
&session_store::FsStore::new(&store_root).unwrap(),
|
||||
session_id,
|
||||
segments[0],
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(after.annotated_history, saved.annotated_history);
|
||||
assert_eq!(gemini_committed_calls(&after, &public_name), calls);
|
||||
assert_eq!(gemini_invocation_rows(&f), invocations);
|
||||
assert_eq!(gemini_receipt_rows(&f), receipts);
|
||||
assert_eq!(model.requests.lock().unwrap().len(), 2);
|
||||
assert_eq!(f.count("plugin_invocations"), 2);
|
||||
assert_eq!(f.count("plugin_effect_receipts"), if https { 2 } else { 0 });
|
||||
if let Some(peer) = &peer {
|
||||
assert_eq!(peer.connection_count(), 2);
|
||||
}
|
||||
drop(completed);
|
||||
f.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gemini_same_response_pure_calls_keep_distinct_inputs_and_completed_restore_is_inert() {
|
||||
gemini_two_calls_through_runtime_and_restore(false, false).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gemini_same_response_https_calls_commit_two_effects_and_completed_restore_is_inert() {
|
||||
gemini_two_calls_through_runtime_and_restore(true, false).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gemini_same_response_pending_pure_calls_restore_both_host_consumers() {
|
||||
gemini_two_calls_through_runtime_and_restore(false, true).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gemini_same_response_pending_https_calls_restore_both_host_consumers_and_effects() {
|
||||
gemini_two_calls_through_runtime_and_restore(true, true).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_pinned_https_tool_commits_normal_history_and_receipt_then_fences_live_denial() {
|
||||
https_through_runtime_and_history(true, false).await;
|
||||
|
||||
Reference in New Issue
Block a user