Scope native artifact uploads to child tasks

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 {

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<String>,
         description: Option<String>,
         artifact: &PreparedUploadArtifact,
     ) -> Result<CreateFileArtifactUploadResponse> {
         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<FileArtifactRecord> {
         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")
     }

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(

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<QF> + Send + 'a>(
+        &'a self,
+        task_id: &AmbientAgentTaskId,
+        operation: O,
+        timeout: Option<Duration>,
+    ) -> BoxFuture<'a, Result<QF>>
+    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`

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<CreateFileArtifactUploadResponse, anyhow::Error>;
 
     async fn confirm_file_artifact_upload(
         &self,
+        task_id: &AmbientAgentTaskId,
         artifact_uid: String,
         checksum: String,
     ) -> anyhow::Result<FileArtifactRecord, anyhow::Error>;
@@ -2718,6 +2720,7 @@ impl AIClient for ServerApi {
 
     async fn create_file_artifact_upload_target(
         &self,
+        task_id: &AmbientAgentTaskId,
         request: CreateFileArtifactUploadRequest,
     ) -> anyhow::Result<CreateFileArtifactUploadResponse, anyhow::Error> {
         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<FileArtifactRecord, anyhow::Error> {
@@ -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) => {

crates/warp_server_client/src/base_client.rs

@@ -322,6 +322,20 @@ impl BaseClient {
     pub async fn graphql_request_options(
         &self,
         timeout: Option<Duration>,
+    ) -> Result<RequestOptions> {
+        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<Duration>,
+        ambient_header_policy: AmbientHeaderPolicy,
     ) -> Result<RequestOptions> {
         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)
     }
 

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();