diff --git a/crates/agen/src/llm_client/scheme/gemini/events.rs b/crates/agen/src/llm_client/scheme/gemini/events.rs index c6c4ea0ea..27bc1feae 100644 --- a/crates/agen/src/llm_client/scheme/gemini/events.rs +++ b/crates/agen/src/llm_client/scheme/gemini/events.rs @@ -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, + /// Partial-argument protocol is not a complete args object. + pub partial_args: Option, + pub will_continue: Option, } /// 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, +} + impl GeminiScheme { /// Convert SSE data to Event. /// @@ -120,7 +132,11 @@ impl GeminiScheme { /// * `Ok(Some(Vec))` - Conversion succeeded. /// * `Ok(None)` - Event ignored. /// * `Err(ClientError)` - Parse error. - pub(crate) fn parse_event(&self, data: &str) -> Result>, ClientError> { + pub(crate) fn parse_event( + &self, + data: &str, + state: &mut GeminiStreamState, + ) -> Result>, 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() + ); } } diff --git a/crates/agen/src/llm_client/scheme/gemini/request.rs b/crates/agen/src/llm_client/scheme/gemini/request.rs index e1e6b4040..f3aebd522 100644 --- a/crates/agen/src/llm_client/scheme/gemini/request.rs +++ b/crates/agen/src/llm_client/scheme/gemini/request.rs @@ -239,6 +239,8 @@ impl GeminiScheme { let mut contents = Vec::new(); let mut pending_model_parts: Vec = Vec::new(); let mut pending_user_parts: Vec = 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] diff --git a/crates/agen/src/llm_client/scheme/gemini/scheme_impl.rs b/crates/agen/src/llm_client/scheme/gemini/scheme_impl.rs index 5d5724ce5..0b3b368d5 100644 --- a/crates/agen/src/llm_client/scheme/gemini/scheme_impl.rs +++ b/crates/agen/src/llm_client/scheme/gemini/scheme_impl.rs @@ -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, ClientError> { - Ok(self.parse_event(data)?.unwrap_or_default()) + Ok(self.parse_event(data, state)?.unwrap_or_default()) } fn default_capability(&self) -> ModelCapability { diff --git a/crates/workspace-server/src/server_plugin_runtime_tests.rs b/crates/workspace-server/src/server_plugin_runtime_tests.rs index 5d5bc117b..127df8be4 100644 --- a/crates/workspace-server/src/server_plugin_runtime_tests.rs +++ b/crates/workspace-server/src/server_plugin_runtime_tests.rs @@ -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>>, + } + + #[async_trait::async_trait] + impl LlmClient for GeminiTwoCallModel { + fn clone_boxed(&self) -> Box { + 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 { + 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 = ::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::>(); + let chunks = if self.separate_frames { + parts.into_iter().map(|part| vec![part]).collect::>() + } 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::>() + } + 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::>(); + 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::>(); + 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::(arguments).unwrap(), + if https { + json!({}) + } else { + json!({"text":text}) + } + ); + let matching = results + .iter() + .filter(|result| result.0 == *id) + .collect::>(); + 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 { + let calls = saved.annotated_history.iter().filter(|entry| matches!( + &entry.item, session_store::LoggedItem::ToolCall { name, .. } if name == public_name + )).cloned().collect::>(); + 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::>>()?) + }) + .unwrap() + } + + fn gemini_receipt_rows(f: &Fixture) -> Vec<(String, String, String, i64, Vec)> { + 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::>>()?) + }) + .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) + .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::>(); + assert_eq!( + live_ids, + calls + .iter() + .map(|entry| entry.metadata.entry_id.clone()) + .collect::>() + ); + 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::>(); + 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::>(), + [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::>(); + 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::(consumer).unwrap() { + PluginConsumer::ToolCall { + tool_call_id, + tool_name, + } => { + assert_eq!(tool_name, f.tool); + tool_call_id + } + other => panic!("unexpected consumer: {other:?}"), + }, + ) + .collect::>(); + 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::>(); + assert_eq!( + completed_ids, + calls + .iter() + .map(|entry| entry.metadata.entry_id.clone()) + .collect::>(), + "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;