diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 3cac1cd6..1cb7fdcf 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -74,10 +74,10 @@ jobs: run: cargo fmt --check - name: Run clippy - run: cargo clippy --all-targets -- -D warnings + run: cargo clippy --workspace --all-targets --all-features --locked -- -D warnings - name: Run tests - run: cargo test + run: cargo test --workspace --all-features --locked - name: Test Codex and Claude launcher contracts run: bash scripts/tests/agentic-launchers-test.sh @@ -85,6 +85,18 @@ jobs: - name: Test agentic CLI end to end run: python3 scripts/tests/agentic-cli-e2e-test.py + - name: Install Python for strict SDK contracts + uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.13' + + - name: Verify Files, Vector Stores, and file-search SDK contracts + run: | + python3 -m venv "${RUNNER_TEMP}/file-search-sdk" + "${RUNNER_TEMP}/file-search-sdk/bin/pip" install 'openai==3.13.0' + cargo build -p agentic-server --all-features --locked + "${RUNNER_TEMP}/file-search-sdk/bin/python" scripts/tests/file-search-sdk-test.py target/debug/agentic-server -v + postgres: runs-on: ubuntu-latest services: @@ -129,20 +141,13 @@ jobs: - name: Provision pgvector extension run: psql "$TEST_POSTGRES_URL" -v ON_ERROR_STOP=1 -c "CREATE EXTENSION IF NOT EXISTS vector" - - name: Verify PostgreSQL schema upgrades and validation - run: cargo test -p agentic-server-core --lib -- --ignored --test-threads=1 - - - name: Verify PostgreSQL migrations and restart persistence - run: cargo test -p agentic-server-core --test postgres_storage_integration -- --ignored --test-threads=1 - - - name: Verify PostgreSQL file search persistence - run: cargo test -p agentic-server-core --test file_search_service postgres_file_search -- --ignored --test-threads=1 - - - name: Verify real pgvector indexed retrieval and transactional publication - run: cargo test -p agentic-server-core --test file_search_service postgres_pgvector -- --ignored --test-threads=1 + # Every ignored core test currently requires this isolated PostgreSQL service. + # Keep new lifecycle/runtime targets in this gate; never add paid/network tests here. + - name: Verify all PostgreSQL storage, retrieval, lifecycle, and worker contracts + run: cargo test -p agentic-server-core --all-features --locked --lib --tests -- --ignored --test-threads=1 - name: Verify bounded PDF ingestion - run: cargo test -p agentic-server-core --features file-search-pdf --test file_search_service pdf + run: cargo test -p agentic-server-core --features file-search-pdf --locked --test file_search_service pdf provider-upstream: runs-on: ubuntu-latest diff --git a/Cargo.lock b/Cargo.lock index 7d48b054..9ac38ea8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10,9 +10,9 @@ checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" [[package]] name = "aes" -version = "0.9.3" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35f0f96ce78e38c3dc6d8948aa8163d06385be74000f3c7a95bf1eef35d3ea32" +checksum = "f8eb277bec05f56a0e0591f155a484cbd0f4f07ff2905051a48c72f004f7ed58" dependencies = [ "cipher", "cpubits", diff --git a/crates/agentic-server-core/src/events/normalize.rs b/crates/agentic-server-core/src/events/normalize.rs index c128f130..2443cbd1 100644 --- a/crates/agentic-server-core/src/events/normalize.rs +++ b/crates/agentic-server-core/src/events/normalize.rs @@ -102,6 +102,7 @@ fn extract_payload(event_type: SSEEventType, json: &Value) -> EventPayload { | SSEEventType::McpListToolsInProgress | SSEEventType::McpListToolsCompleted | SSEEventType::McpListToolsFailed + | SSEEventType::OutputTextAnnotationAdded | SSEEventType::Other => EventPayload::Raw(json.clone()), } } diff --git a/crates/agentic-server-core/src/events/types.rs b/crates/agentic-server-core/src/events/types.rs index c1ec924b..77ea3b77 100644 --- a/crates/agentic-server-core/src/events/types.rs +++ b/crates/agentic-server-core/src/events/types.rs @@ -125,6 +125,7 @@ pub enum SSEEventType { // Text content OutputTextDelta, OutputTextDone, + OutputTextAnnotationAdded, ContentPartAdded, ContentPartDone, @@ -177,6 +178,7 @@ impl From<&str> for SSEEventType { "response.output_item.done" => Self::OutputItemDone, "response.output_text.delta" => Self::OutputTextDelta, "response.output_text.done" => Self::OutputTextDone, + "response.output_text.annotation.added" => Self::OutputTextAnnotationAdded, "response.content_part.added" => Self::ContentPartAdded, "response.content_part.done" => Self::ContentPartDone, "response.function_call_arguments.delta" => Self::FunctionCallArgumentsDelta, @@ -225,6 +227,7 @@ impl TryFrom for &'static str { SSEEventType::OutputItemDone => Ok("response.output_item.done"), SSEEventType::OutputTextDelta => Ok("response.output_text.delta"), SSEEventType::OutputTextDone => Ok("response.output_text.done"), + SSEEventType::OutputTextAnnotationAdded => Ok("response.output_text.annotation.added"), SSEEventType::ContentPartAdded => Ok("response.content_part.added"), SSEEventType::ContentPartDone => Ok("response.content_part.done"), SSEEventType::FunctionCallArgumentsDelta => Ok("response.function_call_arguments.delta"), diff --git a/crates/agentic-server-core/src/events/validate.rs b/crates/agentic-server-core/src/events/validate.rs index a20e8ea6..66b2c99c 100644 --- a/crates/agentic-server-core/src/events/validate.rs +++ b/crates/agentic-server-core/src/events/validate.rs @@ -46,7 +46,7 @@ pub(crate) fn validate_frame(frame: &EventFrame) -> Result, E SSEEventType::OutputItemDone => { validate_output_item(frame, event_name, true).map(|item| ValidatedFrame { item: Some(item) }) } - SSEEventType::Other => Ok(ValidatedFrame { item: None }), + SSEEventType::OutputTextAnnotationAdded | SSEEventType::Other => Ok(ValidatedFrame { item: None }), event_type => { let output_index = required_output_index(frame, event_name)?; let item_id = validate_event_item_id(frame, event_name)?; @@ -110,6 +110,7 @@ pub(crate) fn expected_item_type(event_type: SSEEventType) -> Option None, } } @@ -276,6 +277,7 @@ fn validate_event_fields( | SSEEventType::McpListToolsInProgress | SSEEventType::McpListToolsCompleted | SSEEventType::McpListToolsFailed + | SSEEventType::OutputTextAnnotationAdded | SSEEventType::Other => None, }; if let Some(field) = required { diff --git a/crates/agentic-server-core/src/executor/accumulator/mod.rs b/crates/agentic-server-core/src/executor/accumulator/mod.rs index 74a405e1..dcb891de 100644 --- a/crates/agentic-server-core/src/executor/accumulator/mod.rs +++ b/crates/agentic-server-core/src/executor/accumulator/mod.rs @@ -633,8 +633,9 @@ impl ResponseAccumulator { previous_response_id: previous_response_id.map(str::to_string), conversation_id: self.conversation_id, instructions: instructions.map(str::to_string), - tools: None, - tool_choice: None, + tools: Vec::new(), + tool_choice: crate::types::io::ToolChoice::default(), + parallel_tool_calls: false, } } } diff --git a/crates/agentic-server-core/src/executor/engine.rs b/crates/agentic-server-core/src/executor/engine.rs index f1bf0895..8e6b7829 100644 --- a/crates/agentic-server-core/src/executor/engine.rs +++ b/crates/agentic-server-core/src/executor/engine.rs @@ -469,9 +469,13 @@ async fn run_compaction_trigger( previous_response_id: ctx.original_request.previous_response_id.clone(), conversation_id: ctx.conversation_id.clone(), instructions, - tools: None, - tool_choice: None, + tools: ctx.enriched_request.tools.clone().unwrap_or_default(), + tool_choice: ctx.enriched_request.tool_choice.clone().unwrap_or_default(), + parallel_tool_calls: ctx.enriched_request.parallel_tool_calls.unwrap_or(false), }; + for tool in &mut payload.tools { + tool.sanitize_for_persistence(); + } ctx.inject_ids(&mut payload); Ok(payload) } @@ -606,6 +610,9 @@ struct StreamFailureContext { model: String, previous_response_id: Option, instructions: Option, + tools: Vec, + tool_choice: ToolChoice, + parallel_tool_calls: bool, } impl From<&RequestContext> for StreamFailureContext { @@ -616,6 +623,19 @@ impl From<&RequestContext> for StreamFailureContext { model: ctx.enriched_request.model.clone(), previous_response_id: ctx.original_request.previous_response_id.clone(), instructions: ctx.original_request.instructions.clone(), + tools: ctx + .enriched_request + .tools + .clone() + .unwrap_or_default() + .into_iter() + .map(|mut tool| { + tool.sanitize_for_persistence(); + tool + }) + .collect(), + tool_choice: ctx.enriched_request.tool_choice.clone().unwrap_or_default(), + parallel_tool_calls: ctx.enriched_request.parallel_tool_calls.unwrap_or(false), } } } @@ -639,8 +659,9 @@ impl StreamFailureContext { previous_response_id: self.previous_response_id.clone(), conversation_id: self.conversation_id.clone(), instructions: self.instructions.clone(), - tools: None, - tool_choice: None, + tools: self.tools.clone(), + tool_choice: self.tool_choice.clone(), + parallel_tool_calls: self.parallel_tool_calls, } } } diff --git a/crates/agentic-server-core/src/executor/gateway_accumulator.rs b/crates/agentic-server-core/src/executor/gateway_accumulator.rs index 533113b2..47b90856 100644 --- a/crates/agentic-server-core/src/executor/gateway_accumulator.rs +++ b/crates/agentic-server-core/src/executor/gateway_accumulator.rs @@ -9,6 +9,7 @@ pub struct GatewayStreamAccumulator { next_sequence_number: u64, emitted_created: bool, emitted_in_progress: bool, + citations: super::stream_citations::StreamCitations, } pub(super) struct StreamEvent { @@ -28,6 +29,7 @@ impl GatewayStreamAccumulator { next_sequence_number: 0, emitted_created: false, emitted_in_progress: false, + citations: super::stream_citations::StreamCitations::default(), } } @@ -50,10 +52,20 @@ impl GatewayStreamAccumulator { rebase_output_index(&mut frame.wire, output_offset); } + pub(super) fn citation_frames(&mut self, frame: &EventFrame, offset: usize) -> ExecutorResult> { + self.citations.before_done(frame, offset) + } + pub(crate) fn terminal_response_chunk(&mut self, payload: &ResponsePayload) -> ExecutorResult { let mut frame = terminal_response_frame(payload)?; + let mut chunk = String::new(); + for mut annotation in self.citation_frames(&frame, 0)? { + self.stamp_event(&mut annotation, 0); + chunk.push_str(&checked_stream_event(&annotation)?); + } self.stamp_event(&mut frame, 0); - checked_stream_event(&frame) + chunk.push_str(&checked_stream_event(&frame)?); + Ok(chunk) } #[cfg(test)] diff --git a/crates/agentic-server-core/src/executor/mod.rs b/crates/agentic-server-core/src/executor/mod.rs index 7dcde006..3d98d89e 100644 --- a/crates/agentic-server-core/src/executor/mod.rs +++ b/crates/agentic-server-core/src/executor/mod.rs @@ -21,6 +21,7 @@ pub mod gateway_accumulator; mod pending_calls; mod pipeline; mod response_budget; +mod stream_citations; mod upstream; pub use compaction::compact_response; diff --git a/crates/agentic-server-core/src/executor/pipeline.rs b/crates/agentic-server-core/src/executor/pipeline.rs index 2e0f0904..aa7536e9 100644 --- a/crates/agentic-server-core/src/executor/pipeline.rs +++ b/crates/agentic-server-core/src/executor/pipeline.rs @@ -25,6 +25,7 @@ pub(super) struct StreamPayload { /// Lives for the response, preserving gateway event numbering across inference rounds. pub(super) struct AgentPipeline { pub(super) request: RequestContext, + response_tool_choice: crate::types::io::ToolChoice, tool_search_state: Option, delivery: StreamDelivery, round: Option, @@ -37,6 +38,7 @@ impl AgentPipeline { sender: Option>, ) -> Self { Self { + response_tool_choice: request.enriched_request.tool_choice.clone().unwrap_or_default(), request, tool_search_state, delivery: StreamDelivery::new(sender), @@ -44,6 +46,21 @@ impl AgentPipeline { } } + pub(super) fn response_tool_choice(&self) -> &crate::types::io::ToolChoice { + &self.response_tool_choice + } + + pub(super) fn response_tools(&self) -> Vec { + let mut tools = self.tool_search_state().filter(|state| state.is_active()).map_or_else( + || self.request.enriched_request.tools.clone().unwrap_or_default(), + ToolSearchState::public_response_tools, + ); + for tool in &mut tools { + tool.sanitize_for_persistence(); + } + tools + } + pub(super) fn tool_search_state(&self) -> Option<&ToolSearchState> { self.tool_search_state.as_ref() } @@ -114,6 +131,9 @@ impl AgentPipeline { self.request.original_request.previous_response_id.as_deref(), self.request.original_request.instructions.as_deref(), )?; + payload.tools = self.response_tools(); + payload.tool_choice = self.response_tool_choice.clone(); + payload.parallel_tool_calls = self.request.enriched_request.parallel_tool_calls.unwrap_or(false); self.request.inject_ids(&mut payload); Ok(payload) } diff --git a/crates/agentic-server-core/src/executor/pipeline/delivery.rs b/crates/agentic-server-core/src/executor/pipeline/delivery.rs index 623ae512..0eb2a972 100644 --- a/crates/agentic-server-core/src/executor/pipeline/delivery.rs +++ b/crates/agentic-server-core/src/executor/pipeline/delivery.rs @@ -148,7 +148,27 @@ fn should_defer_stream_event(frame: &EventFrame, defer_from_output_index: Option } async fn emit_stream_frame(frame: &mut EventFrame, emit_ctx: &mut StreamEmitContext<'_>) -> ExecutorResult { + // Completed grounded text supplies authoritative offsets and annotation order. + if frame.event_type == SSEEventType::OutputTextAnnotationAdded + && frame + .wire + .rest + .get("annotation") + .and_then(|annotation| annotation.get("type")) + .and_then(Value::as_str) + == Some("file_citation") + { + return Ok(false); + } apply_context_response_ids(&mut frame.wire, emit_ctx.request); + for mut annotation in emit_ctx.accumulator.citation_frames(frame, emit_ctx.output_offset)? { + if emit_ctx + .accumulator + .process_event(&mut annotation, emit_ctx.output_offset) + { + emit_sse_frame(emit_ctx.sender, &annotation).await?; + } + } let emitted = emit_ctx.accumulator.process_event(frame, emit_ctx.output_offset); if emitted { emit_sse_frame(emit_ctx.sender, frame).await?; diff --git a/crates/agentic-server-core/src/executor/stream_citations.rs b/crates/agentic-server-core/src/executor/stream_citations.rs new file mode 100644 index 00000000..992e3d66 --- /dev/null +++ b/crates/agentic-server-core/src/executor/stream_citations.rs @@ -0,0 +1,146 @@ +//! File citation events are reconciled against grounded completed content. +//! +//! Upstream file annotation events may precede the text needed to validate their +//! offsets. Their completed content is authoritative; emit its restored citations +//! before the corresponding done event, with request-local deduplication. Other +//! annotation kinds, including null annotations, continue to pass through. +use std::collections::HashSet; + +use crate::events::{EventFrame, SSEEventType, WireEvent}; +use crate::types::event::OutputTextFileCitationAdded; +use crate::types::io::FileCitation; +use crate::utils::common::serialize_to_value; + +use super::error::{ExecutorError, ExecutorResult}; + +#[derive(Clone, Default)] +pub(super) struct StreamCitations { + // Global output index avoids collisions when providers reuse IDs between rounds. + emitted: HashSet<(u64, usize, usize)>, +} + +impl StreamCitations { + pub(super) fn before_done(&mut self, frame: &EventFrame, offset: usize) -> ExecutorResult> { + let mut events = Vec::new(); + match frame.event_type { + SSEEventType::ContentPartDone => { + if let (Some(item_id), Some(index), Some(content_index), Some(part)) = ( + frame.wire.rest.get("item_id").and_then(serde_json::Value::as_str), + frame.wire.output_index, + frame.wire.rest.get("content_index").and_then(serde_json::Value::as_u64), + frame.wire.rest.get("part"), + ) { + self.content( + &mut events, + item_id, + index, + usize::try_from(content_index).map_err(|_| { + ExecutorError::StreamError("annotation content index exceeds platform bounds".to_owned()) + })?, + part, + offset, + )?; + } + } + SSEEventType::OutputItemDone => { + if let (Some(index), Some(item)) = (frame.wire.output_index, frame.wire.rest.get("item")) { + self.item(&mut events, index, item, offset)?; + } + } + SSEEventType::ResponseCompleted | SSEEventType::ResponseIncomplete | SSEEventType::ResponseFailed => { + if let Some(output) = frame + .wire + .rest + .get("response") + .and_then(|response| response.get("output")) + .and_then(serde_json::Value::as_array) + { + for (index, item) in output.iter().enumerate() { + self.item( + &mut events, + u64::try_from(index).map_err(|_| { + ExecutorError::StreamError("annotation output index exceeds platform bounds".to_owned()) + })?, + item, + offset, + )?; + } + } + } + _ => {} + } + Ok(events) + } + + fn item( + &mut self, + events: &mut Vec, + index: u64, + item: &serde_json::Value, + offset: usize, + ) -> ExecutorResult<()> { + if item.get("type").and_then(serde_json::Value::as_str) != Some("message") { + return Ok(()); + } + if let (Some(id), Some(content)) = ( + item.get("id").and_then(serde_json::Value::as_str), + item.get("content").and_then(serde_json::Value::as_array), + ) { + for (content_index, part) in content.iter().enumerate() { + self.content(events, id, index, content_index, part, offset)?; + } + } + Ok(()) + } + + fn content( + &mut self, + events: &mut Vec, + item_id: &str, + index: u64, + content_index: usize, + part: &serde_json::Value, + offset: usize, + ) -> ExecutorResult<()> { + if part.get("type").and_then(serde_json::Value::as_str) != Some("output_text") { + return Ok(()); + } + let Some(annotations) = part.get("annotations").and_then(serde_json::Value::as_array) else { + return Ok(()); + }; + let global_index = index.saturating_add(u64::try_from(offset).unwrap_or(u64::MAX)); + for (annotation_index, annotation) in annotations.iter().enumerate() { + if annotation.get("type").and_then(serde_json::Value::as_str) != Some("file_citation") { + continue; + } + let citation: FileCitation = + serde_json::from_value(annotation.clone()).map_err(ExecutorError::JsonError)?; + if !self.emitted.insert((global_index, content_index, annotation_index)) { + continue; + } + // Each entry needs at least this much upstream JSON. The shared 1 MiB + // response budget bounds normal state; retain a defensive independent cap. + if self.emitted.len() > super::response_budget::MAX_EXECUTOR_RESPONSE_BYTES / 24 { + return Err(ExecutorError::StreamError( + "stream annotation state exceeded the response budget".to_owned(), + )); + } + let event = OutputTextFileCitationAdded { + item_id: item_id.to_owned(), + output_index: index, + content_index, + annotation_index, + sequence_number: 0, + annotation: citation, + }; + let wire: WireEvent = serde_json::from_value(serialize_to_value(&event).map_err(ExecutorError::JsonError)?) + .map_err(ExecutorError::JsonError)?; + events.push(EventFrame { + event_type: SSEEventType::OutputTextAnnotationAdded, + payload: crate::events::EventPayload::None, + wire, + }); + } + Ok(()) + } +} diff --git a/crates/agentic-server-core/src/executor/translate/context.rs b/crates/agentic-server-core/src/executor/translate/context.rs index 34a63704..174ecead 100644 --- a/crates/agentic-server-core/src/executor/translate/context.rs +++ b/crates/agentic-server-core/src/executor/translate/context.rs @@ -22,6 +22,7 @@ pub(in crate::executor) struct TranslationContext { custom_tool_map: Option, response_tools: Option>, response_tool_choice: Option, + parallel_tool_calls: Option, } impl std::fmt::Debug for TranslationContext { @@ -73,6 +74,11 @@ impl TranslationContext { self } + pub(in crate::executor) fn with_parallel_tool_calls(mut self, parallel_tool_calls: bool) -> Self { + self.parallel_tool_calls = Some(parallel_tool_calls); + self + } + pub(in crate::executor) fn with_file_search_context(mut self, input: &crate::types::io::ResponsesInput) -> Self { self.file_search_citations = Some(FileSearchCitations::from_input(input)); self @@ -86,19 +92,29 @@ impl TranslationContext { if self.response_tools.is_some() && let Some(choice) = self.response_tool_choice.as_ref() && let Some(response) = wire.rest.get_mut("response").and_then(serde_json::Value::as_object_mut) - && response.contains_key("tool_choice") { response.insert("tool_choice".to_owned(), serde_json::to_value(choice)?); } + if let Some(parallel_tool_calls) = self.parallel_tool_calls + && let Some(response) = wire.rest.get_mut("response").and_then(serde_json::Value::as_object_mut) + { + response.insert( + "parallel_tool_calls".to_owned(), + serde_json::Value::Bool(parallel_tool_calls), + ); + } super::custom::CustomTranslator::restore_response_wire(wire, self.custom_tool_map.as_ref()); let _ = super::namespace::CodexNamespaceTranslator::restore_response_wire(wire, self.namespace_map.as_ref()); Ok(()) } pub(super) fn restore_response_metadata(&self, payload: &mut ResponsePayload) { + if let Some(parallel_tool_calls) = self.parallel_tool_calls { + payload.parallel_tool_calls = parallel_tool_calls; + } if let Some(tools) = &self.response_tools { - payload.tools = Some(tools.clone()); - payload.tool_choice = Some(self.response_tool_choice.clone().unwrap_or_default()); + payload.tools.clone_from(tools); + payload.tool_choice = self.response_tool_choice.clone().unwrap_or_default(); } } diff --git a/crates/agentic-server-core/src/executor/translate/tests.rs b/crates/agentic-server-core/src/executor/translate/tests.rs index ef996018..699cc1fb 100644 --- a/crates/agentic-server-core/src/executor/translate/tests.rs +++ b/crates/agentic-server-core/src/executor/translate/tests.rs @@ -1381,12 +1381,13 @@ fn public_catalog_distinguishes_inactive_search_from_an_empty_active_catalog() { assert_eq!(tools[0]["name"], "ordinary"); } let mut payload = accumulator.finalize("test", None, None); + payload.tools = serde_json::from_value(serde_json::json!([{"type":"function","name":"ordinary"}])).unwrap(); dispatcher .finish() .unwrap() .normalize_response_payload(&mut payload) .unwrap(); - assert_eq!(payload.tools.is_some(), active); + assert_eq!(payload.tools.is_empty(), active); } } diff --git a/crates/agentic-server-core/src/executor/translate/tool_search.rs b/crates/agentic-server-core/src/executor/translate/tool_search.rs index 57217c82..b6cb8502 100644 --- a/crates/agentic-server-core/src/executor/translate/tool_search.rs +++ b/crates/agentic-server-core/src/executor/translate/tool_search.rs @@ -406,9 +406,6 @@ pub(super) fn restore_response_tools( let Some(response) = wire.rest.get_mut("response").and_then(Value::as_object_mut) else { return Ok(()); }; - if !response.contains_key("tools") { - return Ok(()); - } if let Some(tools) = tools { response.insert( "tools".to_owned(), diff --git a/crates/agentic-server-core/src/executor/upstream.rs b/crates/agentic-server-core/src/executor/upstream.rs index f2be23f5..9fc14ece 100644 --- a/crates/agentic-server-core/src/executor/upstream.rs +++ b/crates/agentic-server-core/src/executor/upstream.rs @@ -37,24 +37,10 @@ fn translation_context(registry: &ToolRegistry, agent: &AgentPipeline) -> Transl .with_response_metadata( registry.namespace_map().cloned(), registry.custom_tool_map().cloned(), - state - .filter(|state| state.is_active()) - .map(crate::tool::ToolSearchState::public_response_tools) - .or_else(|| { - agent - .request - .enriched_request - .tools - .as_ref() - .filter(|tools| { - tools - .iter() - .any(|tool| matches!(tool, crate::types::tools::ResponsesTool::Shell(_))) - }) - .cloned() - }), - agent.request.enriched_request.tool_choice.clone(), + Some(agent.response_tools()), + Some(agent.response_tool_choice().clone()), ) + .with_parallel_tool_calls(agent.request.enriched_request.parallel_tool_calls.unwrap_or(false)) } /// Builds the JSON body sent upstream: history inlined, continuation and storage @@ -227,6 +213,8 @@ pub(super) mod tests { .unwrap(), ); request.enriched_request.parallel_tool_calls = Some(false); + request.enriched_request.tool_choice = + Some(serde_json::from_value(serde_json::json!({"type":"custom","name":"raw_echo"})).unwrap()); let state = ToolSearchHandler::prepare_request(&mut request.enriched_request, &[], false).unwrap(); let registry = ToolRegistry::build_with_handlers( request.enriched_request.tools.as_mut().unwrap(), diff --git a/crates/agentic-server-core/src/storage/file_search.rs b/crates/agentic-server-core/src/storage/file_search.rs index 29ca5ee6..c4ccbf62 100644 --- a/crates/agentic-server-core/src/storage/file_search.rs +++ b/crates/agentic-server-core/src/storage/file_search.rs @@ -58,6 +58,48 @@ pub(crate) struct StoredChunk { pub attributes: crate::types::file_search::FileAttributes, } +/// Private retrieval identity; never serialized into a public search result. +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize)] +pub(crate) struct ChunkOrigin { + pub store_id: String, + pub file_id: String, + pub chunk_index: i64, + pub generation: String, +} + +pub(crate) struct RetrievedChunk { + pub chunk: StoredChunk, + pub origin: ChunkOrigin, +} + +#[derive(FromRow)] +pub(crate) struct ChunkRow { + pub store_id: String, + pub file_id: String, + pub chunk_index: i64, + pub generation: String, + pub data: String, +} + +impl ChunkRow { + fn origin(self) -> ChunkOrigin { + ChunkOrigin { + store_id: self.store_id, + file_id: self.file_id, + chunk_index: self.chunk_index, + generation: self.generation, + } + } + + pub(crate) fn decode(self) -> Result { + let chunk = serde_json::from_str(&self.data)?; + Ok(RetrievedChunk { + chunk, + origin: self.origin(), + }) + } +} + pub(crate) struct PreparedAttachment { pub object: VectorStoreFileObject, pub chunks: Vec, @@ -199,35 +241,46 @@ impl FileSearchStorage { Ok(()) } - pub(crate) async fn visible_result_files( + /// Revalidate exact attachment generations and chunks in one database snapshot. + /// Attributes are refreshed here because they may change during model work. + pub(crate) async fn visible_result_origins( &self, - stores: &[String], - results: &[crate::types::file_search::SearchResult], - ) -> Result, FileSearchError> { - if results.is_empty() { - return Ok(std::collections::HashSet::new()); + origins: &[&ChunkOrigin], + filter: Option<&crate::types::file_search::SearchFilter>, + ) -> Result, FileSearchError> + { + use futures::TryStreamExt; + let mut visible = std::collections::HashMap::new(); + if origins.is_empty() { + return Ok(visible); } - let store_placeholders = (1..=stores.len()) - .map(|index| format!("${index}")) - .collect::>() - .join(", "); - let file_placeholders = (stores.len() + 1..=stores.len() + results.len()) - .map(|index| format!("${index}")) - .collect::>() - .join(", "); + // One JSON parameter avoids backend bind-parameter limits for the bounded + // 10000-origin union. SQL and backend-specific JSON decoding stay in storage. + let requested = if self.pool.acquire().await?.backend_name() == "PostgreSQL" { + "jsonb_to_recordset($1::jsonb) AS requested(store_id text, file_id text, chunk_index bigint, generation text)" + } else { + "(SELECT json_extract(value, '$.store_id') AS store_id, json_extract(value, '$.file_id') AS file_id, json_extract(value, '$.chunk_index') AS chunk_index, json_extract(value, '$.generation') AS generation FROM json_each($1)) AS requested" + }; + let visibility = self.file_visibility().await?; let sql = format!( - "SELECT DISTINCT file_id FROM file_search_attachments WHERE status = 'completed' AND store_id IN ({store_placeholders}) AND file_id IN ({file_placeholders}) AND store_id IN (SELECT id FROM file_search_stores WHERE lifecycle_status != 'expired' AND {}) AND file_id IN (SELECT id FROM file_search_files WHERE {})", - self.file_visibility().await?, - self.file_visibility().await? + "SELECT c.store_id, c.file_id, c.chunk_index, a.generation, a.data FROM {requested} JOIN file_search_chunks c ON c.store_id=requested.store_id AND c.file_id=requested.file_id AND c.chunk_index=requested.chunk_index JOIN file_search_attachments a ON a.store_id=c.store_id AND a.file_id=c.file_id AND a.generation=requested.generation WHERE a.status='completed' AND c.store_id IN (SELECT id FROM file_search_stores WHERE lifecycle_status!='expired' AND {visibility}) AND c.file_id IN (SELECT id FROM file_search_files WHERE {visibility})" ); - let mut query = sqlx::query_scalar::<_, String>(&sql); - for store in stores { - query = query.bind(store); - } - for result in results { - query = query.bind(&result.file_id); + let query = sqlx::query_as::<_, ChunkRow>(&sql).bind(serde_json::to_string(origins)?); + let mut rows = query.fetch(self.pool.as_ref()); + let mut bytes = 0usize; + while let Some(row) = rows.try_next().await? { + bytes = bytes.saturating_add(row.data.len()); + if bytes > 64 * 1024 * 1024 { + return Err(FileSearchError::Unavailable( + "Final search visibility exceeds 64 MiB".into(), + )); + } + let object: VectorStoreFileObject = serde_json::from_str(&row.data)?; + if filter.is_none_or(|filter| filter.matches(&object.attributes)) { + visible.insert(row.origin(), object.attributes); + } } - Ok(query.fetch_all(self.pool.as_ref()).await?.into_iter().collect()) + Ok(visible) } pub(crate) async fn has_chunks(&self, stores: &[String]) -> Result { @@ -253,7 +306,7 @@ impl FileSearchStorage { vectors: &[Vec], mode: crate::types::file_search::SearchMode, filter: Option<&crate::types::file_search::SearchFilter>, - ) -> Result, FileSearchError> { + ) -> Result, FileSearchError> { match &self.pgvector { Some(pgvector) => { pgvector @@ -352,7 +405,16 @@ impl FileSearchStorage { .bind(&object.id).bind(object.created_at).bind(serde_json::to_string(object)?).bind(identity).bind(now).bind(days).bind(days.map(|days| now + days * 86400)) .execute(&mut *tx).await?; for (attachment, (chunks, storage_bytes)) in attachments.iter().zip(encoded) { - publish_attachment(&mut tx, &object.id, identity, attachment, &chunks, storage_bytes).await?; + publish_attachment( + &mut tx, + &object.id, + identity, + attachment, + &chunks, + storage_bytes, + &uuid::Uuid::now_v7().to_string(), + ) + .await?; } tx.commit().await?; Ok(()) @@ -367,7 +429,16 @@ impl FileSearchStorage { self.initialize().await?; let (chunks, storage_bytes) = serialize_chunks(&attachment.chunks).await?; let mut tx = self.pool.begin().await?; - publish_attachment(&mut tx, store_id, identity, attachment, &chunks, storage_bytes).await?; + publish_attachment( + &mut tx, + store_id, + identity, + attachment, + &chunks, + storage_bytes, + &uuid::Uuid::now_v7().to_string(), + ) + .await?; tx.commit().await?; Ok(()) } @@ -516,7 +587,7 @@ impl FileSearchStorage { } /// Streaming row decoding bounds corpus memory before deserializing vectors. - pub(crate) async fn chunks(&self, store_ids: &[String]) -> Result, FileSearchError> { + pub(crate) async fn chunks(&self, store_ids: &[String]) -> Result, FileSearchError> { use futures::TryStreamExt; let placeholders = (1..=store_ids.len()) .map(|index| format!("${index}")) @@ -524,17 +595,17 @@ impl FileSearchStorage { .join(", "); let visibility = self.file_visibility().await?; let sql = format!( - "SELECT data FROM file_search_chunks WHERE (store_id, file_id) IN (SELECT store_id, file_id FROM file_search_attachments WHERE status = 'completed') AND store_id IN ({placeholders}) AND store_id IN (SELECT id FROM file_search_stores WHERE lifecycle_status != 'expired' AND {visibility}) AND file_id IN (SELECT id FROM file_search_files WHERE {visibility}) ORDER BY store_id, file_id, chunk_index" + "SELECT store_id, file_id, chunk_index, (SELECT generation FROM file_search_attachments a WHERE a.store_id=file_search_chunks.store_id AND a.file_id=file_search_chunks.file_id) AS generation, data FROM file_search_chunks WHERE (store_id, file_id) IN (SELECT store_id, file_id FROM file_search_attachments WHERE status = 'completed') AND store_id IN ({placeholders}) AND store_id IN (SELECT id FROM file_search_stores WHERE lifecycle_status != 'expired' AND {visibility}) AND file_id IN (SELECT id FROM file_search_files WHERE {visibility}) ORDER BY store_id, file_id, chunk_index" ); - let mut query = sqlx::query_scalar::<_, String>(&sql); + let mut query = sqlx::query_as::<_, ChunkRow>(&sql); for id in store_ids { query = query.bind(id); } let mut rows = query.fetch(self.pool.as_ref()); let mut chunks = Vec::new(); let mut total_bytes = 0usize; - while let Some(data) = rows.try_next().await? { - total_bytes = total_bytes.saturating_add(data.len()); + while let Some(row) = rows.try_next().await? { + total_bytes = total_bytes.saturating_add(row.data.len()); if i64::try_from(total_bytes).unwrap_or(i64::MAX) > MAX_CORPUS_BYTES || i64::try_from(chunks.len()).unwrap_or(i64::MAX) >= MAX_CORPUS_CHUNKS { @@ -542,7 +613,7 @@ impl FileSearchStorage { "Selected corpus exceeds the exact retrieval capacity; search fewer vector stores or files".into(), )); } - chunks.push(serde_json::from_str(&data)?); + chunks.push(row.decode()?); } Ok(chunks) } @@ -602,6 +673,7 @@ async fn publish_attachment( attachment: &PreparedAttachment, chunks: &[String], storage_bytes: i64, + generation: &str, ) -> Result<(), FileSearchError> { let file_id = &attachment.object.id; let locked = sqlx::query("UPDATE file_search_files SET id = id WHERE id = $1") @@ -650,9 +722,9 @@ async fn publish_attachment( count.saturating_add(i64::try_from(chunks.len()).unwrap_or(i64::MAX)), )?; let object = &attachment.object; - let inserted = sqlx::query("INSERT INTO file_search_attachments (store_id, file_id, created_at, usage_bytes, data, storage_bytes, status, parsed_content) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)") + let inserted = sqlx::query("INSERT INTO file_search_attachments (store_id, file_id, created_at, usage_bytes, data, storage_bytes, status, parsed_content, generation) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)") .bind(store_id).bind(&object.id).bind(object.created_at).bind(object.usage_bytes).bind(serde_json::to_string(object)?) - .bind(storage_bytes).bind(object.status.as_str()).bind(&attachment.parsed_content) + .bind(storage_bytes).bind(object.status.as_str()).bind(&attachment.parsed_content).bind(generation) .execute(&mut **tx).await; if let Err(sqlx::Error::Database(error)) = &inserted { if error.is_unique_violation() { diff --git a/crates/agentic-server-core/src/storage/pgvector.rs b/crates/agentic-server-core/src/storage/pgvector.rs index 2e5868f1..707a5fa4 100644 --- a/crates/agentic-server-core/src/storage/pgvector.rs +++ b/crates/agentic-server-core/src/storage/pgvector.rs @@ -5,7 +5,10 @@ use futures::TryStreamExt; use std::{collections::HashSet, sync::Arc}; use tokio::sync::OnceCell; -use super::{DbPool, file_search::StoredChunk}; +use super::{ + DbPool, + file_search::{ChunkRow, RetrievedChunk}, +}; use crate::types::file_search::{ AttributeValue, ComparisonOperator, CompoundOperator, FileSearchBackend, FileSearchError, FilterValue, PgvectorIndex, SearchFilter, SearchMode, invalid, @@ -155,7 +158,7 @@ impl PgvectorStorage { vectors: &[Vec], mode: SearchMode, filter: Option<&SearchFilter>, - ) -> Result, FileSearchError> { + ) -> Result, FileSearchError> { self.initialize(pool).await?; if mode != SearchMode::Keyword && matches!(self.index, PgvectorIndex::Ivfflat { .. }) { self.maintain_index(pool).await?; @@ -188,7 +191,7 @@ impl PgvectorStorage { continue; } let mut sql = CandidateSql::new( - "SELECT data FROM file_search_chunks WHERE (store_id, file_id) IN (SELECT store_id, file_id FROM file_search_attachments WHERE status = 'completed') AND store_id IN (", + "SELECT store_id, file_id, chunk_index, (SELECT generation FROM file_search_attachments a WHERE a.store_id=file_search_chunks.store_id AND a.file_id=file_search_chunks.file_id) AS generation, data FROM file_search_chunks WHERE (store_id, file_id) IN (SELECT store_id, file_id FROM file_search_attachments WHERE status = 'completed') AND store_id IN (", ); for (i, store) in stores.iter().enumerate() { if i > 0 { @@ -225,19 +228,19 @@ impl PgvectorStorage { .push(")) DESC"); } sql.push(format!(" LIMIT {}", self.limit)); - let mut query = sqlx::query_scalar::<_, String>(&sql.text); + let mut query = sqlx::query_as::<_, ChunkRow>(&sql.text); for value in &sql.parameters { query = query.bind(value); } let mut rows = query.fetch(&mut *tx); - while let Some(data) = rows.try_next().await? { - candidate_bytes = candidate_bytes.saturating_add(data.len()); + while let Some(row) = rows.try_next().await? { + candidate_bytes = candidate_bytes.saturating_add(row.data.len()); candidate_count += 1; if candidate_bytes > 64 * 1024 * 1024 || candidate_count > 10_000 { return Err(FileSearchError::Unavailable("Indexed candidate union exceeds 64 MiB or 10000 rows; reduce candidate_limit or the number of queries".into())); } - let chunk: StoredChunk = serde_json::from_str(&data)?; - if seen.insert((chunk.file_id.clone(), chunk.text.clone())) { + let chunk = row.decode()?; + if seen.insert(chunk.origin.clone()) { chunks.push(chunk); } } diff --git a/crates/agentic-server-core/src/storage/vector_store_batches.rs b/crates/agentic-server-core/src/storage/vector_store_batches.rs index 72543420..ae57d986 100644 --- a/crates/agentic-server-core/src/storage/vector_store_batches.rs +++ b/crates/agentic-server-core/src/storage/vector_store_batches.rs @@ -450,13 +450,16 @@ impl FileSearchStorage { .await?; let (chunks, bytes) = encoded.ok_or_else(|| FileSearchError::Unavailable("Missing prepared chunks".into()))?; - publish_attachment(&mut tx, &job.store_id, &job.identity, &prepared, &chunks, bytes).await?; - sqlx::query("UPDATE file_search_attachments SET generation=$3 WHERE store_id=$1 AND file_id=$2") - .bind(&job.store_id) - .bind(&job.options.file_id) - .bind(&job.generation.0) - .execute(&mut *tx) - .await?; + publish_attachment( + &mut tx, + &job.store_id, + &job.identity, + &prepared, + &chunks, + bytes, + &job.generation.0, + ) + .await?; object = prepared.object; } Err(error) => { diff --git a/crates/agentic-server-core/src/tool/file_search/ingest.rs b/crates/agentic-server-core/src/tool/file_search/ingest.rs index cdffbb57..20b7ca36 100644 --- a/crates/agentic-server-core/src/tool/file_search/ingest.rs +++ b/crates/agentic-server-core/src/tool/file_search/ingest.rs @@ -253,10 +253,10 @@ fn chunks(text: &str, config: &StaticChunking, cancelled: &AtomicBool) -> Result /// Keep complete source chunks that fit the remaining model-context budget. pub(super) fn limit_context( - results: Vec, + results: Vec, mut budget: usize, cancelled: &AtomicBool, -) -> Result, FileSearchError> { +) -> Result, FileSearchError> { if cancelled.load(Ordering::Relaxed) { return Err(FileSearchError::Unavailable( "File search context preparation was cancelled".into(), @@ -266,7 +266,7 @@ pub(super) fn limit_context( let mut selected = Vec::with_capacity(results.len()); 'passages: for result in results { let mut tokens = 0usize; - for content in &result.content { + for content in &result.result.content { let mut text = content.text.as_str(); while !text.is_empty() { if cancelled.load(Ordering::Relaxed) { diff --git a/crates/agentic-server-core/src/tool/file_search/models.rs b/crates/agentic-server-core/src/tool/file_search/models.rs index 0b8cd354..cbd24af0 100644 --- a/crates/agentic-server-core/src/tool/file_search/models.rs +++ b/crates/agentic-server-core/src/tool/file_search/models.rs @@ -1,8 +1,7 @@ //! Cancellation-safe bounded text generation and vLLM reranking. use crate::types::{ file_search::{ - ContextualChunking, FileSearchError, ModelProvider, ScoreInterpretation, SearchResult, VectorStoresConfig, - invalid, + ContextualChunking, FileSearchError, ModelProvider, ScoreInterpretation, VectorStoresConfig, invalid, }, retrieval_models::{ RetrievalChatRequest, RetrievalChatResponse, RetrievalMessage, RetrievalRole, RetrievalText, RetrievalTextPart, @@ -150,9 +149,9 @@ impl Models { pub(super) async fn rerank( &self, query: &str, - mut candidates: Vec, + mut candidates: Vec, selection: Option<&str>, - ) -> Result, FileSearchError> { + ) -> Result, FileSearchError> { let (provider, model) = self .config .resolve(selection, self.config.default_reranker_model.as_ref())?; @@ -161,7 +160,7 @@ impl Models { } let documents = candidates .iter() - .map(|result| result.content[0].text.as_str()) + .map(|result| result.result.content[0].text.as_str()) .collect(); // Score every bounded candidate, then validate the complete permutation before publishing any result. let request = TextRerankRequest { @@ -193,9 +192,13 @@ impl Models { exp / (1.0 + exp) } }; - candidates[result.index].score = score; + let candidate = &mut candidates[result.index]; + candidate.result.score = score; + for origin in &mut candidate.origins { + origin.score = score; + } } - candidates.sort_by(|left, right| right.score.total_cmp(&left.score)); + candidates.sort_by(|left, right| right.result.score.total_cmp(&left.result.score)); Ok(candidates) } async fn post( diff --git a/crates/agentic-server-core/src/tool/file_search/ranking.rs b/crates/agentic-server-core/src/tool/file_search/ranking.rs index a58a365d..4d60a307 100644 --- a/crates/agentic-server-core/src/tool/file_search/ranking.rs +++ b/crates/agentic-server-core/src/tool/file_search/ranking.rs @@ -3,28 +3,44 @@ use std::collections::{HashMap, HashSet}; use crate::{ - storage::file_search::StoredChunk, + storage::file_search::{ChunkOrigin, RetrievedChunk, StoredChunk}, types::file_search::{ChunkRetrievalParams, Ranker, SearchContent, SearchMode, SearchRequest, SearchResult}, }; +/// Keeps all matching origins of a deduplicated passage until final visibility. +pub(super) struct RankedCandidate { + pub result: SearchResult, + pub origins: Vec, +} + +pub(super) struct ScoredOrigin { + pub origin: ChunkOrigin, + pub score: f64, +} + pub(super) fn rank( - chunks: Vec, + chunks: Vec, queries: &[String], query_embeddings: &[Vec], mode: SearchMode, request: &SearchRequest, ranker: Ranker, defaults: &ChunkRetrievalParams, -) -> Vec { +) -> Vec { let chunks: Vec<_> = chunks .into_iter() .filter(|chunk| { request .filters .as_ref() - .is_none_or(|filter| filter.matches(&chunk.attributes)) + .is_none_or(|filter| filter.matches(&chunk.chunk.attributes)) }) .collect(); + let (chunks, origins): (Vec<_>, Vec<_>) = chunks + .into_iter() + .map(|retrieved| (retrieved.chunk, retrieved.origin)) + .unzip(); + let mut origin_scores = vec![None::; chunks.len()]; let mut best: HashMap<(&str, &str), (usize, f64)> = HashMap::new(); let threshold = request .ranking_options @@ -59,6 +75,7 @@ pub(super) fn rank( if !ranker.uses_model() && (score <= 0.0 || score < threshold) { continue; } + origin_scores[index] = Some(origin_scores[index].map_or(score, |previous| previous.max(score))); let key = (chunks[index].file_id.as_str(), chunks[index].text.as_str()); let entry = best.entry(key).or_insert((index, score)); if score > entry.1 { @@ -66,6 +83,15 @@ pub(super) fn rank( } } } + let mut origin_groups: HashMap<(&str, &str), Vec> = HashMap::new(); + for ((chunk, origin), score) in chunks.iter().zip(origins).zip(origin_scores) { + if let Some(score) = score { + origin_groups + .entry((&chunk.file_id, &chunk.text)) + .or_default() + .push(ScoredOrigin { origin, score }); + } + } let mut matches: Vec<_> = best.into_values().collect(); matches.sort_by(|(left, left_score), (right, right_score)| { right_score @@ -78,15 +104,20 @@ pub(super) fn rank( .into_iter() .map(|(index, score)| { let chunk = &chunks[index]; - SearchResult { - file_id: chunk.file_id.clone(), - filename: chunk.filename.clone(), - score, - attributes: chunk.attributes.clone(), - content: vec![SearchContent { - type_: "text".into(), - text: chunk.text.clone(), - }], + RankedCandidate { + origins: origin_groups + .remove(&(chunk.file_id.as_str(), chunk.text.as_str())) + .unwrap_or_default(), + result: SearchResult { + file_id: chunk.file_id.clone(), + filename: chunk.filename.clone(), + score, + attributes: chunk.attributes.clone(), + content: vec![SearchContent { + type_: "text".into(), + text: chunk.text.clone(), + }], + }, } }) .collect() @@ -264,6 +295,49 @@ fn weighted(semantic: &[f64], keyword: &[f64], embedding_weight: f64, text_weigh mod tests { use super::cosine; + #[test] + fn deduplicated_origins_must_independently_meet_the_score_threshold() { + use super::*; + let chunks = [vec![1.0, 0.0], vec![0.5, 0.75_f64.sqrt()]] + .into_iter() + .enumerate() + .map(|(index, embedding)| RetrievedChunk { + origin: ChunkOrigin { + store_id: format!("store-{index}"), + file_id: "file".into(), + chunk_index: 0, + generation: "generation".into(), + }, + chunk: StoredChunk { + file_id: "file".into(), + filename: "source.txt".into(), + chunk_index: 0, + text: "same source".into(), + embedding_text: None, + embedding: Some(embedding), + attributes: std::collections::BTreeMap::default(), + }, + }) + .collect(); + let request: SearchRequest = + serde_json::from_str(r#"{"query":"source","ranking_options":{"score_threshold":0.8}}"#).unwrap(); + let results = rank( + chunks, + &["source".into()], + &[vec![1.0, 0.0]], + SearchMode::Semantic, + &request, + Ranker::None, + &ChunkRetrievalParams::default(), + ); + assert_eq!(results.len(), 1); + assert_eq!( + results[0].origins.len(), + 1, + "a below-threshold attachment cannot validate another origin's result" + ); + } + #[test] fn cosine_is_stable_for_finite_vectors_of_extreme_magnitude() { for scale in [1e-300, 1.0, 1e300] { diff --git a/crates/agentic-server-core/src/tool/file_search/service.rs b/crates/agentic-server-core/src/tool/file_search/service.rs index 8ce260b8..710804b7 100644 --- a/crates/agentic-server-core/src/tool/file_search/service.rs +++ b/crates/agentic-server-core/src/tool/file_search/service.rs @@ -112,9 +112,9 @@ impl FileSearchService { async fn prepare_context( &self, - data: Vec, + data: Vec, permit: Arc, - ) -> Result, FileSearchError> { + ) -> Result, FileSearchError> { let budget = self.config.chunk_retrieval_params.max_tokens_in_context; let cancelled = Arc::new(AtomicBool::new(false)); let _cancel_on_drop = CancelIngestionOnDrop(cancelled.clone()); @@ -662,6 +662,33 @@ impl FileSearchService { self.search_impl(store_ids, request, true).await } + async fn visible_results( + &self, + data: Vec, + filter: Option<&crate::types::file_search::SearchFilter>, + ) -> Result, FileSearchError> { + let origins: Vec<_> = data + .iter() + .flat_map(|candidate| candidate.origins.iter().map(|scored| &scored.origin)) + .collect(); + let visible = self.storage.visible_result_origins(&origins, filter).await?; + let mut data: Vec<_> = data + .into_iter() + .filter_map(|mut candidate| { + let (origin, attributes) = candidate + .origins + .iter() + .filter_map(|scored| visible.get(&scored.origin).map(|attributes| (scored, attributes))) + .max_by(|(left, _), (right, _)| left.score.total_cmp(&right.score))?; + candidate.result.attributes.clone_from(attributes); + candidate.result.score = origin.score; + Some(candidate.result) + }) + .collect(); + data.sort_by(|left, right| right.score.total_cmp(&left.score)); + Ok(data) + } + async fn search_impl( &self, store_ids: &[String], @@ -702,7 +729,7 @@ impl FileSearchService { }; if mode != SearchMode::Keyword { for chunk in &chunks { - let vector = chunk.embedding.as_ref().ok_or_else(|| { + let vector = chunk.chunk.embedding.as_ref().ok_or_else(|| { FileSearchError::Unavailable("Stored embeddings are missing; recreate this vector store".into()) })?; if dimensions.is_some_and(|expected| expected != vector.len()) @@ -749,15 +776,14 @@ impl FileSearchService { .models .rerank(&queries.join(" "), data, options.model.as_deref()) .await?; - data.retain(|result| result.score >= options.score_threshold.unwrap_or(0.0)); + data.retain(|candidate| candidate.result.score >= options.score_threshold.unwrap_or(0.0)); } data.truncate(limit); if prepare_context { data = self.prepare_context(data, permit).await?; } self.storage.refresh_activity(store_ids).await?; - let visible = self.storage.visible_result_files(store_ids, &data).await?; - data.retain(|result| visible.contains(&result.file_id)); + let data = self.visible_results(data, request.filters.as_ref()).await?; Ok(SearchResponse { object: "vector_store.search_results.page".into(), search_query: queries, diff --git a/crates/agentic-server-core/src/tool/file_search/stores.rs b/crates/agentic-server-core/src/tool/file_search/stores.rs index 6288e789..3e1fc896 100644 --- a/crates/agentic-server-core/src/tool/file_search/stores.rs +++ b/crates/agentic-server-core/src/tool/file_search/stores.rs @@ -16,9 +16,11 @@ pub(super) fn validate_store_fields( metadata.len() > 16 || metadata .iter() - .any(|(key, value)| key.is_empty() || key.len() > 64 || value.len() > 512) + .any(|(key, value)| key.is_empty() || key.chars().count() > 64 || value.chars().count() > 512) }) { - return invalid("metadata accepts at most 16 entries, with 1 to 64 byte keys and values up to 512 bytes"); + return invalid( + "metadata accepts at most 16 entries, with 1 to 64 character keys and values up to 512 characters", + ); } if let Some(expiration) = expiration { expiration.validate()?; diff --git a/crates/agentic-server-core/src/types/event.rs b/crates/agentic-server-core/src/types/event.rs index ac154ae5..970c2ab4 100644 --- a/crates/agentic-server-core/src/types/event.rs +++ b/crates/agentic-server-core/src/types/event.rs @@ -113,3 +113,16 @@ mod tests { assert_eq!("unknown".parse::().unwrap(), MessageStatus::InProgress); } } + +/// Grounded file citation emitted before its completed output text content. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename = "response.output_text.annotation.added")] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +pub struct OutputTextFileCitationAdded { + pub item_id: String, + pub output_index: u64, + pub content_index: usize, + pub annotation_index: usize, + pub sequence_number: u64, + pub annotation: crate::types::io::FileCitation, +} diff --git a/crates/agentic-server-core/src/types/file_search.rs b/crates/agentic-server-core/src/types/file_search.rs index f7c5fba0..cc1c28e1 100644 --- a/crates/agentic-server-core/src/types/file_search.rs +++ b/crates/agentic-server-core/src/types/file_search.rs @@ -679,8 +679,8 @@ pub(crate) fn validate_attributes(attributes: &FileAttributes) -> Result<(), Fil return invalid("attributes accepts at most 16 keys"); } for (key, value) in attributes { - if key.is_empty() || key.len() > 64 { - return invalid("attribute keys must contain 1 to 64 bytes"); + if key.is_empty() || key.chars().count() > 64 { + return invalid("attribute keys must contain 1 to 64 characters"); } validate_attribute_value(value)?; } @@ -689,7 +689,9 @@ pub(crate) fn validate_attributes(attributes: &FileAttributes) -> Result<(), Fil fn validate_attribute_value(value: &AttributeValue) -> Result<(), FileSearchError> { match value { - AttributeValue::String(value) if value.len() > 512 => invalid("attribute values must not exceed 512 bytes"), + AttributeValue::String(value) if value.chars().count() > 512 => { + invalid("attribute values must not exceed 512 characters") + } AttributeValue::Number(value) if !value.is_finite() => invalid("attribute numbers must be finite"), _ => Ok(()), } @@ -711,8 +713,8 @@ impl SearchFilter { } } Self::Comparison(filter) => { - if filter.key.is_empty() || filter.key.len() > 64 { - return invalid("filter keys must contain 1 to 64 bytes"); + if filter.key.is_empty() || filter.key.chars().count() > 64 { + return invalid("filter keys must contain 1 to 64 characters"); } match (&filter.operator, &filter.value) { (ComparisonOperator::In | ComparisonOperator::Nin, FilterValue::List(values)) @@ -981,3 +983,26 @@ pub struct FileBatchObject { pub status: BatchStatus, pub file_counts: FileCounts, } + +#[cfg(test)] +mod unicode_contract_tests { + use super::*; + + #[test] + fn attribute_and_filter_limits_count_characters_without_loosening_boundaries() { + for (key_length, value_length, valid) in [(64, 512, true), (65, 512, false), (64, 513, false)] { + let key = "界".repeat(key_length); + let value = "文".repeat(value_length); + let attributes = [(key.clone(), AttributeValue::String(value.clone()))].into(); + assert_eq!(validate_attributes(&attributes).is_ok(), valid); + for filter_value in [serde_json::json!(value), serde_json::json!([value])] { + let operator = if filter_value.is_array() { "in" } else { "eq" }; + let request: SearchRequest = serde_json::from_value(serde_json::json!({ + "query":"policy", "filters":{"type":operator,"key":key,"value":filter_value} + })) + .unwrap(); + assert_eq!(request.validate().is_ok(), valid); + } + } + } +} diff --git a/crates/agentic-server-core/src/types/request_response.rs b/crates/agentic-server-core/src/types/request_response.rs index 1a0de021..e991e524 100644 --- a/crates/agentic-server-core/src/types/request_response.rs +++ b/crates/agentic-server-core/src/types/request_response.rs @@ -486,10 +486,16 @@ pub struct ResponsePayload { pub previous_response_id: Option, pub conversation_id: Option, pub instructions: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tools: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tool_choice: Option, + /// Always serialized; defaults permit reading older response payloads. + #[serde(default)] + #[cfg_attr(feature = "openapi", schema(required = true))] + pub tools: Vec, + #[serde(default)] + #[cfg_attr(feature = "openapi", schema(required = true))] + pub tool_choice: ToolChoice, + #[serde(default)] + #[cfg_attr(feature = "openapi", schema(required = true))] + pub parallel_tool_calls: bool, } impl ResponsePayload { @@ -1253,8 +1259,9 @@ mod tests { previous_response_id: None, conversation_id: None, instructions: None, - tools: None, - tool_choice: None, + tools: Vec::new(), + tool_choice: ToolChoice::Auto, + parallel_tool_calls: false, }; for (status, expected_type) in [ @@ -1288,8 +1295,9 @@ mod tests { previous_response_id: None, conversation_id: None, instructions: None, - tools: None, - tool_choice: None, + tools: Vec::new(), + tool_choice: ToolChoice::Auto, + parallel_tool_calls: false, }; let chunk = payload.as_created_response_chunk(); diff --git a/crates/agentic-server-core/tests/file_search_models.rs b/crates/agentic-server-core/tests/file_search_models.rs index f6593d28..a8ef9807 100644 --- a/crates/agentic-server-core/tests/file_search_models.rs +++ b/crates/agentic-server-core/tests/file_search_models.rs @@ -940,3 +940,219 @@ async fn store_expiration_during_reranking_rejects_cached_candidates() { Some(1) ); } + +#[derive(Clone, Copy, Debug)] +enum AttachmentMutation { + DifferentStoreExcluded, + DifferentStoreMatching, + SameStoreReplacement, + AttributesExcluded, + AttributesMatching, +} + +async fn attachment_mutation_during_rerank(backend: Option, mutation: AttachmentMutation) { + let mode = if backend.is_some() { "semantic" } else { "keyword" }; + let mut setup = setup(backend.is_some()).await; + if let Some(backend) = backend { + let url = std::env::var("TEST_POSTGRES_URL").unwrap(); + setup.pool = create_pool_with_schema(Some(&url)).await.unwrap(); + setup.config.backend = backend; + setup.service = FileSearchService::new( + setup.pool.clone(), + Arc::new(reqwest::Client::new()), + setup.config.clone(), + ) + .unwrap(); + } + let service = &setup.service; + let mut stores = Vec::new(); + for _ in 0..2 { + stores.push( + service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap() + .id, + ); + } + stores.sort(); + let file_id = attach(service, &stores[0], "coral preferred", None).await; + let original: FileAttributes = serde_json::from_value(json!({"scope":"included", "label":"original"})).unwrap(); + service + .update_vector_store_file( + &stores[0], + &file_id, + UpdateVectorStoreFileRequest { + attributes: original.clone(), + }, + ) + .await + .unwrap(); + if matches!( + mutation, + AttachmentMutation::DifferentStoreExcluded | AttachmentMutation::DifferentStoreMatching + ) { + let scope = if matches!(mutation, AttachmentMutation::DifferentStoreMatching) { + "included" + } else { + "excluded" + }; + service + .attach_file( + &stores[1], + AttachFileRequest { + file_id: file_id.clone(), + attributes: serde_json::from_value(json!({"scope":scope, "label":"survivor"})).unwrap(), + ..Default::default() + }, + ) + .await + .unwrap(); + } + *setup.state.rerank_pause.lock().unwrap() = true; + let search_service = service.clone(); + let search_stores = stores.clone(); + let request: SearchRequest = serde_json::from_value(json!({ + "query":"coral", "search_mode":mode, "ranking_options":{"ranker":"neural"}, + "filters":{"type":"eq","key":"scope","value":"included"} + })) + .unwrap(); + let search = tokio::spawn(async move { search_service.search(&search_stores, &request).await }); + tokio::time::timeout(std::time::Duration::from_secs(5), setup.state.rerank_started.notified()) + .await + .unwrap(); + mutate_attachment(service, &stores[0], &file_id, original, mutation).await; + setup.state.rerank_resume.notify_one(); + let response = tokio::time::timeout(std::time::Duration::from_secs(5), search) + .await + .unwrap() + .unwrap() + .unwrap(); + for store in stores { + service.delete_vector_store(&store).await.unwrap(); + } + service.delete_file(&file_id).await.unwrap(); + assert_attachment_search(&response, &file_id, mutation); +} + +fn assert_attachment_search(response: &SearchResponse, file_id: &str, mutation: AttachmentMutation) { + match mutation { + AttachmentMutation::DifferentStoreMatching | AttachmentMutation::AttributesMatching => { + assert_eq!( + response.data.len(), + 1, + "{mutation:?}: preserve one valid deduplicated result" + ); + assert_eq!(response.data[0].file_id, file_id); + assert_eq!(response.data[0].content[0].text, "coral preferred"); + let label = if matches!(mutation, AttachmentMutation::DifferentStoreMatching) { + "survivor" + } else { + "updated" + }; + assert_eq!( + serde_json::to_value(&response.data[0].attributes).unwrap(), + json!({"scope":"included", "label":label}) + ); + } + _ => assert!( + response.data.is_empty(), + "{mutation:?}: stale attachment must not validate cached evidence" + ), + } +} + +async fn mutate_attachment( + service: &FileSearchService, + store: &str, + file_id: &str, + original: FileAttributes, + mutation: AttachmentMutation, +) { + match mutation { + AttachmentMutation::DifferentStoreExcluded + | AttachmentMutation::DifferentStoreMatching + | AttachmentMutation::SameStoreReplacement => { + service.detach_file(store, file_id).await.unwrap(); + if matches!(mutation, AttachmentMutation::SameStoreReplacement) { + // Identical source, attributes and chunk boundaries still belong to a new attachment. + service + .attach_file( + store, + AttachFileRequest { + file_id: file_id.to_owned(), + attributes: original, + ..Default::default() + }, + ) + .await + .unwrap(); + } + } + AttachmentMutation::AttributesExcluded | AttachmentMutation::AttributesMatching => { + let scope = if matches!(mutation, AttachmentMutation::AttributesMatching) { + "included" + } else { + "excluded" + }; + service + .update_vector_store_file( + store, + file_id, + UpdateVectorStoreFileRequest { + attributes: serde_json::from_value(json!({"scope":scope, "label":"updated"})).unwrap(), + }, + ) + .await + .unwrap(); + } + } +} + +#[tokio::test] +async fn rerank_provenance_excludes_other_store_attributes() { + attachment_mutation_during_rerank(None, AttachmentMutation::DifferentStoreExcluded).await; +} +#[tokio::test] +async fn rerank_provenance_preserves_matching_deduplicated_origin() { + attachment_mutation_during_rerank(None, AttachmentMutation::DifferentStoreMatching).await; +} +#[tokio::test] +async fn rerank_provenance_rejects_same_store_replacement() { + attachment_mutation_during_rerank(None, AttachmentMutation::SameStoreReplacement).await; +} +#[tokio::test] +async fn rerank_provenance_rechecks_updated_filter() { + attachment_mutation_during_rerank(None, AttachmentMutation::AttributesExcluded).await; +} +#[tokio::test] +async fn rerank_provenance_returns_current_matching_attributes() { + attachment_mutation_during_rerank(None, AttachmentMutation::AttributesMatching).await; +} + +#[tokio::test] +#[ignore = "requires isolated TEST_POSTGRES_URL with pgvector"] +async fn postgres_rerank_provenance_exact_and_pgvector() { + for backend in [ + FileSearchBackend::Exact, + FileSearchBackend::Pgvector { + dimensions: 2, + index: PgvectorIndex::Hnsw { + m: 16, + ef_construction: 64, + ef_search: 100, + }, + candidate_limit: 50, + }, + ] { + for mutation in [ + AttachmentMutation::DifferentStoreExcluded, + AttachmentMutation::DifferentStoreMatching, + AttachmentMutation::SameStoreReplacement, + AttachmentMutation::AttributesExcluded, + AttachmentMutation::AttributesMatching, + ] { + attachment_mutation_during_rerank(Some(backend.clone()), mutation).await; + } + } +} diff --git a/crates/agentic-server-core/tests/file_search_service.rs b/crates/agentic-server-core/tests/file_search_service.rs index ef9ce10b..f0bd11eb 100644 --- a/crates/agentic-server-core/tests/file_search_service.rs +++ b/crates/agentic-server-core/tests/file_search_service.rs @@ -1449,6 +1449,7 @@ async fn postgres_pgvector_indexed_semantics_restart_filters_and_deletion() { "coral reef", [ ("region".into(), AttributeValue::String("sea".into())), + ("界".repeat(64), AttributeValue::String("文".repeat(512))), ("rank".into(), AttributeValue::Number(42.0)), ("active".into(), AttributeValue::Boolean(true)), ] @@ -1509,6 +1510,10 @@ async fn postgres_pgvector_indexed_semantics_restart_filters_and_deletion() { .is_empty() ); for (filter, expected) in [ + ( + serde_json::json!({"type":"eq", "key":"界".repeat(64), "value":"文".repeat(512)}), + 1, + ), (serde_json::json!({"type":"eq", "key":"region", "value":"sea"}), 1), (serde_json::json!({"type":"ne", "key":"region", "value":"sea"}), 0), (serde_json::json!({"type":"ne", "key":"region", "value":5}), 0), diff --git a/crates/agentic-server-core/tests/file_search_tool_test.rs b/crates/agentic-server-core/tests/file_search_tool_test.rs index 0e6f28db..b455aad9 100644 --- a/crates/agentic-server-core/tests/file_search_tool_test.rs +++ b/crates/agentic-server-core/tests/file_search_tool_test.rs @@ -83,11 +83,19 @@ fn upstream(items: Vec, streaming: bool) -> support::MockResponse { } events.push(json!({"type":"response.output_item.added","output_index":index,"item":started})); if item["type"] == "message" { - let text = item["content"][0]["text"].as_str().unwrap(); - events.push(json!({"type":"response.content_part.added","output_index":index,"content_index":0,"item_id":item["id"],"part":{"type":"output_text","text":"","annotations":[]}})); - events.push(json!({"type":"response.output_text.delta","output_index":index,"content_index":0,"item_id":item["id"],"delta":text})); - events.push(json!({"type":"response.output_text.done","output_index":index,"content_index":0,"item_id":item["id"],"text":text})); - events.push(json!({"type":"response.content_part.done","output_index":index,"content_index":0,"item_id":item["id"],"part":item["content"][0]})); + for (content_index, part) in item["content"].as_array().unwrap().iter().enumerate() { + let text = part["text"].as_str().unwrap(); + events.push(json!({"type":"response.content_part.added","output_index":index,"content_index":content_index,"item_id":item["id"],"part":{"type":"output_text","text":"","annotations":[]}})); + let split = text.find('【').map_or(text.len(), |offset| offset + '【'.len_utf8()); + for delta in [&text[..split], &text[split..]] { + events.push(json!({"type":"response.output_text.delta","output_index":index,"content_index":content_index,"item_id":item["id"],"delta":delta})); + } + for (annotation_index, annotation) in part["annotations"].as_array().unwrap().iter().enumerate() { + events.push(json!({"type":"response.output_text.annotation.added","output_index":index,"content_index":content_index,"item_id":item["id"],"annotation_index":annotation_index,"annotation":annotation})); + } + events.push(json!({"type":"response.output_text.done","output_index":index,"content_index":content_index,"item_id":item["id"],"text":text})); + events.push(json!({"type":"response.content_part.done","output_index":index,"content_index":content_index,"item_id":item["id"],"part":part})); + } } events.push(json!({"type":"response.output_item.done","output_index":index,"item":item})); } @@ -155,6 +163,21 @@ fn assert_file_search_stream(events: &[Value], output: &Value) { .find(|event| event["type"] == "response.content_part.done") .unwrap(); assert_eq!(content_done["part"], output[1]["content"][0]); + let annotations: Vec<_> = events + .iter() + .filter(|event| event["type"] == "response.output_text.annotation.added") + .collect(); + assert_eq!( + annotations.len(), + 1, + "one grounded annotation across all done representations" + ); + assert_eq!(annotations[0]["annotation"], output[1]["content"][0]["annotations"][0]); + assert_eq!(annotations[0]["item_id"], "msg_answer"); + assert_eq!(annotations[0]["output_index"], 1); + assert_eq!(annotations[0]["content_index"], 0); + assert_eq!(annotations[0]["annotation_index"], 0); + assert!(annotations[0]["sequence_number"].as_u64() < content_done["sequence_number"].as_u64()); assert_eq!( events .iter() @@ -164,6 +187,27 @@ fn assert_file_search_stream(events: &[Value], output: &Value) { ); } +async fn assert_stored_effective_tools(ctx: &ExecutionContext, response_id: &str, tools: &Value) { + let lookup = support::make_request("", true, false, Some(response_id.to_owned()), None); + let stored = ctx + .resp_handler + .get(&agentic_core::executor::RequestContext { + original_request: lookup.clone(), + enriched_request: lookup, + new_input_items: vec![], + response_id: String::new(), + conversation_id: None, + conversation_version: None, + }) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(stored.metadata.effective_tool_choice).unwrap(), + "auto" + ); + assert_eq!(serde_json::to_value(stored.metadata.effective_tools).unwrap(), *tools); +} + #[tokio::test] async fn file_search_blocking_and_streaming_preserve_context_include_and_citations() { for streaming in [false, true] { @@ -172,15 +216,23 @@ async fn file_search_blocking_and_streaming_preserve_context_include_and_citatio let llm = support::MockServer::start_deque(vec![ upstream(vec![search_call()], streaming), upstream(vec![answer(&file_id)], streaming), - upstream(vec![answer(&file_id)], false), + upstream(vec![answer(&file_id)], streaming), ]) .await; let ctx = context(service, llm.url()).await; let mut request = support::make_request("What does the coral policy protect?", true, streaming, None, None); request.tools = Some(serde_json::from_value(json!([{"type":"file_search","vector_store_ids":[store_id]}])).unwrap()); + request.tool_choice = Some(agentic_core::types::io::ToolChoice::Required); request.include = include_results.then(|| vec!["file_search_call.results".to_owned()]); let (response, events) = collect(ExecuteRequest::new(request, Arc::clone(&ctx)).run().await.unwrap()).await; + let public = serde_json::to_value(&response).unwrap(); + assert_eq!( + public["tools"], + json!([{"type":"file_search", "vector_store_ids":[store_id]}]) + ); + assert_eq!(public["tool_choice"], "required"); + assert_eq!(public["parallel_tool_calls"], false); let output = serde_json::to_value(&response.output).unwrap(); assert_eq!(output[0]["type"], "file_search_call"); assert_eq!(output[0]["status"], "completed"); @@ -215,8 +267,19 @@ async fn file_search_blocking_and_streaming_preserve_context_include_and_citatio if streaming { assert_file_search_stream(&events, &output); } - let continuation = support::make_request("Explain that policy", true, false, Some(response.id), None); - let (continued, _) = collect(ExecuteRequest::new(continuation, ctx).run().await.unwrap()).await; + assert_stored_effective_tools(&ctx, &response.id, &public["tools"]).await; + let continuation = support::make_request("Explain that policy", true, streaming, Some(response.id), None); + let (continued, continued_events) = + collect(ExecuteRequest::new(continuation, ctx).run().await.unwrap()).await; + if streaming { + let annotations: Vec<_> = continued_events + .iter() + .filter(|event| event["type"] == "response.output_text.annotation.added") + .collect(); + assert_eq!(annotations.len(), 1); + assert_eq!(annotations[0]["output_index"], 0); + assert_eq!(annotations[0]["annotation"]["file_id"], file_id); + } let continued = serde_json::to_value(continued).unwrap(); assert_eq!( continued["output"][0]["content"][0]["annotations"][0]["file_id"], @@ -388,3 +451,70 @@ async fn file_search_selector_rejects_missing_declaration_before_inference() { assert!(llm.request_bodies().await.is_empty()); } } + +#[tokio::test] +async fn stream_citations_reconcile_upstream_events_and_preserve_other_annotations() { + let (service, store_id, file_id, _files) = fixture().await; + let url = + json!({"type":"url_citation","start_index":0,"end_index":1,"url":"https://example.org","title":"Reference"}); + let mut message = answer(&file_id); + message["content"] = json!([ + {"type":"output_text", "text":format!("【{file_id}】 and 【{file_id}】 【forged】"), "annotations":[url, + {"type":"file_citation","file_id":file_id,"filename":"untrusted.txt","index":0}, + {"type":"file_citation","file_id":"forged","filename":"forged.txt","index":0}]}, + {"type":"output_text", "text":format!("Second 【{file_id}】"), "annotations":[]} + ]); + let mut mock = upstream(vec![message], true); + if let support::MockResponse::Sse(body) = &mut mock { + let null = json!({"type":"response.output_text.annotation.added","item_id":"msg_answer","output_index":0,"content_index":0,"annotation_index":9,"annotation":null}); + let done = body.find("\"type\":\"response.content_part.done\"").unwrap(); + let event_start = body[..done].rfind("data: ").unwrap(); + body.insert_str(event_start, &format!("data: {null}\n\n")); + } + let llm = support::MockServer::start_deque(vec![upstream(vec![search_call()], true), mock]).await; + let ctx = context(service, llm.url()).await; + let mut request = support::make_request("Search", false, true, None, None); + request.tools = + Some(serde_json::from_value(json!([{"type":"file_search","vector_store_ids":[store_id]}])).unwrap()); + request.parallel_tool_calls = Some(true); + let (response, events) = collect(ExecuteRequest::new(request, ctx).run().await.unwrap()).await; + assert!(response.parallel_tool_calls); + let output = serde_json::to_value(response.output).unwrap(); + let annotations: Vec<_> = events + .iter() + .filter(|event| event["type"] == "response.output_text.annotation.added") + .collect(); + assert_eq!(annotations.len(), 5, "three grounded files, one URL, one null"); + assert_eq!(annotations.iter().filter(|event| event["annotation"] == url).count(), 1); + assert_eq!( + annotations.iter().filter(|event| event["annotation"].is_null()).count(), + 1 + ); + let files: Vec<_> = annotations + .iter() + .filter(|event| event["annotation"]["type"] == "file_citation") + .collect(); + assert_eq!(files.len(), 3); + for (event, (content, annotation)) in files.iter().zip([(0, 1), (0, 2), (1, 0)]) { + assert_eq!(event["output_index"], 1); + assert_eq!(event["content_index"], content); + assert_eq!(event["annotation_index"], annotation); + assert_eq!( + event["annotation"], + output[1]["content"][content]["annotations"][annotation] + ); + assert_eq!(event["annotation"]["file_id"], file_id); + assert_eq!(event["annotation"]["filename"], "policy.txt"); + let done = events + .iter() + .find(|candidate| { + candidate["type"] == "response.content_part.done" && candidate["content_index"] == content + }) + .unwrap(); + assert!(event["sequence_number"].as_u64() < done["sequence_number"].as_u64()); + } + assert_eq!(output[1]["content"][0]["annotations"][0], url); + for (sequence, event) in events.iter().enumerate() { + assert_eq!(event["sequence_number"], sequence); + } +} diff --git a/crates/agentic-server-core/tests/tool_search_test.rs b/crates/agentic-server-core/tests/tool_search_test.rs index 16a94530..b5e7dd1f 100644 --- a/crates/agentic-server-core/tests/tool_search_test.rs +++ b/crates/agentic-server-core/tests/tool_search_test.rs @@ -1311,12 +1311,11 @@ async fn namespace_nonstreaming_manual_flow_reuses_flattening_and_restoration() .expect("loaded namespace member") .remove("defer_loading"); assert_eq!( - serde_json::to_value(second.tools.as_ref().expect("response tools")).expect("response tools serialize"), + serde_json::to_value(&second.tools).expect("response tools serialize"), json!([available_namespace]) ); assert_eq!( - serde_json::to_value(second.tool_choice.as_ref().expect("response tool choice")) - .expect("response tool choice serializes"), + serde_json::to_value(&second.tool_choice).expect("response tool choice serializes"), public_choice ); diff --git a/crates/agentic-server/src/openapi.rs b/crates/agentic-server/src/openapi.rs index ae083e79..a5a7d891 100644 --- a/crates/agentic-server/src/openapi.rs +++ b/crates/agentic-server/src/openapi.rs @@ -47,6 +47,7 @@ use utoipa::OpenApi; agentic_core::types::io::FileSearchCall, agentic_core::types::io::FileSearchCallResult, agentic_core::types::io::FileCitation, + agentic_core::types::event::OutputTextFileCitationAdded, agentic_core::types::request_response::RequestPayload, agentic_core::types::request_response::ResponsePayload, agentic_core::types::request_response::CompactRequest, diff --git a/docs/api/file-search.md b/docs/api/file-search.md index 5bbc563c..d05ba3b0 100644 --- a/docs/api/file-search.md +++ b/docs/api/file-search.md @@ -111,7 +111,8 @@ embeddings exist. `file_ingestion_params.default_chunk_size_tokens` and 800/400 behavior. Static overlap remains limited to half the chunk size. `file_batch_params` validates `max_concurrent_files_per_batch` (default 3, range 1–32), `file_batch_chunk_size` (10, 1–1000), and `cleanup_interval_seconds` (86400, -1–604800); asynchronous workers are a later layer. +1–604800). These settings control the current +[durable batch workers](#durable-file-batches-and-workers). ### Contextual ingestion @@ -329,7 +330,12 @@ PDF ingestion is an optional build feature: cargo build -p agentic-server --features file-search-pdf ``` -This feature requires Rust 1.88 or newer for the parser's decompression limits. +The locked dependency graph supports this feature on Rust 1.88; the newly added +optional `aes` dependency is locked to 0.9.2. Use `--locked` for reproducible builds. +The repository still declares Rust 1.85, but existing main-branch dependencies +(`home` and ICU require 1.88; `process-wrap` requires 1.87) prevent that toolchain +from building the current workspace. This stack does not change that policy or +downgrade those unrelated dependencies. CI uses Rust 1.98. PDFs must contain extractable text; scanned PDFs require OCR before ingestion. Encrypted PDFs and documents exceeding parsing, decompression, or extracted-text limits are rejected. The default build can store and download PDFs, but returns @@ -505,7 +511,7 @@ explicit capacity errors when limits are exceeded. Each store is limited to 10,000 chunks and 64 MiB of serialized chunk data, including embeddings; ingestion enforces these limits before publication. Searching multiple stores shares the same aggregate retrieval budget. -It does not expose OGX's provider catalog or asynchronous file batches. +Provider endpoints and credentials are deployment controlled; the HTTP API does not expose OGX's provider catalog. ### Durable file batches and workers @@ -561,3 +567,93 @@ for pool acquisition/statement execution and five seconds for lock waits; SQLite uses a five-second busy timeout. Embedding HTTP calls have a 45-second timeout; contextual calls use the bounded configured timeout, and both respond to runtime cancellation. Library callers must explicitly consume the runtime with `shutdown`. + + +## Compatibility and verification + +Compatibility covers the exercised Files, Vector Stores, and Responses file-search +contracts below. It does not imply conformance with every OpenAI API or equivalent +retrieval quality to OpenAI's hosted models. The method schemas and OpenAI Python +SDK **3.13.0** provide independent wire contracts; generated OpenAPI alone is not +conformance evidence. Contract references are the +[Files schema](https://developers.openai.com/api/reference/resources/files/methods/create), +[Vector Stores search schema](https://developers.openai.com/api/reference/resources/vector_stores/methods/search), +[batch schema](https://developers.openai.com/api/reference/resources/vector_stores/subresources/file_batches/methods/create), +[Responses annotation event](https://developers.openai.com/api/reference/resources/responses/websocket-events#response.output_text.annotation.added), +and [retrieval guide](https://developers.openai.com/api/docs/guides/retrieval). + +| Surface | Implemented contract | Verification and boundaries | +| --- | --- | --- | +| Files | Create, retrieve, content, list, delete; six upload purposes; explicit expiry and batch default; purpose and keyset pagination; `object: "file"` deletion | Strict SDK plus HTTP/service tests cover file-first multipart and streamed binary hashes above the former 20 MiB upload cap. `evals` uses exact raw-wire assertions because this SDK's request enum accepts it but its response enum omits it. Only that contradictory enum case bypasses typed SDK parsing. | +| Vector Stores | Create, retrieve, update, list, delete; nullable metadata/name/policy updates; activity/deadline/status and counts | SDK and SQLite/PostgreSQL lifecycle tests verify updates, terminal expiry, publication races, and cleanup. The activity refresh policy is explicitly defined above; expired metadata persists while search data is removed. | +| Vector Store files | Attach, retrieve, attribute update, filtered/paginated list, detach, parsed content | SDK and HTTP/service tests exercise null attributes and required update fields, original extracted text in a `data` content page, overlap/context exclusion, and preservation of the independent upload. | +| File batches | Create, retrieve, cancel, filtered/paginated member lists; per-file options; `vector_store.files_batch` and five counts | Strict SDK covers completed/failed members, blocked-model cancellation, graceful shutdown/restart, and hard-crash recovery after the real lease expires. SQLite and PostgreSQL runtime tests additionally cover competing claims, expired leases, stale publication, and durable recovery. SQL and the Files mount must be shared across replicas. Model calls may repeat after recovery; fenced publication prevents stale attempts from committing. | +| Search | String or query list; attribute-key filters; 1–50 results; actual rewritten `search_query`; finite 0–1 scores; `none`, `auto`, and documented dated ranker selectors | SDK and local model fixtures exercise rewrite/embedding/rerank flow; exact SQL and real pgvector tests cover isolation, Unicode filters, and deletion. Selector compatibility does not reproduce the hosted ranking models. Bounded candidate retrieval and capacity errors apply. | +| Responses `file_search` | Public tool declarations/choice, result includes, call lifecycle, grounded file citations, `response.output_text.annotation.added`, continuation | Rust stream tests and strict SDK decoding verify contiguous sequences, rebased output indexes, content/annotation indexes, split markers, duplicate/forged citations, and agreement with done/final output. Required `tools`, `tool_choice`, and `parallel_tool_calls` are serialized on generated response objects. Each response reports its initial effective tool choice despite later internal `auto` selection; existing continuation persistence keeps the effective post-loop choice. Omitted parallel calls retain this server's effective `false` default. | + +Native `GET /v1/responses/{id}` is not implemented. Response storage supports item +history and effective settings for continuation; it does not retain a complete +Response object for the SDK's `responses.retrieve`. That general Responses CRUD +surface is outside this file-search conformance scope. + +File citation events are emitted when completed content establishes valid offsets +and final annotation order, before its content/item done event. Repeated done +representations produce one event per final file annotation. Upstream file-citation +events are reconciled through that grounded completed content; non-file and null +annotation events pass through. A citation must refer to a retrieved file in this +request or its continuation history. `annotation_index` identifies the final +annotation-array entry. The existing file citation `index` convention remains a +Unicode character offset in the output text; the method schema's prose calls it +an index in a file list without an executable example resolving that difference. +This layer verifies event/final consistency and does not claim that ambiguous +semantic detail is independently established as hosted-service parity. + +Metadata and attachment attributes allow at most 16 entries, keys of 1–64 Unicode +characters, and string values of at most 512 Unicode characters. Filter keys and +string operands use the same character limits, including array operands. Character +counts are Unicode scalar values, not UTF-8 bytes; serialized-byte resource budgets +remain independent. The retrieval guide's filename-filter `property` selector is +not implemented: the current method schema and pinned SDK require attribute `key`. +The SDK Search ranking schema includes `ranker` and `score_threshold`; guide-only +hybrid weighting examples are not treated as SDK contract evidence. + +| OGX extension | Deployment or request behavior | +| --- | --- | +| Grouped vector-store configuration | Registered provider/model identities, contextual ingestion defaults, reranker selection, rewriting, and batch settings; endpoint URLs and credentials stay deployment controlled. | +| Contextual chunking | A configured model produces document-aware embedding context while stored original source remains suitable for parsed content and citations. | +| Retrieval and fusion controls | Keyword, semantic/vector, hybrid, weighted/RRF/normalized modes and explicit registered model selectors supplement OpenAI-compatible selectors. | +| Neural/classifier reranking | Configured vLLM/Cohere transports and probability/logit score interpretation; model-specific behavior, not an OpenAI ranking-model replica. | + +| Operational resource | Bound or behavior | +| --- | --- | +| Upload and download | 512 MiB upload ceiling with bounded streaming. A download linearizes at open; later deletion/expiry cannot retract bytes already sent. | +| Ingestion | 20 MiB source input; 16 MiB extracted text; 2,048 chunks per file. Supported text formats and optional extractable-text PDF only; no OCR. | +| Chunking | Default 800 tokens with 400 overlap; bounded-block `cl100k_base` tokenization can differ from whole-document tokenization. Static and contextual limits are validated before publication. | +| Store and candidate transfer | 10,000 chunks and 64 MiB serialized chunk data per store; aggregate search/candidate budgets and explicit capacity errors apply. | +| Workers and model requests | Four owned worker/admission slots, per-batch concurrency, 30-second renewable claims, bounded model calls and parser work; graceful shutdown joins owned work and preserves resumability. | +| Response streaming | Existing 1 MiB executor-response/event budgets and bounded backpressure apply to citation events as well. No unbounded duplicate text buffer is introduced. | + +The maintained loopback suite is `scripts/tests/file-search-sdk-test.py`, with +owned fixtures in `scripts/tests/sdk_file_search_fixtures.py`. Run it against an +actual built server in an isolated environment: + +```bash +cargo build -p agentic-server --all-features --locked +uv run --no-project --python 3.13 --with openai==3.13.0 python \ + scripts/tests/file-search-sdk-test.py target/debug/agentic-server -v +``` + +Local verification of this stack passed 1,339 workspace tests, all 31 PostgreSQL +integration tests, and 17 strict SDK cases with deterministic loopback models. +Rust 1.88 passed the locked all-feature workspace check; Rust 1.98 passed the +all-feature tests (including PDF extraction) and clippy. + +The SDK uses strict response validation and disabled environment proxies. Provider +fixtures bind only to loopback, server environments exclude inherited provider +credentials, polling/streams have total deadlines, and forced cleanup never counts +as a successful shutdown test. CI installs the exact SDK pin in a separate virtual +environment, runs all-feature Rust tests/clippy, and runs every currently ignored +core PostgreSQL integration test serially against PostgreSQL 17 and pgvector 0.8.6. +Those ignored tests are explicitly database tests; future paid/network scenarios +must use a separate opt-in gate. These checks establish the listed API behavior, +not model quality or platform certification. diff --git a/docs/deploying/container.md b/docs/deploying/container.md index be3b692d..1df19600 100644 --- a/docs/deploying/container.md +++ b/docs/deploying/container.md @@ -137,10 +137,15 @@ Drain replicas running an older release before enabling writes through this rele Stored requests now fail if their response or conversation state cannot be persisted. For streaming requests, the gateway sends an error event instead of `response.completed`. Most client responses use the generic message `failed to persist response`; the underlying database error is written only to gateway logs. The exception is an optimistic conversation conflict, which returns status `400`, type `invalid_request_error`, code `conversation_locked`, and param `conversation`. No part of the stale turn is persisted, so the client can retry the request against the conversation's latest state. This prevents clients from receiving a response ID that cannot be continued after a lock timeout or other database failure without exposing database schema or constraint details. -The file search migration (`0005_file_search.sql`) adds four tables for files, -vector stores, attachments, and chunks. Supervisor-managed deployments must apply -this migration and grant the runtime role `SELECT`, `INSERT`, `UPDATE`, and `DELETE` -on the new tables before starting the upgraded gateway. File metadata and +Supervisor-managed deployments must apply every applicable migration in order +through `0008_vector_store_batches.sql`: `0005_file_search.sql`, +`0006_file_expiration.sql`, `0007_vector_store_lifecycle.sql`, and +`0008_vector_store_batches.sql` add file search, expiration/blob cleanup, store +lifecycle, and durable file batches. Grant the runtime role `SELECT`, +`INSERT`, `UPDATE`, and `DELETE` on all file search tables before starting the +upgraded gateway: `file_search_files`, `file_search_stores`, +`file_search_attachments`, `file_search_chunks`, `file_search_blob_cleanup`, +`file_search_batches`, and `file_search_jobs`. File metadata and embeddings use the same database as conversation state. Uploaded bytes use the local filesystem configured by `AGENTIC_FILES_STORAGE_DIR`. Mount persistent writable storage for the Files API, including PostgreSQL deployments. Replicas @@ -182,7 +187,15 @@ curl --fail http://127.0.0.1:9000/ready The container CI workflow builds the image, verifies that build tools are absent, launches the gateway against a mock upstream, checks both probes, and exercises a stored Responses API request through SQLite persistence. HTTP streaming and WebSockets use the same gateway binary and exposed port; the image does not add a transport proxy. -On `SIGTERM`, the gateway stops accepting connections and gives in-flight requests up to eight seconds to drain before closing the remaining connections. Set an orchestrator termination grace period longer than eight seconds; the default 30-second Kubernetes grace period and the documented 10-second Docker stop timeout both satisfy this requirement. +On `SIGTERM`, the gateway stops accepting connections and gives in-flight HTTP +requests up to eight seconds to drain before closing the remaining connections. +Full process shutdown also joins the file-search workers; cooperative parser, +filesystem, and SQL/COMMIT work can outlast that HTTP drain. Size the orchestrator +termination grace period for both drains using the configured +[database timeouts](#postgresql-production-settings) and the +[file-search worker shutdown guidance](../api/file-search.md#durable-file-batches-and-workers). +The ten-second Docker example and Kubernetes's default 30 seconds are not +unconditional full-process shutdown guarantees. ## Kubernetes and OpenShift security context diff --git a/scripts/tests/file-search-sdk-test.py b/scripts/tests/file-search-sdk-test.py new file mode 100644 index 00000000..2f4808f4 --- /dev/null +++ b/scripts/tests/file-search-sdk-test.py @@ -0,0 +1,564 @@ +"""Strict Files/Vector Stores/Responses contracts against the actual local binary. + +Run: uv run --no-project --python 3.13 --with openai==3.13.0 python + scripts/tests/file-search-sdk-test.py /path/to/agentic-server +All provider and SDK traffic uses owned loopback fixtures and sanitized environments. +""" + +import hashlib +import signal +import sys +import tempfile +import time +import unittest + +from sdk_file_search_fixtures import Gateway, ResponsesModel, RetrievalModels + +import httpx2 as httpx +from openai import BadRequestError, OpenAI + + +class FileSearchSDKContract(unittest.TestCase): + client: OpenAI + models: RetrievalModels + + def setUp(self): + # Per-request timeouts cannot bound a stream that keeps sending partial + # bytes. Bound the entire test too, including SDK decoding and polling. + def expired(_signal, _frame): + raise TimeoutError("SDK contract test exceeded its 90-second deadline") + + previous = signal.signal(signal.SIGALRM, expired) + self.addCleanup(signal.signal, signal.SIGALRM, previous) + self.addCleanup(signal.setitimer, signal.ITIMER_REAL, 0) + signal.setitimer(signal.ITIMER_REAL, 90) + self.files = [] + self.stores = [] + + def tearDown(self): + self.models.unblock() + for store in self.stores: + self.client.vector_stores.delete(store) + for file in self.files: + self.client.files.delete(file) + + def upload( + self, + name="policy.txt", + text=b"The lunar return policy permits thirty days.", + **kwargs, + ): + obj = self.client.files.create( + file=(name, text), purpose=kwargs.pop("purpose", "assistants"), **kwargs + ) + self.files.append(obj.id) + return obj + + def store(self, **kwargs): + obj = self.client.vector_stores.create(**kwargs) + self.stores.append(obj.id) + return obj + + def wait_for_ingestion(self, initial, retrieve, total_seconds=10): + deadline = time.monotonic() + total_seconds + current = initial + while current.status == "in_progress": + remaining = deadline - time.monotonic() + self.assertGreater( + remaining, + 0, + f"Ingestion {current.id} did not finish within {total_seconds} seconds", + ) + time.sleep(min(0.1, remaining)) + current = retrieve(min(remaining, 3)) + return current + + def test_files_purpose_expiration_and_download(self): + payload = b"\x00\x01binary\xff" + file = self.upload( + "data.bin", + payload, + purpose="user_data", + expires_after={"anchor": "created_at", "seconds": 3600}, + ) + self.assertEqual(file.expires_at, file.created_at + 3600) + self.assertEqual(self.client.files.content(file.id).read(), payload) + page = self.client.files.list(purpose="user_data", limit=10000) + self.assertIn(file.id, [item.id for item in page.data]) + self.assertTrue(all(item.purpose == "user_data" for item in page.data)) + + def test_file_upload_above_old_limit_and_streamed_download(self): + size = 21 * 1024 * 1024 + with tempfile.TemporaryFile() as source: + source.seek(size - 1) + source.write(b"z") + source.seek(0) + expected = hashlib.file_digest(source, "sha256").hexdigest() + source.seek(0) + file = self.client.files.create( + file=("large.bin", source), purpose="user_data" + ) + self.files.append(file.id) + self.assertEqual(file.bytes, size) + received = 0 + digest = hashlib.sha256() + with self.client.files.with_streaming_response.content(file.id) as response: + for chunk in response.iter_bytes(chunk_size=65536): + received += len(chunk) + digest.update(chunk) + self.assertEqual(received, size) + self.assertEqual(digest.hexdigest(), expected) + + def test_file_delete_response_schema(self): + file = self.upload() + self.files.remove(file.id) + deleted = self.client.files.delete(file.id) + self.assertEqual(deleted.object, "file") + self.assertTrue(deleted.deleted) + + def test_store_nullable_updates_and_expiry(self): + store = self.store( + name="before", + metadata=None, + expires_after={"anchor": "last_active_at", "days": 1}, + ) + self.assertEqual(store.expires_at, store.last_active_at + 86400) + changed = self.client.vector_stores.update( + store.id, name="after", metadata={"department": "support"} + ) + self.assertEqual(changed.name, "after") + self.assertEqual(changed.metadata, {"department": "support"}) + self.assertIsNotNone(changed.expires_after) + cleared = self.client.vector_stores.update( + store.id, expires_after=None, metadata=None + ) + self.assertIsNone(cleared.expires_after) + self.assertEqual(cleared.name, "after") + self.assertFalse(cleared.metadata) + + def test_file_attributes_and_parsed_content(self): + file = self.upload() + store = self.store() + attachment = self.client.vector_stores.files.create( + vector_store_id=store.id, file_id=file.id, attributes=None + ) + self.wait_for_ingestion( + attachment, + lambda timeout: self.client.vector_stores.files.retrieve( + file.id, vector_store_id=store.id, timeout=timeout + ), + ) + changed = self.client.vector_stores.files.update( + file.id, vector_store_id=store.id, attributes={"department": "support"} + ) + self.assertEqual(changed.attributes, {"department": "support"}) + content = self.client.vector_stores.files.content( + file.id, vector_store_id=store.id + ) + self.assertIn( + "lunar return policy", " ".join(item.text or "" for item in content.data) + ) + hit = self.client.vector_stores.search( + store.id, + query="lunar", + filters={"type": "eq", "key": "department", "value": "support"}, + ) + self.assertEqual([item.file_id for item in hit.data], [file.id]) + self.client.vector_stores.files.update( + file.id, vector_store_id=store.id, attributes=None + ) + miss = self.client.vector_stores.search( + store.id, + query="lunar", + filters={"type": "eq", "key": "department", "value": "support"}, + ) + self.assertEqual(miss.data, []) + + def test_metadata_limits_count_unicode_characters(self): + key = "界" * 64 + value = "文" * 512 + store = self.store(metadata={key: value}) + self.assertEqual( + self.client.vector_stores.retrieve(store.id).metadata, {key: value} + ) + changed = self.client.vector_stores.update(store.id, metadata={key: "章" * 512}) + self.assertEqual(changed.metadata, {key: "章" * 512}) + with self.assertRaises(BadRequestError): + self.store(metadata={key + "界": value}) + with self.assertRaises(BadRequestError): + self.client.vector_stores.update(store.id, metadata={key: value + "文"}) + self.assertEqual( + self.client.vector_stores.retrieve(store.id).metadata, changed.metadata + ) + + def test_attribute_and_filter_limits_count_unicode_characters(self): + key = "界" * 64 + value = "文" * 512 + file = self.upload() + store = self.store() + attachment = self.client.vector_stores.files.create( + vector_store_id=store.id, file_id=file.id, attributes={key: value} + ) + done = self.wait_for_ingestion( + attachment, + lambda timeout: self.client.vector_stores.files.retrieve( + file.id, vector_store_id=store.id, timeout=timeout + ), + ) + self.assertEqual(done.attributes, {key: value}) + hits = self.client.vector_stores.search( + store.id, query="lunar", filters={"type": "eq", "key": key, "value": value} + ) + self.assertEqual([item.file_id for item in hits.data], [file.id]) + self.assertEqual(hits.data[0].attributes, {key: value}) + with self.assertRaises(BadRequestError): + self.client.vector_stores.files.update( + file.id, vector_store_id=store.id, attributes={key + "界": value} + ) + with self.assertRaises(BadRequestError): + self.client.vector_stores.files.update( + file.id, vector_store_id=store.id, attributes={key: value + "文"} + ) + self.assertEqual( + self.client.vector_stores.files.retrieve( + file.id, vector_store_id=store.id + ).attributes, + {key: value}, + ) + + def test_search_named_ranker_and_result_schema(self): + file = self.upload() + store = self.store(file_ids=[file.id]) + hits = self.client.vector_stores.search( + store.id, + query=["lunar"], + ranking_options={"ranker": "default-2024-11-15"}, + max_num_results=1, + ) + self.assertEqual([hit.file_id for hit in hits.data], [file.id]) + self.assertTrue(0 <= hits.data[0].score <= 1) + self.assertEqual(hits.model_dump()["search_query"], ["lunar"]) + + def test_batch_per_file_options_and_list_filter(self): + first = self.upload("first.txt", b"First lunar document.") + second = self.upload("second.txt", b"Second orbital document.") + store = self.store() + batch = self.client.vector_stores.file_batches.create( + store.id, + files=[ + {"file_id": first.id, "attributes": {"kind": "lunar"}}, + { + "file_id": second.id, + "chunking_strategy": { + "type": "static", + "static": { + "max_chunk_size_tokens": 100, + "chunk_overlap_tokens": 0, + }, + }, + }, + ], + ) + batch = self.wait_for_ingestion( + batch, + lambda timeout: self.client.vector_stores.file_batches.retrieve( + batch.id, vector_store_id=store.id, timeout=timeout + ), + ) + self.assertEqual(batch.object, "vector_store.files_batch") + self.assertEqual(batch.status, "completed") + self.assertEqual(batch.file_counts.completed, 2) + self.assertEqual(batch.file_counts.total, 2) + page = self.client.vector_stores.file_batches.list_files( + batch.id, vector_store_id=store.id, filter="completed", limit=1, order="asc" + ) + self.assertEqual(len(page.data), 1) + self.assertTrue(page.has_more) + following = self.client.vector_stores.file_batches.list_files( + batch.id, + vector_store_id=store.id, + filter="completed", + limit=1, + order="asc", + after=page.data[-1].id, + ) + self.assertEqual(len(following.data), 1) + self.assertNotEqual(following.data[0].id, page.data[0].id) + filtered = self.client.vector_stores.files.list(store.id, filter="failed") + self.assertEqual(filtered.data, []) + + def test_evals_raw_response_despite_sdk_enum_omission(self): + # SDK3.13.0 accepts evals input but omits it from FileObject's purpose enum. + raw = self.client.files.with_raw_response.create( + file=("evals.bin", b"evaluation payload"), purpose="evals" + ).http_response.json() + self.files.append(raw["id"]) + self.assertEqual(raw["object"], "file") + self.assertEqual(raw["purpose"], "evals") + got = self.client.files.with_raw_response.retrieve( + raw["id"] + ).http_response.json() + self.assertEqual(got["purpose"], "evals") + page = self.client.files.with_raw_response.list( + purpose="evals", limit=10000 + ).http_response.json() + self.assertIn(raw["id"], [item["id"] for item in page["data"]]) + self.assertTrue(all(item["purpose"] == "evals" for item in page["data"])) + self.assertEqual( + self.client.files.content(raw["id"]).read(), b"evaluation payload" + ) + + def test_search_rewrite_reaches_embedding_and_reports_actual_query(self): + file = self.upload() + store = self.store(file_ids=[file.id]) + offset = len(self.models.snapshot()) + hits = self.client.vector_stores.search( + store.id, + query="what is the return window", + rewrite_query=True, + ranking_options={"ranker": "default-2024-11-15"}, + ) + self.assertEqual(hits.search_query, ["lunar return policy"]) + self.assertEqual([hit.file_id for hit in hits.data], [file.id]) + requests = self.models.snapshot()[offset:] + self.assertTrue(any(path == "/v1/chat/completions" for path, _ in requests)) + self.assertTrue( + any( + path == "/v1/embeddings" and "lunar return policy" in body["input"] + for path, body in requests + ) + ) + self.assertTrue(any(path == "/rerank" for path, _ in requests)) + + def test_batch_cancel_while_embedding_is_blocked(self): + file = self.upload() + store = self.store() + self.models.block_embeddings() + batch = self.client.vector_stores.file_batches.create( + store.id, file_ids=[file.id] + ) + self.assertEqual(batch.status, "in_progress") + self.assertTrue( + self.models.entered.wait(timeout=5), "Batch worker did not call embeddings" + ) + cancelled = self.client.vector_stores.file_batches.cancel( + batch.id, vector_store_id=store.id + ) + self.assertEqual(cancelled.status, "cancelled") + self.models.unblock() + # Verify durable cancellation through the API; runtime fencing has separate Rust tests. + final = self.client.vector_stores.file_batches.retrieve( + batch.id, vector_store_id=store.id + ) + self.assertEqual(final.status, "cancelled") + self.assertEqual(final.file_counts.cancelled, 1) + self.assertEqual(final.file_counts.completed, 0) + members = self.client.vector_stores.file_batches.list_files( + batch.id, vector_store_id=store.id, filter="cancelled" + ) + self.assertEqual([item.id for item in members.data], [file.id]) + self.assertEqual( + self.client.vector_stores.search(store.id, query="lunar").data, [] + ) + self.assertEqual( + self.client.files.content(file.id).read(), + b"The lunar return policy permits thirty days.", + ) + + def test_streamed_citation_events_match_strict_completed_response(self): + request_offset = len(self.upstream.requests) + file = self.upload() + self.upstream.file_id = file.id + store = self.store(file_ids=[file.id]) + with self.client.responses.create( + model="test-model", + input="What is the lunar return policy?", + tools=[{"type": "file_search", "vector_store_ids": [store.id]}], + include=["file_search_call.results"], + stream=True, + ) as stream: + events = [] + stream_deadline = time.monotonic() + 30 + for event in stream: + assert len(events) < 256 and time.monotonic() < stream_deadline, ( + "Stream contract exceeded its bounds" + ) + events.append(event.model_dump(mode="json")) + assert [event["sequence_number"] for event in events] == list( + range(len(events)) + ) + completed = [event for event in events if event["type"] == "response.completed"] + assert len(completed) == 1, events + output = completed[0]["response"]["output"] + message_index = next( + i for i, item in enumerate(output) if item["type"] == "message" + ) + message = output[message_index] + annotations = message["content"][0]["annotations"] + assert len(annotations) == 1 and annotations[0]["file_id"] == file.id, ( + annotations + ) + added = [ + event + for event in events + if event["type"] == "response.output_text.annotation.added" + ] + assert len(added) == 1, ( + f"Expected one annotation event matching final citations; received {len(added)}" + ) + event = added[0] + assert event["annotation"] == annotations[0] + assert ( + event["item_id"], + event["output_index"], + event["content_index"], + event["annotation_index"], + ) == (message["id"], message_index, 0, 0) + content_done_index = next( + i for i, e in enumerate(events) if e["type"] == "response.content_part.done" + ) + assert events.index(event) < content_done_index + assert len(self.upstream.requests) - request_offset == 2 + response = completed[0]["response"] + self.assertFalse(response["parallel_tool_calls"]) + self.assertEqual(response["tool_choice"], "auto") + self.assertEqual(response["tools"][0]["type"], "file_search") + self.assertEqual(response["tools"][0]["vector_store_ids"], [store.id]) + # Native response retrieval is not a route in this server. Persistence is + # exercised through continuation; Rust tests inspect stored metadata. + continued = self.client.responses.create( + model="test-model", + input="Explain the policy", + previous_response_id=response["id"], + parallel_tool_calls=True, + tool_choice="none", + ) + self.assertTrue(continued.parallel_tool_calls) + self.assertEqual(continued.tool_choice, "none") + self.assertEqual(continued.tools[0].type, "file_search") + self.assertEqual(continued.output[0].content[0].annotations[0].file_id, file.id) + + def test_restart_resumes_ingestion_after_orderly_shutdown(self): + file = self.upload() + store = self.store() + self.models.block_embeddings() + batch = self.client.vector_stores.file_batches.create( + store.id, file_ids=[file.id] + ) + self.assertTrue( + self.models.entered.wait(timeout=5), + "worker must own a blocked model request", + ) + # stop() asserts zero exit and fails if emergency kill was necessary. + self.gateway.stop() + self.models.unblock() + self.gateway.start() + batch = self.wait_for_ingestion( + batch, + lambda timeout: self.client.vector_stores.file_batches.retrieve( + batch.id, vector_store_id=store.id, timeout=timeout + ), + ) + self.assertEqual(batch.status, "completed") + self.assertEqual(batch.file_counts.completed, 1) + self.assertEqual(batch.file_counts.total, 1) + members = self.client.vector_stores.file_batches.list_files( + batch.id, vector_store_id=store.id + ) + self.assertEqual([item.id for item in members.data], [file.id]) + self.assertEqual( + self.client.files.content(file.id).read(), + b"The lunar return policy permits thirty days.", + ) + + def test_crash_restart_recovers_expired_worker_claim(self): + file = self.upload() + store = self.store() + self.models.block_embeddings() + batch = self.client.vector_stores.file_batches.create( + store.id, file_ids=[file.id] + ) + self.assertTrue(self.models.entered.wait(timeout=5)) + self.gateway.crash() + self.models.unblock() + self.gateway.start() + # A hard crash cannot release its claim; wait for the real 30-second lease. + batch = self.wait_for_ingestion( + batch, + lambda timeout: self.client.vector_stores.file_batches.retrieve( + batch.id, vector_store_id=store.id, timeout=timeout + ), + total_seconds=45, + ) + self.assertEqual(batch.status, "completed") + self.assertEqual(batch.file_counts.completed, 1) + self.assertEqual(batch.file_counts.total, 1) + found = self.client.vector_stores.search(store.id, query="lunar") + self.assertEqual([item.file_id for item in found.data], [file.id]) + + def test_partial_batch_failure_keeps_valid_member(self): + valid = self.upload() + invalid = self.upload("unsupported.bin", b"unsupported") + store = self.store() + initial = self.client.vector_stores.file_batches.create( + store.id, file_ids=[valid.id, invalid.id] + ) + batch = self.wait_for_ingestion( + initial, + lambda timeout: self.client.vector_stores.file_batches.retrieve( + initial.id, vector_store_id=store.id, timeout=timeout + ), + ) + self.assertEqual(batch.status, "completed") + self.assertEqual(batch.file_counts.completed, 1) + self.assertEqual(batch.file_counts.failed, 1) + failed = self.client.vector_stores.file_batches.list_files( + batch.id, vector_store_id=store.id, filter="failed" + ) + self.assertEqual([item.id for item in failed.data], [invalid.id]) + self.assertEqual(failed.data[0].last_error.code, "unsupported_file") + + def test_files_purposes_and_pagination(self): + for purpose in ["assistants", "batch", "fine-tune", "vision", "user_data"]: + with self.subTest(purpose=purpose): + first = self.upload(purpose=purpose) + second = self.upload(purpose=purpose) + self.assertEqual(self.client.files.retrieve(first.id).purpose, purpose) + if purpose == "batch": + self.assertEqual(first.expires_at, first.created_at + 30 * 86400) + page = self.client.files.list(purpose=purpose, limit=1, order="asc") + self.assertTrue(page.has_more) + self.assertEqual([item.id for item in page.data], [first.id]) + following = self.client.files.list( + purpose=purpose, limit=1, order="asc", after=first.id + ) + self.assertEqual([item.id for item in following.data], [second.id]) + + +def main(): + binary = sys.argv.pop(1) + with ( + ResponsesModel() as upstream, + RetrievalModels() as models, + tempfile.TemporaryDirectory(prefix="agentic-sdk-contract-") as directory, + ): + models.configure(directory) + with Gateway(binary, directory, upstream) as gateway: + with OpenAI( + base_url=f"{gateway.base}/v1", + api_key="sdk-contract-test", + max_retries=0, + http_client=httpx.Client(trust_env=False), + timeout=10, + _strict_response_validation=True, + ) as client: + FileSearchSDKContract.client = client + FileSearchSDKContract.models = models + FileSearchSDKContract.upstream = upstream + FileSearchSDKContract.gateway = gateway + outcome = unittest.main(exit=False) + return 0 if outcome.result.wasSuccessful() else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/tests/sdk_file_search_fixtures.py b/scripts/tests/sdk_file_search_fixtures.py new file mode 100644 index 00000000..b86c4ff8 --- /dev/null +++ b/scripts/tests/sdk_file_search_fixtures.py @@ -0,0 +1,422 @@ +"""Loopback-only deterministic retrieval models for independent SDK contracts.""" + +import os +import socket +import subprocess +import time + +import httpx2 as httpx +from openai.types.responses import Response + +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +from pathlib import Path +import threading + + +class RetrievalModels: + def __init__(self): + self.lock = threading.Lock() + self.requests = [] + self.block = None + self.entered = threading.Event() + fixture = self + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *_args): + pass + + def setup(self): + super().setup() + self.connection.settimeout(5) + + def do_POST(self): + length = int(self.headers.get("Content-Length", "0")) + if length < 0 or length > 4 * 1024 * 1024: + self.send_error(413) + return + request = json.loads(self.rfile.read(length)) + with fixture.lock: + fixture.requests.append((self.path, request)) + block = fixture.block + if self.path == "/v1/embeddings": + if block is not None: + fixture.entered.set() + if not block.wait(10): + self.send_error(504) + return + inputs = request["input"] + if isinstance(inputs, str): + inputs = [inputs] + response = { + "object": "list", + "model": request["model"], + "data": [ + { + "object": "embedding", + "index": index, + "embedding": [1.0, 0.0], + } + for index, _ in enumerate(inputs) + ], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + elif self.path == "/v1/chat/completions": + response = { + "id": "chatcmpl-local-rewrite", + "object": "chat.completion", + "created": 0, + "model": request["model"], + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "lunar return policy", + }, + "finish_reason": "stop", + } + ], + } + elif self.path == "/rerank": + response = { + "id": "rerank-local", + "model": request["model"], + "results": [ + {"index": index, "relevance_score": 1.0 / (index + 1)} + for index, _ in enumerate(request["documents"]) + ], + } + else: + self.send_error(404) + return + body = json.dumps(response).encode() + try: + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + except (BrokenPipeError, ConnectionResetError): + pass # A cancelled ingestion deliberately drops its model request. + + self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self.server.daemon_threads = False + self.thread = threading.Thread( + target=self.server.serve_forever, name="sdk-retrieval-models" + ) + + def __enter__(self): + self.thread.start() + return self + + def __exit__(self, *_args): + self.unblock() + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=5) + if self.thread.is_alive(): + raise RuntimeError("Retrieval model fixture did not stop") + + def block_embeddings(self): + with self.lock: + self.entered.clear() + self.block = threading.Event() + + def unblock(self): + with self.lock: + if self.block is not None: + self.block.set() + self.block = None + + def snapshot(self): + with self.lock: + return list(self.requests) + + def configure(self, directory): + base = f"http://127.0.0.1:{self.server.server_port}/v1" + Path(directory, "config.toml").write_text(f'''[file_search.vector_stores] +default_provider_id = "sdk" + +[file_search.vector_stores.providers.sdk] +base_url = "{base}" +models = ["embedding", "rewrite", "rerank"] +protocol = "vllm" +score_interpretation = "probability" + +[file_search.vector_stores.default_embedding_model] +provider_id = "sdk" +model_id = "embedding" +embedding_dimensions = 2 + +[file_search.vector_stores.default_reranker_model] +provider_id = "sdk" +model_id = "rerank" + +[file_search.vector_stores.rewrite_query_params] +model = {{ provider_id = "sdk", model_id = "rewrite" }} +max_tokens = 100 +temperature = 0.0 + +[file_search.vector_stores.file_batch_params] +cleanup_interval_seconds = 1 +''') + + +class ResponsesModel: + def __init__(self): + self.file_id = None + self.requests = [] + fixture = self + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *_args): + pass + + def setup(self): + super().setup() + self.connection.settimeout(5) + + def do_POST(self): + length = int(self.headers.get("Content-Length", "0")) + if self.path != "/v1/responses" or not 0 < length <= 4 * 1024 * 1024: + self.send_error(400) + return + request = json.loads(self.rfile.read(length)) + fixture.requests.append(request) + tool_output = any( + item.get("type") == "function_call_output" + for item in request["input"] + ) + if tool_output: + text = ( + f"Returns are accepted for thirty days. 【{fixture.file_id}】" + ) + item = { + "type": "message", + "id": "msg_sdk", + "role": "assistant", + "status": "completed", + "content": [ + {"type": "output_text", "text": text, "annotations": []} + ], + } + else: + item = { + "type": "function_call", + "id": "fc_sdk", + "call_id": "call_sdk", + "name": "file_search", + "arguments": json.dumps({"queries": ["lunar"]}), + "status": "completed", + } + response = { + "id": "resp_fixture", + "object": "response", + "created_at": 0, + "model": "test-model", + "status": "completed", + "output": [item], + "parallel_tool_calls": request.get("parallel_tool_calls", True), + "tool_choice": "auto", + "tools": request.get("tools", []), + } + Response.model_validate(response) + if not request.get("stream"): + payload = json.dumps(response).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + return + pending = dict(response, status="in_progress", output=[]) + events = [ + {"type": "response.created", "response": pending}, + {"type": "response.in_progress", "response": pending}, + ] + added = dict(item, status="in_progress") + if tool_output: + added["content"] = [] + events.append( + { + "type": "response.output_item.added", + "output_index": 0, + "item": added, + } + ) + if tool_output: + indexes = { + "output_index": 0, + "content_index": 0, + "item_id": item["id"], + } + events.append( + { + "type": "response.content_part.added", + **indexes, + "part": { + "type": "output_text", + "text": "", + "annotations": [], + }, + } + ) + split = text.index("【") + 7 + for delta in (text[:split], text[split:]): + events.append( + { + "type": "response.output_text.delta", + **indexes, + "delta": delta, + "logprobs": [], + } + ) + events.append( + { + "type": "response.output_text.done", + **indexes, + "text": text, + "logprobs": [], + } + ) + events.append( + { + "type": "response.content_part.done", + **indexes, + "part": item["content"][0], + } + ) + events.append( + { + "type": "response.output_item.done", + "output_index": 0, + "item": item, + } + ) + events.append({"type": "response.completed", "response": response}) + for sequence, event in enumerate(events): + event["sequence_number"] = sequence + payload = ( + "".join(f"data: {json.dumps(event)}\n\n" for event in events) + + "data: [DONE]\n\n" + ).encode() + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self.server.daemon_threads = False + self.thread = threading.Thread( + target=self.server.serve_forever, name="sdk-responses-model" + ) + + def __enter__(self): + self.thread.start() + return self + + def __exit__(self, *_args): + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=5) + if self.thread.is_alive(): + raise RuntimeError("Responses fixture failed to stop") + + +class Gateway: + """Own one loopback process; reuse its SQL/files state for restart coverage.""" + + def __init__(self, binary, directory, upstream): + self.binary = str(Path(binary).resolve()) + self.directory = directory + self.upstream = upstream + self.process = None + self.log = None + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + self.port = listener.getsockname()[1] + self.base = f"http://127.0.0.1:{self.port}" + + def start(self): + assert self.process is None + env = { + "PATH": os.environ.get("PATH", ""), + "AGENTIC_API_HOME": self.directory, + "AGENTIC_FILES_STORAGE_DIR": f"{self.directory}/files", + "RUST_LOG": "warn", + } + self.log = open(f"{self.directory}/server.log", "a+") + self.process = subprocess.Popen( + [ + self.binary, + "--llm-api-base", + f"http://127.0.0.1:{self.upstream.server.server_port}/v1", + "--skip-llm-ready-check", + "--gateway-host", + "127.0.0.1", + "--gateway-port", + str(self.port), + ], + env=env, + stdout=self.log, + stderr=self.log, + ) + with httpx.Client(timeout=1, trust_env=False) as health: + deadline = time.monotonic() + 15 + while time.monotonic() < deadline: + if self.process.poll() is not None: + self.log.seek(0) + raise RuntimeError(self.log.read()) + try: + if health.get(f"{self.base}/health").is_success: + return + except httpx.TransportError: + pass + time.sleep(0.03) + raise RuntimeError("Server startup deadline exceeded") + + def crash(self): + """Deliberate crash recovery test; never used by graceful-shutdown checks.""" + process, self.process = self.process, None + assert process is not None and process.poll() is None + try: + process.kill() + if process.wait(timeout=5) >= 0: + raise RuntimeError("Crash fixture did not terminate by signal") + finally: + self.log.close() + + def stop(self): + process, self.process = self.process, None + if process is None: + return + try: + process.terminate() + try: + code = process.wait(timeout=12) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + raise RuntimeError( + "Server failed the graceful shutdown deadline" + ) from None + if code != 0: + self.log.seek(0) + raise RuntimeError( + f"Server did not exit cleanly ({code}): {self.log.read()}" + ) + finally: + self.log.close() + + def __enter__(self): + try: + self.start() + except BaseException: + self.stop() + raise + return self + + def __exit__(self, *_args): + self.stop()