From 12f7ffdb6683403a93558bc5aef8b416291e9c78 Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Mon, 16 Feb 2026 23:14:53 -0800 Subject: [PATCH] Make share metadata creation atomic --- codex-rs/core/src/session_share.rs | 111 +++++++++++++++++++++++++++-- 1 file changed, 106 insertions(+), 5 deletions(-) diff --git a/codex-rs/core/src/session_share.rs b/codex-rs/core/src/session_share.rs index ffcad1c6db..500654da8a 100644 --- a/codex-rs/core/src/session_share.rs +++ b/codex-rs/core/src/session_share.rs @@ -106,6 +106,22 @@ impl SessionObjectStore { } } + async fn put_object_if_absent( + &self, + key: &str, + data: Vec, + content_type: &str, + ) -> anyhow::Result { + match self { + SessionObjectStore::Http(store) => { + store.put_object_if_absent(key, data, content_type).await + } + SessionObjectStore::Azure(store) => { + store.put_object_if_absent(key, data, content_type).await + } + } + } + async fn put_object_file( &self, key: &str, @@ -194,11 +210,32 @@ pub async fn upload_rollout_with_owner( created_at: now, updated_at: now, }; - upload_meta(&store, &meta_key, &meta).await?; - store - .put_object_file(&key, rollout_path, "application/x-ndjson") - .await - .with_context(|| format!("failed to upload rollout for id {session_id}"))?; + let created = create_meta_if_absent(&store, &meta_key, &meta).await?; + if created { + store + .put_object_file(&key, rollout_path, "application/x-ndjson") + .await + .with_context(|| format!("failed to upload rollout for id {session_id}"))?; + } else { + let meta = fetch_meta(&store, &meta_key).await?.ok_or_else(|| { + anyhow::anyhow!("failed to atomically create share metadata; try again") + })?; + if meta.owner != owner { + return Err(anyhow::anyhow!( + "remote session metadata already exists and belongs to another user" + )); + } + store + .put_object_file(&key, rollout_path, "application/x-ndjson") + .await + .with_context(|| format!("failed to upload rollout for id {session_id}"))?; + let updated = SessionShareMeta { + owner: meta.owner, + created_at: meta.created_at, + updated_at: now, + }; + upload_meta(&store, &meta_key, &updated).await?; + } } } @@ -284,6 +321,17 @@ async fn upload_meta( Ok(()) } +async fn create_meta_if_absent( + store: &SessionObjectStore, + key: &str, + meta: &SessionShareMeta, +) -> anyhow::Result { + let payload = serde_json::to_vec(meta).with_context(|| "failed to serialize metadata")?; + store + .put_object_if_absent(key, payload, "application/json") + .await +} + fn build_rollout_download_path(codex_home: &Path, session_id: ThreadId) -> anyhow::Result { let timestamp = OffsetDateTime::now_local() .map_err(|e| anyhow::anyhow!("failed to get local time: {e}"))?; @@ -373,6 +421,30 @@ impl HttpObjectStore { } } + async fn put_object_if_absent( + &self, + key: &str, + data: Vec, + content_type: &str, + ) -> anyhow::Result { + let url = self.object_url(key)?; + let response = self + .client + .put(url) + .header(reqwest::header::IF_NONE_MATCH, "*") + .header(reqwest::header::CONTENT_TYPE, content_type) + .body(data) + .send() + .await?; + match response.status() { + status if status.is_success() => Ok(true), + StatusCode::PRECONDITION_FAILED | StatusCode::CONFLICT => Ok(false), + status => Err(anyhow::anyhow!( + "object store conditional PUT failed with status {status}" + )), + } + } + async fn put_object_file( &self, key: &str, @@ -533,6 +605,35 @@ impl AzureObjectStore { } } + async fn put_object_if_absent( + &self, + key: &str, + data: Vec, + content_type: &str, + ) -> anyhow::Result { + let url = self.object_url(key)?; + let response = self + .authorized_request( + self.client + .put(url) + .header("x-ms-blob-type", "BlockBlob") + .header(reqwest::header::IF_NONE_MATCH, "*") + .header(reqwest::header::CONTENT_TYPE, content_type) + .body(data), + ) + .await? + .send() + .await?; + match response.status() { + status if status.is_success() => Ok(true), + StatusCode::PRECONDITION_FAILED => Ok(false), + status => Err(anyhow::anyhow!( + "azure blob conditional PUT failed with status {status}{}", + azure_response_context(response.headers()) + )), + } + } + async fn put_object_file( &self, key: &str,