Initial public release of Warp.

Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b
2026-04-28 08:43:33 -05:00
commit 0dbd3d567a
4982 changed files with 1431549 additions and 0 deletions
@@ -0,0 +1,386 @@
+use std::env;
+use std::fs::File;
+use std::io::Read as _;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+
+use anyhow::{anyhow, bail, Context, Result};
+use blocking::unblock;
+use mime_guess::from_path;
+use warp_cli::artifact::UploadArtifactArgs;
+
+use super::common::parse_ambient_task_id;
+use crate::ai::agent::api::ServerConversationToken;
+use crate::ai::agent::conversation::ServerAIConversationMetadata;
+use crate::ai::ambient_agents::AmbientAgentTaskId;
+use crate::server::server_api::ai::{
+    AIClient, CreateFileArtifactUploadRequest, CreateFileArtifactUploadResponse,
+    FileArtifactRecord, FileArtifactUploadTargetInfo,
+};
+use crate::server::server_api::presigned_upload::upload_file_to_target;
+use crate::server::server_api::ServerApi;
+
+const MIME_SNIFF_BYTES: usize = 8 * 1024;
+const OZ_RUN_ID_ENV_VAR: &str = "OZ_RUN_ID";
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub(crate) struct FileArtifactUploadRequest {
+    pub(crate) path: PathBuf,
+    pub(crate) run_id: Option<AmbientAgentTaskId>,
+    pub(crate) conversation_id: Option<ServerConversationToken>,
+    pub(crate) description: Option<String>,
+}
+
+impl TryFrom<UploadArtifactArgs> for FileArtifactUploadRequest {
+    type Error = anyhow::Error;
+
+    fn try_from(value: UploadArtifactArgs) -> Result<Self> {
+        let run_id = match value.run_id {
+            Some(run_id) => Some(parse_run_id(&run_id, "Invalid run ID")?),
+            None => None,
+        };
+
+        Ok(Self {
+            path: value.path,
+            run_id,
+            conversation_id: value.conversation_id.map(ServerConversationToken::new),
+            description: value.description,
+        })
+    }
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct CompletedFileArtifactUpload {
+    pub(crate) artifact: FileArtifactRecord,
+    pub(crate) size_bytes: i64,
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub(crate) struct ResolvedUploadAssociation {
+    conversation_id: Option<ServerConversationToken>,
+    run_id: Option<AmbientAgentTaskId>,
+    pub(crate) ambient_task_id: AmbientAgentTaskId,
+}
+
+#[derive(Debug, Clone)]
+struct PreparedUploadArtifact {
+    path: PathBuf,
+    filepath: String,
+    mime_type: String,
+    file_size: u64,
+}
+
+impl PreparedUploadArtifact {
+    fn from_path(path: PathBuf) -> Result<Self> {
+        // `infer` only needs leading signature bytes, so avoid buffering the whole artifact
+        // before we stream the file body to the upload target.
+        let (file_size, mime_sniff_bytes) = file_size_and_prefix_for_path(&path, MIME_SNIFF_BYTES)?;
+
+        Ok(Self {
+            filepath: normalize_artifact_filepath(&path),
+            mime_type: infer_mime_type(&path, &mime_sniff_bytes),
+            file_size,
+            path,
+        })
+    }
+
+    fn graphql_size_bytes(&self) -> Option<i32> {
+        checked_graphql_size_bytes_for_upload(&self.path, self.file_size)
+    }
+}
+
+pub(crate) struct FileArtifactUploader {
+    ai_client: Arc<dyn AIClient>,
+    server_api: Arc<ServerApi>,
+}
+
+impl FileArtifactUploader {
+    pub(crate) fn new(ai_client: Arc<dyn AIClient>, server_api: Arc<ServerApi>) -> Self {
+        Self {
+            ai_client,
+            server_api,
+        }
+    }
+
+    pub(crate) async fn upload_with_association(
+        &self,
+        request: FileArtifactUploadRequest,
+        association: ResolvedUploadAssociation,
+    ) -> Result<CompletedFileArtifactUpload> {
+        let FileArtifactUploadRequest {
+            path, description, ..
+        } = request;
+
+        let artifact = self.prepare_upload_artifact(path).await?;
+        let create_response = self
+            .create_upload_target(association, 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)
+            .await?;
+        let size_bytes = i64::try_from(artifact.file_size)
+            .context("Artifact file size exceeds supported range")?;
+
+        Ok(CompletedFileArtifactUpload {
+            artifact: uploaded_artifact,
+            size_bytes,
+        })
+    }
+
+    async fn prepare_upload_artifact(&self, path: PathBuf) -> Result<PreparedUploadArtifact> {
+        unblock(move || PreparedUploadArtifact::from_path(path)).await
+    }
+
+    async fn create_upload_target(
+        &self,
+        association: ResolvedUploadAssociation,
+        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(),
+                description,
+                mime_type: Some(artifact.mime_type.clone()),
+                size_bytes: artifact.graphql_size_bytes(),
+            })
+            .await
+            .context("Failed to create file artifact upload target")
+    }
+
+    async fn upload_artifact_bytes(
+        &self,
+        target: &FileArtifactUploadTargetInfo,
+        artifact: &PreparedUploadArtifact,
+    ) -> Result<String> {
+        upload_file_to_target(
+            self.server_api.http_client(),
+            target,
+            &artifact.path,
+            artifact.file_size,
+        )
+        .await
+    }
+
+    async fn confirm_upload(
+        &self,
+        artifact_uid: String,
+        checksum: String,
+    ) -> Result<FileArtifactRecord> {
+        self.ai_client
+            .confirm_file_artifact_upload(artifact_uid, checksum)
+            .await
+            .context("Failed to confirm file artifact upload")
+    }
+
+    pub(crate) async fn resolve_upload_association(
+        &self,
+        request: &FileArtifactUploadRequest,
+    ) -> Result<ResolvedUploadAssociation> {
+        let conversation_task_id = match (request.run_id.as_ref(), request.conversation_id.as_ref())
+        {
+            // we were given a conversation id, so we need to resolve the task id from the conversation via the api
+            (None, Some(conversation_id)) => {
+                Some(self.resolve_conversation_task_id(conversation_id).await)
+            }
+            _ => None,
+        };
+
+        resolve_upload_association_from_sources(
+            request.run_id,
+            request.conversation_id.clone(),
+            conversation_task_id,
+            load_env_run_id()?,
+        )
+    }
+
+    async fn resolve_conversation_task_id(
+        &self,
+        conversation_id: &ServerConversationToken,
+    ) -> Result<AmbientAgentTaskId> {
+        let metadata = self
+            .ai_client
+            .list_ai_conversation_metadata(Some(vec![conversation_id.as_str().to_string()]))
+            .await
+            .with_context(|| {
+                format!(
+                    "Failed to load conversation '{}' to resolve artifact upload headers",
+                    conversation_id.as_str()
+                )
+            })?;
+
+        let metadata = single_conversation_metadata(conversation_id.as_str(), metadata)
+            .with_context(|| {
+                format!(
+                    "Failed to load conversation '{}' to resolve artifact upload headers",
+                    conversation_id.as_str()
+                )
+            })?;
+
+        ambient_task_id_from_conversation_metadata(conversation_id.as_str(), metadata)
+    }
+}
+
+fn normalize_artifact_filepath(path: &Path) -> String {
+    path.to_string_lossy().replace('\\', "/")
+}
+
+fn infer_mime_type(path: &Path, file_bytes: &[u8]) -> String {
+    infer::get(file_bytes)
+        .map(|kind| kind.mime_type().to_string())
+        .unwrap_or_else(|| from_path(path).first_or_octet_stream().to_string())
+}
+
+fn file_size_and_prefix_for_path(path: &Path, max_bytes: usize) -> Result<(u64, Vec<u8>)> {
+    let mut file = File::open(path)
+        .with_context(|| format!("Failed to open artifact file '{}'", path.display()))?;
+    let file_size = file
+        .metadata()
+        .with_context(|| format!("Failed to stat artifact file '{}'", path.display()))?
+        .len();
+    let mut bytes = vec![0; max_bytes];
+    let bytes_read = file
+        .read(&mut bytes)
+        .with_context(|| format!("Failed to read artifact file '{}'", path.display()))?;
+    bytes.truncate(bytes_read);
+    Ok((file_size, bytes))
+}
+
+fn checked_graphql_size_bytes_for_upload(path: &Path, size_bytes: u64) -> Option<i32> {
+    let graphql_size_bytes = i32::try_from(size_bytes).ok();
+    if graphql_size_bytes.is_none() {
+        // The backing upload can handle large files, but the GraphQL field is still `Int`.
+        // Dropping `size_bytes` preserves the upload request instead of failing on conversion.
+        log::warn!(
+            "Artifact file '{}' is {} bytes, which exceeds the GraphQL size_bytes limit of {} bytes; omitting size_bytes from the upload target request",
+            path.display(),
+            size_bytes,
+            i32::MAX,
+        );
+    }
+
+    graphql_size_bytes
+}
+
+fn single_conversation_metadata(
+    conversation_id: &str,
+    mut metadata: Vec<ServerAIConversationMetadata>,
+) -> Result<ServerAIConversationMetadata> {
+    match metadata.len() {
+        0 => bail!("Conversation not found"),
+        1 => Ok(metadata.pop().expect("metadata length checked")),
+        _ => bail!("Multiple conversations found for '{conversation_id}'"),
+    }
+}
+
+fn ambient_task_id_from_conversation_metadata(
+    conversation_id: &str,
+    metadata: ServerAIConversationMetadata,
+) -> Result<AmbientAgentTaskId> {
+    metadata.ambient_agent_task_id.ok_or_else(|| {
+        anyhow!("Conversation '{conversation_id}' is not backed by a cloud agent task")
+    })
+}
+
+fn parse_run_id(run_id: &str, error_prefix: &str) -> Result<AmbientAgentTaskId> {
+    parse_ambient_task_id(run_id, error_prefix)
+}
+
+fn load_env_run_id() -> Result<Option<String>> {
+    match env::var(OZ_RUN_ID_ENV_VAR) {
+        Ok(run_id) => Ok(Some(run_id)),
+        Err(env::VarError::NotPresent) => Ok(None),
+        Err(env::VarError::NotUnicode(_)) => Err(anyhow!(
+            "{OZ_RUN_ID_ENV_VAR} is set but is not valid Unicode"
+        )),
+    }
+}
+
+fn resolve_env_run_id(env_run_id: Option<String>) -> Result<AmbientAgentTaskId> {
+    let Some(run_id) = env_run_id else {
+        bail!("{OZ_RUN_ID_ENV_VAR} is not set");
+    };
+
+    parse_run_id(&run_id, "Invalid OZ_RUN_ID")
+}
+
+fn resolve_upload_association_from_sources(
+    explicit_run_id: Option<AmbientAgentTaskId>,
+    explicit_conversation_id: Option<ServerConversationToken>,
+    conversation_task_id: Option<Result<AmbientAgentTaskId>>,
+    env_run_id: Option<String>,
+) -> Result<ResolvedUploadAssociation> {
+    // Precedence is deliberate:
+    // 1. An explicit run ID is authoritative and must not silently fall back.
+    // 2. A conversation ID stays attached to the artifact even if we have to borrow the ambient
+    //    task ID from `OZ_RUN_ID` because the conversation lacks cloud-task metadata.
+    // 3. `OZ_RUN_ID` becomes the sole source of truth only when the caller supplied nothing else.
+    if let Some(run_id) = explicit_run_id {
+        let ambient_task_id = run_id;
+        return Ok(ResolvedUploadAssociation {
+            conversation_id: None,
+            run_id: Some(run_id),
+            ambient_task_id,
+        });
+    }
+
+    if let Some(conversation_id) = explicit_conversation_id {
+        match conversation_task_id
+            .ok_or_else(|| anyhow!("conversation resolution should be provided"))?
+        {
+            Ok(ambient_task_id) => {
+                return Ok(ResolvedUploadAssociation {
+                    conversation_id: Some(conversation_id),
+                    run_id: None,
+                    ambient_task_id,
+                });
+            }
+            Err(conversation_err) => {
+                let env_err = match resolve_env_run_id(env_run_id) {
+                    Ok(ambient_task_id) => {
+                        log::warn!(
+                            "Conversation '{}' task resolution failed ({conversation_err}); falling back to {OZ_RUN_ID_ENV_VAR} for ambient task context",
+                            conversation_id.as_str()
+                        );
+                        return Ok(ResolvedUploadAssociation {
+                            conversation_id: Some(conversation_id),
+                            run_id: None,
+                            ambient_task_id,
+                        });
+                    }
+                    Err(env_err) => env_err,
+                };
+
+                return Err(anyhow!(
+                    "Failed to resolve artifact upload association for conversation '{}': {conversation_err}; also failed to use {OZ_RUN_ID_ENV_VAR}: {env_err}",
+                    conversation_id.as_str()
+                ));
+            }
+        }
+    }
+
+    let ambient_task_id = resolve_env_run_id(env_run_id).map_err(|env_err| {
+        anyhow!(
+            "Failed to resolve artifact upload association: no usable --run-id or --conversation-id was provided, and {OZ_RUN_ID_ENV_VAR}: {env_err}"
+        )
+    })?;
+
+    Ok(ResolvedUploadAssociation {
+        conversation_id: None,
+        run_id: Some(ambient_task_id),
+        ambient_task_id,
+    })
+}
+
+#[cfg(test)]
+#[path = "artifact_upload_tests.rs"]
+mod tests;