diff --git a/app/src/ai/agent_sdk/artifact.rs b/app/src/ai/agent_sdk/artifact.rs index e9341118a1e..7783e4177d1 100644 --- a/app/src/ai/agent_sdk/artifact.rs +++ b/app/src/ai/agent_sdk/artifact.rs @@ -107,13 +107,12 @@ impl ArtifactCommandRunner { ) { let server_api = ServerApiProvider::as_ref(ctx).get(); let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client(); - let uploader = FileArtifactUploader::new(ai_client, server_api.clone()); + let uploader = FileArtifactUploader::new(ai_client, server_api); ctx.spawn( async move { let request = FileArtifactUploadRequest::try_from(args)?; let association = uploader.resolve_upload_association(&request).await?; - server_api.set_ambient_agent_task_id(Some(association.ambient_task_id)); uploader.upload_with_association(request, association).await }, move |_, result, ctx| match result { diff --git a/app/src/ai/agent_sdk/artifact_upload.rs b/app/src/ai/agent_sdk/artifact_upload.rs index 7077f4d15a1..36d2e1e5f33 100644 --- a/app/src/ai/agent_sdk/artifact_upload.rs +++ b/app/src/ai/agent_sdk/artifact_upload.rs @@ -119,14 +119,18 @@ impl FileArtifactUploader { let artifact = self.prepare_upload_artifact(path).await?; let create_response = self - .create_upload_target(association, title, description, &artifact) + .create_upload_target(&association, title, description, &artifact) .await?; let checksum = self .upload_artifact_bytes(&create_response.upload_target, &artifact) .await?; let uploaded_artifact = self - .confirm_upload(create_response.artifact.artifact_uid, checksum) + .confirm_upload( + &association.ambient_task_id, + create_response.artifact.artifact_uid, + checksum, + ) .await?; let size_bytes = i64::try_from(artifact.file_size) .context("Artifact file size exceeds supported range")?; @@ -143,24 +147,27 @@ impl FileArtifactUploader { async fn create_upload_target( &self, - association: ResolvedUploadAssociation, + association: &ResolvedUploadAssociation, title: Option, description: Option, artifact: &PreparedUploadArtifact, ) -> Result { self.ai_client - .create_file_artifact_upload_target(CreateFileArtifactUploadRequest { - conversation_id: association - .conversation_id - .as_ref() - .map(|token| token.as_str().to_string()), - run_id: association.run_id.as_ref().map(ToString::to_string), - filepath: artifact.filepath.clone(), - title, - description, - mime_type: Some(artifact.mime_type.clone()), - size_bytes: artifact.graphql_size_bytes(), - }) + .create_file_artifact_upload_target( + &association.ambient_task_id, + CreateFileArtifactUploadRequest { + conversation_id: association + .conversation_id + .as_ref() + .map(|token| token.as_str().to_string()), + run_id: association.run_id.as_ref().map(ToString::to_string), + filepath: artifact.filepath.clone(), + title, + description, + mime_type: Some(artifact.mime_type.clone()), + size_bytes: artifact.graphql_size_bytes(), + }, + ) .await .context("Failed to create file artifact upload target") } @@ -180,11 +187,12 @@ impl FileArtifactUploader { async fn confirm_upload( &self, + task_id: &AmbientAgentTaskId, artifact_uid: String, checksum: String, ) -> Result { self.ai_client - .confirm_file_artifact_upload(artifact_uid, checksum) + .confirm_file_artifact_upload(task_id, artifact_uid, checksum) .await .context("Failed to confirm file artifact upload") } diff --git a/app/src/ai/agent_sdk/artifact_upload_tests.rs b/app/src/ai/agent_sdk/artifact_upload_tests.rs index f0200bee7fe..cbdb6b6fa09 100644 --- a/app/src/ai/agent_sdk/artifact_upload_tests.rs +++ b/app/src/ai/agent_sdk/artifact_upload_tests.rs @@ -1,7 +1,9 @@ use std::path::PathBuf; +use std::sync::Arc; use std::{env, fs}; use chrono::Utc; +use mockito::Server; use tempfile::tempdir; use warp_cli::artifact::UploadArtifactArgs; @@ -11,6 +13,11 @@ use crate::ai::agent::conversation::{AIAgentHarness, ServerAIConversationMetadat use crate::cloud_object::{Revision, ServerMetadata, ServerPermissions}; use crate::persistence::model::ConversationUsageMetadata; use crate::server::ids::ServerId; +use crate::server::server_api::ServerApi; +use crate::server::server_api::ai::{ + CreateFileArtifactUploadResponse, FileArtifactRecord, FileArtifactUploadTargetInfo, + MockAIClient, +}; fn create_mock_server_metadata() -> ServerMetadata { ServerMetadata { @@ -101,6 +108,86 @@ fn file_size_and_prefix_for_path_returns_full_contents_when_prefix_exceeds_file( ); } +#[test] +fn upload_uses_resolved_task_identity_for_create_and_confirm() { + let tempdir = tempdir().unwrap(); + let path = tempdir.path().join("artifact.txt"); + fs::write(&path, b"artifact contents").unwrap(); + let mut server = Server::new(); + let upload = server + .mock("PUT", "/upload") + .match_body("artifact contents") + .with_status(200) + .create(); + + let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(); + let mut ai_client = MockAIClient::new(); + ai_client + .expect_create_file_artifact_upload_target() + .withf(|actual_task_id, request| { + actual_task_id.to_string() == "550e8400-e29b-41d4-a716-446655440000" + && request.conversation_id.as_deref() == Some("conversation-123") + }) + .once() + .returning(move |_, _| { + Ok(CreateFileArtifactUploadResponse { + artifact: FileArtifactRecord { + artifact_uid: "artifact-123".to_string(), + filepath: "artifact.txt".to_string(), + description: None, + mime_type: "text/plain".to_string(), + size_bytes: Some(17), + }, + upload_target: FileArtifactUploadTargetInfo { + url: format!("{}/upload", server.url()), + method: "PUT".to_string(), + headers: Vec::new(), + fields: Vec::new(), + }, + }) + }); + ai_client + .expect_confirm_file_artifact_upload() + .withf(|actual_task_id, artifact_uid, checksum| { + actual_task_id.to_string() == "550e8400-e29b-41d4-a716-446655440000" + && artifact_uid == "artifact-123" + && !checksum.is_empty() + }) + .once() + .returning(|_, _, _| { + Ok(FileArtifactRecord { + artifact_uid: "artifact-123".to_string(), + filepath: "artifact.txt".to_string(), + description: None, + mime_type: "text/plain".to_string(), + size_bytes: Some(17), + }) + }); + + let uploader = + FileArtifactUploader::new(Arc::new(ai_client), Arc::new(ServerApi::new_for_test())); + let request = FileArtifactUploadRequest { + path, + run_id: None, + conversation_id: Some(ServerConversationToken::new("conversation-123".to_string())), + title: None, + description: None, + }; + let association = ResolvedUploadAssociation { + conversation_id: request.conversation_id.clone(), + run_id: None, + ambient_task_id: task_id, + }; + + let completed = tokio::runtime::Runtime::new() + .unwrap() + .block_on(uploader.upload_with_association(request, association)) + .unwrap(); + + upload.assert(); + assert_eq!(completed.artifact.artifact_uid, "artifact-123"); + assert_eq!(completed.size_bytes, 17); +} #[test] fn single_conversation_metadata_returns_the_only_metadata_record() { let metadata = single_conversation_metadata( diff --git a/app/src/server/server_api.rs b/app/src/server/server_api.rs index 4f26bbd6210..60e6a456d9f 100644 --- a/app/src/server/server_api.rs +++ b/app/src/server/server_api.rs @@ -461,7 +461,7 @@ impl ServerApi { } #[cfg(any(test, all(feature = "tui", feature = "test-util")))] - fn new_for_test() -> Self { + pub(crate) fn new_for_test() -> Self { let (tx, _) = async_channel::unbounded(); let auth_state = Arc::new(AuthState::new_for_test()); let client = Arc::new(http_client::Client::new_for_test()); @@ -523,6 +523,23 @@ impl ServerApi { ) } + fn send_graphql_request_for_task<'a, QF, O: warp_graphql::client::Operation + Send + 'a>( + &'a self, + task_id: &AmbientAgentTaskId, + operation: O, + timeout: Option, + ) -> BoxFuture<'a, Result> + where + QF: 'a, + { + warp_server_client::graphql_helpers::send_graphql_request_with_ambient_policy( + &self.base_client, + operation, + timeout, + AmbientHeaderPolicy::for_task(task_id.to_string()), + ) + } + /// Opens an SSE stream to the agent event-push endpoint. /// /// The returned `EventSourceStream` yields `reqwest_eventsource::Event` diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index a15c768743d..25036faad36 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -1434,11 +1434,13 @@ pub trait AIClient: 'static + Send + Sync { async fn create_file_artifact_upload_target( &self, + task_id: &AmbientAgentTaskId, request: CreateFileArtifactUploadRequest, ) -> anyhow::Result; async fn confirm_file_artifact_upload( &self, + task_id: &AmbientAgentTaskId, artifact_uid: String, checksum: String, ) -> anyhow::Result; @@ -2718,6 +2720,7 @@ impl AIClient for ServerApi { async fn create_file_artifact_upload_target( &self, + task_id: &AmbientAgentTaskId, request: CreateFileArtifactUploadRequest, ) -> anyhow::Result { let variables = CreateFileArtifactUploadTargetVariables { @@ -2733,7 +2736,9 @@ impl AIClient for ServerApi { request_context: get_request_context(), }; let operation = CreateFileArtifactUploadTarget::build(variables); - let response = self.send_graphql_request(operation, None).await?; + let response = self + .send_graphql_request_for_task(task_id, operation, None) + .await?; match response.create_file_artifact_upload_target { CreateFileArtifactUploadTargetResult::CreateFileArtifactUploadTargetOutput(output) => { @@ -2773,6 +2778,7 @@ impl AIClient for ServerApi { async fn confirm_file_artifact_upload( &self, + task_id: &AmbientAgentTaskId, artifact_uid: String, checksum: String, ) -> anyhow::Result { @@ -2784,7 +2790,9 @@ impl AIClient for ServerApi { request_context: get_request_context(), }; let operation = ConfirmFileArtifactUpload::build(variables); - let response = self.send_graphql_request(operation, None).await?; + let response = self + .send_graphql_request_for_task(task_id, operation, None) + .await?; match response.confirm_file_artifact_upload { ConfirmFileArtifactUploadResult::ConfirmFileArtifactUploadOutput(output) => { diff --git a/crates/warp_server_client/src/base_client.rs b/crates/warp_server_client/src/base_client.rs index 108717dbec6..ff9b9481977 100644 --- a/crates/warp_server_client/src/base_client.rs +++ b/crates/warp_server_client/src/base_client.rs @@ -322,6 +322,20 @@ impl BaseClient { pub async fn graphql_request_options( &self, timeout: Option, + ) -> Result { + self.graphql_request_options_with_ambient_policy( + timeout, + AmbientHeaderPolicy::inherit_all(), + ) + .await + } + + /// Returns GraphQL options for a session-authenticated operation with request-local ambient + /// agent headers. + pub async fn graphql_request_options_with_ambient_policy( + &self, + timeout: Option, + ambient_header_policy: AmbientHeaderPolicy, ) -> Result { let auth_token = self .get_or_refresh_access_token() @@ -330,10 +344,9 @@ impl BaseClient { let mut options = self.graphql_request_options_with_token(auth_token.bearer_token()); options.timeout = timeout; options.headers = self.authenticated_graphql.headers.clone(); - options.headers.extend( - self.ambient_headers(AmbientHeaderPolicy::inherit_all()) - .await?, - ); + options + .headers + .extend(self.ambient_headers(ambient_header_policy).await?); Ok(options) } diff --git a/crates/warp_server_client/src/base_client_tests.rs b/crates/warp_server_client/src/base_client_tests.rs index a28dba6d8f1..bae6b094ba0 100644 --- a/crates/warp_server_client/src/base_client_tests.rs +++ b/crates/warp_server_client/src/base_client_tests.rs @@ -138,6 +138,43 @@ fn authenticated_graphql_options_include_configured_and_ambient_headers() { ); } +#[test] +fn authenticated_graphql_options_can_override_task_id_for_one_request() { + let client = client(); + client.set_ambient_agent_task_id(Some("ambient-task".to_string())); + + let task_scoped = block_on(client.graphql_request_options_with_ambient_policy( + None, + AmbientHeaderPolicy { + workload_token: HeaderOverride::Set("workload-token".to_string()), + ..AmbientHeaderPolicy::for_task("child-task") + }, + )) + .unwrap(); + let inherited = block_on(client.graphql_request_options_with_ambient_policy( + None, + AmbientHeaderPolicy { + workload_token: HeaderOverride::Set("workload-token".to_string()), + ..AmbientHeaderPolicy::inherit_all() + }, + )) + .unwrap(); + + assert_eq!( + task_scoped + .headers + .get(CLOUD_AGENT_ID_HEADER) + .map(String::as_str), + Some("child-task") + ); + assert_eq!( + inherited + .headers + .get(CLOUD_AGENT_ID_HEADER) + .map(String::as_str), + Some("ambient-task") + ); +} #[test] fn authenticated_graphql_configuration_cannot_override_base_client_owned_headers() { let (event_sender, _) = async_channel::unbounded(); diff --git a/crates/warp_server_client/src/graphql_helpers.rs b/crates/warp_server_client/src/graphql_helpers.rs index a60e898d922..29e04bb13d5 100644 --- a/crates/warp_server_client/src/graphql_helpers.rs +++ b/crates/warp_server_client/src/graphql_helpers.rs @@ -8,7 +8,7 @@ use warp_graphql::client::{GraphQLError, Operation}; use warpui_core::r#async::BoxFuture; use crate::auth::AuthEvent; -use crate::base_client::BaseClient; +use crate::base_client::{AmbientHeaderPolicy, BaseClient}; /// Sends a GraphQL operation through a base client supplied by the application. /// @@ -19,12 +19,32 @@ pub fn send_graphql_request<'a, QF: 'a, O>( operation: O, timeout: Option, ) -> BoxFuture<'a, Result> +where + O: Operation + Send + 'a, +{ + send_graphql_request_with_ambient_policy( + base_client, + operation, + timeout, + AmbientHeaderPolicy::inherit_all(), + ) +} + +/// Sends a GraphQL operation with request-local ambient agent headers. +pub fn send_graphql_request_with_ambient_policy<'a, QF: 'a, O>( + base_client: &'a BaseClient, + operation: O, + timeout: Option, + ambient_header_policy: AmbientHeaderPolicy, +) -> BoxFuture<'a, Result> where O: Operation + Send + 'a, { Box::pin(async move { let operation_name = operation.operation_name().map(Cow::into_owned); - let options = base_client.graphql_request_options(timeout).await?; + let options = base_client + .graphql_request_options_with_ambient_policy(timeout, ambient_header_policy) + .await?; let response = match operation .send_request(base_client.owned_http_client(), options) .await