diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index c9a82df965..5c4358ee34 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -43,22 +43,22 @@ pub use audit::read_thread_state_audit_rows; /// Most consumers should prefer [`StateRuntime`]. pub use extract::apply_rollout_item; pub use extract::rollout_item_affects_thread_metadata; +pub use model::AddThreadAttachmentOutcome; pub use model::Anchor; pub use model::BackfillState; pub use model::BackfillStats; pub use model::BackfillStatus; pub use model::DirectionalThreadSpawnEdgeStatus; pub use model::ExtractionOutcome; +pub use model::RemoveThreadAttachmentOutcome; pub use model::SortDirection; pub use model::SortKey; pub use model::Stage1JobClaim; pub use model::Stage1JobClaimOutcome; pub use model::Stage1Output; pub use model::Stage1StartupClaimParams; -pub use model::ThreadArtifact; -pub use model::ThreadArtifactAttachmentOutcome; -pub use model::ThreadArtifactPage; -pub use model::ThreadArtifactRemovalOutcome; +pub use model::ThreadAttachment; +pub use model::ThreadAttachmentPage; pub use model::ThreadGoal; pub use model::ThreadGoalStatus; pub use model::ThreadMetadata; @@ -99,6 +99,18 @@ pub use telemetry::record_fallback; /// Maximum number of pending user submissions permitted for one thread. pub const MAX_QUEUE_ITEMS: usize = 100; +/// Maximum serialized size of one persisted thread-attachment payload. +pub const MAX_THREAD_ATTACHMENT_PAYLOAD_BYTES: usize = 64 * 1024; + +/// Maximum byte length of a persisted attachment type. +pub const MAX_THREAD_ATTACHMENT_TYPE_BYTES: usize = 256; + +/// Maximum byte length of a persisted stable attachment identity key. +pub const MAX_THREAD_ATTACHMENT_IDENTITY_KEY_BYTES: usize = 256; + +/// Maximum number of active attachments retained for one thread. +pub const MAX_THREAD_ATTACHMENTS_PER_THREAD: usize = 100; + /// Stable UUIDv7 identifying the built-in pinned thread section. pub const PINNED_THREAD_SECTION_ID: &str = "01984de2-8f74-7c91-a3b2-5c5e937cf318"; diff --git a/codex-rs/state/src/model/mod.rs b/codex-rs/state/src/model/mod.rs index 7adc1afefc..1693af9933 100644 --- a/codex-rs/state/src/model/mod.rs +++ b/codex-rs/state/src/model/mod.rs @@ -5,7 +5,7 @@ mod memories; mod project; mod queued_item; mod rollout_migration_state; -mod thread_artifact; +mod thread_attachment; mod thread_goal; mod thread_metadata; @@ -29,10 +29,10 @@ pub use queued_item::QueuedUserSubmissionRecord; pub use rollout_migration_state::RolloutMigrationCursor; pub use rollout_migration_state::RolloutMigrationSkippedRollout; pub use rollout_migration_state::RolloutMigrationState; -pub use thread_artifact::ThreadArtifact; -pub use thread_artifact::ThreadArtifactAttachmentOutcome; -pub use thread_artifact::ThreadArtifactPage; -pub use thread_artifact::ThreadArtifactRemovalOutcome; +pub use thread_attachment::AddThreadAttachmentOutcome; +pub use thread_attachment::RemoveThreadAttachmentOutcome; +pub use thread_attachment::ThreadAttachment; +pub use thread_attachment::ThreadAttachmentPage; pub use thread_goal::ThreadGoal; pub use thread_goal::ThreadGoalStatus; pub use thread_metadata::Anchor; diff --git a/codex-rs/state/src/model/thread_artifact.rs b/codex-rs/state/src/model/thread_artifact.rs deleted file mode 100644 index cf04445de2..0000000000 --- a/codex-rs/state/src/model/thread_artifact.rs +++ /dev/null @@ -1,46 +0,0 @@ -use codex_protocol::ThreadId; -use serde_json::Value; - -/// A bounded artifact durably associated with one thread. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ThreadArtifact { - /// Stable, server-assigned UUIDv7 artifact identity. - pub id: String, - /// Thread that owns this artifact. - pub thread_id: ThreadId, - /// Client-defined artifact category. - pub artifact_type: String, - /// Client-defined stable identity within the owning thread and artifact category. - pub identity_key: String, - /// Bounded, client-defined artifact metadata. - pub payload: Value, - /// Integer Unix timestamp in seconds when the artifact was attached. - pub created_at: i64, -} - -/// Result of attaching one uniquely identified thread artifact. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ThreadArtifactAttachmentOutcome { - /// A new durable artifact was created. - Created(ThreadArtifact), - /// The artifact was already attached; its payload and creation time are unchanged. - Existing(ThreadArtifact), -} - -/// Result of removing a thread artifact. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ThreadArtifactRemovalOutcome { - /// An attached artifact was removed. - Removed(ThreadArtifact), - /// No artifact with the requested identity was attached. - NotFound, -} - -/// One deterministically ordered page of artifacts across selected threads. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ThreadArtifactPage { - /// Artifacts ordered by thread identity, creation time, and artifact identity. - pub artifacts: Vec, - /// Opaque cursor for the next page, or `None` when the selection is exhausted. - pub next_cursor: Option, -} diff --git a/codex-rs/state/src/model/thread_attachment.rs b/codex-rs/state/src/model/thread_attachment.rs new file mode 100644 index 0000000000..0712b21386 --- /dev/null +++ b/codex-rs/state/src/model/thread_attachment.rs @@ -0,0 +1,48 @@ +//! Thread attachment records and outcomes; membership changes independently of resource contents. + +use codex_protocol::ThreadId; +use serde_json::Value; + +/// A bounded attachment durably associated with one thread. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ThreadAttachment { + /// Stable, server-assigned UUIDv7 attachment identity. + pub id: String, + /// Thread that owns this attachment. + pub thread_id: ThreadId, + /// Client-defined attachment category. + pub attachment_type: String, + /// Client-defined stable identity within the owning thread and attachment category. + pub identity_key: String, + /// Bounded, client-defined attachment metadata. + pub payload: Value, + /// Integer Unix timestamp in seconds when the attachment was attached. + pub created_at: i64, +} + +/// Result of attaching one uniquely identified thread attachment. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AddThreadAttachmentOutcome { + /// A new durable attachment was created. + Created(ThreadAttachment), + /// The attachment was already attached; its payload and creation time are unchanged. + Existing(ThreadAttachment), +} + +/// Result of removing a thread attachment. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RemoveThreadAttachmentOutcome { + /// An attached attachment was removed. + Removed(ThreadAttachment), + /// No attachment with the requested identity was attached. + NotFound, +} + +/// One deterministically ordered page of attachments across selected threads. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ThreadAttachmentPage { + /// Attachments ordered by thread identity, creation time, and attachment identity. + pub attachments: Vec, + /// Opaque cursor for the next page, or `None` when the selection is exhausted. + pub next_cursor: Option, +} diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 1c571c0104..376ad682ed 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -52,6 +52,7 @@ mod remote_control; mod rollout_migration; #[cfg(test)] pub(crate) mod test_support; +mod thread_attachments; mod thread_section_order; mod thread_sections; mod threads; diff --git a/codex-rs/state/src/runtime/thread_attachments.rs b/codex-rs/state/src/runtime/thread_attachments.rs new file mode 100644 index 0000000000..abe2c34423 --- /dev/null +++ b/codex-rs/state/src/runtime/thread_attachments.rs @@ -0,0 +1,173 @@ +//! Transactional thread attachment storage using the attachment SQL schema. + +use super::StateRuntime; +use crate::AddThreadAttachmentOutcome; +use crate::MAX_THREAD_ATTACHMENT_IDENTITY_KEY_BYTES; +use crate::MAX_THREAD_ATTACHMENT_PAYLOAD_BYTES; +use crate::MAX_THREAD_ATTACHMENT_TYPE_BYTES; +use crate::MAX_THREAD_ATTACHMENTS_PER_THREAD; +use crate::RemoveThreadAttachmentOutcome; +use crate::ThreadAttachment; +use anyhow::Context; +use chrono::Utc; +use codex_protocol::ThreadId; +use serde_json::Value; +use sqlx::Row; +use sqlx::sqlite::SqliteRow; +use uuid::Uuid; + +impl StateRuntime { + /// Attach an attachment once, returning an existing attachment for repeated requests. + pub async fn add_thread_attachment( + &self, + thread_id: ThreadId, + attachment_type: &str, + identity_key: &str, + payload: &Value, + ) -> anyhow::Result { + validate_attachment_identity(attachment_type, identity_key)?; + let serialized_payload = serde_json::to_string(payload).context( + "invalid thread attachment request: attachment payload cannot be serialized", + )?; + if serialized_payload.len() > MAX_THREAD_ATTACHMENT_PAYLOAD_BYTES { + anyhow::bail!( + "invalid thread attachment request: attachment payload exceeds {MAX_THREAD_ATTACHMENT_PAYLOAD_BYTES} bytes" + ); + } + + let mut transaction = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let thread_id_string = thread_id.to_string(); + let thread_exists = sqlx::query_scalar::<_, i64>("SELECT 1 FROM threads WHERE id = ?") + .bind(&thread_id_string) + .fetch_optional(&mut *transaction) + .await? + .is_some(); + if !thread_exists { + anyhow::bail!("thread not found: {thread_id}"); + } + + let existing = sqlx::query( + "SELECT id, thread_id, attachment_type, identity_key, payload, created_at FROM thread_attachments WHERE thread_id = ? AND attachment_type = ? AND identity_key = ?", + ) + .bind(&thread_id_string) + .bind(attachment_type) + .bind(identity_key) + .fetch_optional(&mut *transaction) + .await?; + if let Some(existing) = existing { + let attachment = attachment_from_row(&existing)?; + transaction.commit().await?; + return Ok(AddThreadAttachmentOutcome::Existing(attachment)); + } + + let identity_count = sqlx::query_scalar::<_, i64>( + "SELECT COUNT(*) FROM thread_attachments WHERE thread_id = ?", + ) + .bind(&thread_id_string) + .fetch_one(&mut *transaction) + .await?; + if usize::try_from(identity_count)? >= MAX_THREAD_ATTACHMENTS_PER_THREAD { + anyhow::bail!( + "invalid thread attachment request: thread attachment identity count exceeds {MAX_THREAD_ATTACHMENTS_PER_THREAD}" + ); + } + + let attachment = ThreadAttachment { + id: Uuid::now_v7().to_string(), + thread_id, + attachment_type: attachment_type.to_string(), + identity_key: identity_key.to_string(), + payload: payload.clone(), + created_at: Utc::now().timestamp(), + }; + sqlx::query( + "INSERT INTO thread_attachments (id, thread_id, attachment_type, identity_key, payload, created_at) VALUES (?, ?, ?, ?, ?, ?)", + ) + .bind(&attachment.id) + .bind(&thread_id_string) + .bind(attachment_type) + .bind(identity_key) + .bind(&serialized_payload) + .bind(attachment.created_at) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(AddThreadAttachmentOutcome::Created(attachment)) + } + + /// Remove an attached attachment, immediately freeing its slot. + pub async fn remove_thread_attachment( + &self, + thread_id: ThreadId, + attachment_type: &str, + identity_key: &str, + ) -> anyhow::Result { + validate_attachment_identity(attachment_type, identity_key)?; + let mut transaction = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let thread_id_string = thread_id.to_string(); + let thread_exists = sqlx::query_scalar::<_, i64>("SELECT 1 FROM threads WHERE id = ?") + .bind(&thread_id_string) + .fetch_optional(&mut *transaction) + .await? + .is_some(); + if !thread_exists { + anyhow::bail!("thread not found: {thread_id}"); + } + + let removed = sqlx::query( + "DELETE FROM thread_attachments WHERE thread_id = ? AND attachment_type = ? AND identity_key = ? RETURNING id, thread_id, attachment_type, identity_key, payload, created_at", + ) + .bind(&thread_id_string) + .bind(attachment_type) + .bind(identity_key) + .fetch_optional(&mut *transaction) + .await?; + let outcome = match removed { + Some(row) => RemoveThreadAttachmentOutcome::Removed(attachment_from_row(&row)?), + None => RemoveThreadAttachmentOutcome::NotFound, + }; + transaction.commit().await?; + Ok(outcome) + } +} + +fn validate_attachment_identity(attachment_type: &str, identity_key: &str) -> anyhow::Result<()> { + if attachment_type.trim().is_empty() { + anyhow::bail!("invalid thread attachment request: attachment type must not be empty"); + } + if attachment_type.len() > MAX_THREAD_ATTACHMENT_TYPE_BYTES { + anyhow::bail!( + "invalid thread attachment request: attachment type exceeds {MAX_THREAD_ATTACHMENT_TYPE_BYTES} bytes" + ); + } + if identity_key.trim().is_empty() { + anyhow::bail!( + "invalid thread attachment request: attachment identity key must not be empty" + ); + } + if identity_key.len() > MAX_THREAD_ATTACHMENT_IDENTITY_KEY_BYTES { + anyhow::bail!( + "invalid thread attachment request: attachment identity key exceeds {MAX_THREAD_ATTACHMENT_IDENTITY_KEY_BYTES} bytes" + ); + } + Ok(()) +} + +fn attachment_from_row(row: &SqliteRow) -> anyhow::Result { + let thread_id: String = row.try_get("thread_id")?; + let payload: String = row.try_get("payload")?; + Ok(ThreadAttachment { + id: row.try_get("id")?, + thread_id: ThreadId::from_string(&thread_id) + .context("invalid persisted thread attachment owner")?, + attachment_type: row.try_get("attachment_type")?, + identity_key: row.try_get("identity_key")?, + payload: serde_json::from_str(&payload) + .context("invalid persisted thread attachment payload")?, + created_at: row.try_get("created_at")?, + }) +} + +#[cfg(test)] +#[path = "thread_attachments_tests.rs"] +mod tests; diff --git a/codex-rs/state/src/runtime/thread_attachments_tests.rs b/codex-rs/state/src/runtime/thread_attachments_tests.rs new file mode 100644 index 0000000000..24afdf5c00 --- /dev/null +++ b/codex-rs/state/src/runtime/thread_attachments_tests.rs @@ -0,0 +1,250 @@ +//! Coverage for attachment identity, capacity, pagination, and durable lifecycle behavior. + +use super::StateRuntime; +use crate::AddThreadAttachmentOutcome; +use crate::MAX_THREAD_ATTACHMENT_IDENTITY_KEY_BYTES; +use crate::MAX_THREAD_ATTACHMENT_PAYLOAD_BYTES; +use crate::MAX_THREAD_ATTACHMENT_TYPE_BYTES; +use crate::MAX_THREAD_ATTACHMENTS_PER_THREAD; +use crate::RemoveThreadAttachmentOutcome; +use crate::runtime::test_support::test_thread_metadata; +use crate::runtime::test_support::unique_temp_dir; +use anyhow::Result; +use codex_protocol::ThreadId; +use codex_utils_absolute_path::test_support::PathExt; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::path::PathBuf; +use std::sync::Arc; + +async fn runtime_with_threads(count: usize) -> Result<(Arc, PathBuf, Vec)> { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init( + crate::SqliteConfig::new_for_testing(codex_home.as_path().abs()), + "test-provider".to_string(), + ) + .await?; + let mut thread_ids = Vec::with_capacity(count); + for _ in 0..count { + let thread_id = ThreadId::new(); + runtime + .upsert_thread(&test_thread_metadata( + &codex_home, + thread_id, + codex_home.clone(), + )) + .await?; + thread_ids.push(thread_id); + } + Ok((runtime, codex_home, thread_ids)) +} + +#[tokio::test] +async fn attachment_attachments_are_idempotent_and_scoped_to_their_thread() -> Result<()> { + let (runtime, _codex_home, thread_ids) = runtime_with_threads(/*count*/ 2).await?; + let first = runtime + .add_thread_attachment( + thread_ids[0], + "pull_request", + "openai/codex#123", + &json!({ "url": "https://github.com/openai/codex/pull/123" }), + ) + .await?; + let AddThreadAttachmentOutcome::Created(first) = first else { + anyhow::bail!("first attachment should create an attachment"); + }; + assert_eq!(first.thread_id, thread_ids[0]); + + let repeated = runtime + .add_thread_attachment( + thread_ids[0], + "pull_request", + "openai/codex#123", + &json!({ "url": "https://github.com/openai/codex/pull/456" }), + ) + .await?; + assert_eq!( + repeated, + AddThreadAttachmentOutcome::Existing(first.clone()) + ); + + let other = runtime + .add_thread_attachment( + thread_ids[1], + "pull_request", + "openai/codex#123", + &json!({ "url": "https://github.com/openai/codex/pull/123" }), + ) + .await?; + let AddThreadAttachmentOutcome::Created(other) = other else { + anyhow::bail!("the same identity on another thread should create its own attachment"); + }; + assert_ne!(first.id, other.id); + Ok(()) +} + +#[tokio::test] +async fn attachment_removals_return_not_found_or_the_removed_record() -> Result<()> { + let (runtime, _codex_home, thread_ids) = runtime_with_threads(/*count*/ 1).await?; + let thread_id = thread_ids[0]; + let removed_before_attach = runtime + .remove_thread_attachment(thread_id, "pull_request", "openai/codex#123") + .await?; + assert_eq!( + removed_before_attach, + RemoveThreadAttachmentOutcome::NotFound + ); + + let explicit = runtime + .add_thread_attachment( + thread_id, + "pull_request", + "openai/codex#123", + &json!({ "url": "https://github.com/openai/codex/pull/123" }), + ) + .await?; + let AddThreadAttachmentOutcome::Created(explicit) = explicit else { + anyhow::bail!("attachment should create an attachment"); + }; + assert_eq!( + runtime + .remove_thread_attachment(thread_id, "pull_request", "openai/codex#123") + .await?, + RemoveThreadAttachmentOutcome::Removed(explicit) + ); + Ok(()) +} + +#[tokio::test] +async fn active_attachment_limit_is_freed_by_removal() -> Result<()> { + let (runtime, _codex_home, thread_ids) = runtime_with_threads(/*count*/ 1).await?; + let thread_id = thread_ids[0]; + + for index in 0..MAX_THREAD_ATTACHMENTS_PER_THREAD { + runtime + .add_thread_attachment( + thread_id, + "pull_request", + &format!("attachment-{index}"), + &json!({ "index": index }), + ) + .await?; + } + + let existing = runtime + .add_thread_attachment( + thread_id, + "pull_request", + "attachment-0", + &json!({ "index": "different" }), + ) + .await?; + assert!(matches!(existing, AddThreadAttachmentOutcome::Existing(_))); + let attachment_limit = runtime + .add_thread_attachment(thread_id, "pull_request", "one-too-many", &json!({})) + .await + .expect_err("a new attachment must not exceed the per-thread limit"); + assert!( + attachment_limit.to_string().contains( + "invalid thread attachment request: thread attachment identity count exceeds" + ) + ); + + let removed = runtime + .remove_thread_attachment(thread_id, "pull_request", "attachment-0") + .await?; + assert!(matches!(removed, RemoveThreadAttachmentOutcome::Removed(_))); + assert_eq!( + runtime + .remove_thread_attachment(thread_id, "pull_request", "attachment-0") + .await?, + RemoveThreadAttachmentOutcome::NotFound + ); + let replacement = runtime + .add_thread_attachment(thread_id, "pull_request", "one-too-many", &json!({})) + .await?; + assert!(matches!( + replacement, + AddThreadAttachmentOutcome::Created(_) + )); + + Ok(()) +} + +#[tokio::test] +async fn attachment_mutations_reject_invalid_identity_payload_and_unknown_threads() -> Result<()> { + let (runtime, _codex_home, thread_ids) = runtime_with_threads(/*count*/ 1).await?; + let thread_id = thread_ids[0]; + for (attachment_type, identity_key, expected) in [ + (" ".to_string(), "pr".to_string(), "type must not be empty"), + ( + "x".repeat(MAX_THREAD_ATTACHMENT_TYPE_BYTES + 1), + "pr".to_string(), + "type exceeds", + ), + ( + "pull_request".to_string(), + " ".to_string(), + "key must not be empty", + ), + ( + "pull_request".to_string(), + "x".repeat(MAX_THREAD_ATTACHMENT_IDENTITY_KEY_BYTES + 1), + "key exceeds", + ), + ] { + let error = runtime + .add_thread_attachment(thread_id, &attachment_type, &identity_key, &json!({})) + .await + .expect_err("invalid attachment identities must be rejected"); + assert!(error.to_string().contains(expected)); + } + + let too_large = runtime + .add_thread_attachment( + thread_id, + "pull_request", + "pr", + &json!("x".repeat(MAX_THREAD_ATTACHMENT_PAYLOAD_BYTES)), + ) + .await + .expect_err("oversized attachment payload must be rejected"); + assert!(too_large.to_string().contains("payload exceeds")); + let missing = ThreadId::new(); + let missing_error = runtime + .add_thread_attachment(missing, "pull_request", "pr", &json!({})) + .await + .expect_err("missing owners must be rejected"); + assert!(missing_error.to_string().contains("thread not found")); + Ok(()) +} + +#[tokio::test] +async fn concurrent_attachment_attachments_preserve_one_deterministic_identity() -> Result<()> { + let (runtime, _codex_home, thread_ids) = runtime_with_threads(/*count*/ 1).await?; + let thread_id = thread_ids[0]; + let mut joins = Vec::new(); + for _ in 0..8 { + let runtime = Arc::clone(&runtime); + joins.push(tokio::spawn(async move { + runtime + .add_thread_attachment( + thread_id, + "pull_request", + "openai/codex#123", + &json!({ "url": "https://github.com/openai/codex/pull/123" }), + ) + .await + })); + } + + let mut created = 0; + for join in joins { + match join.await?? { + AddThreadAttachmentOutcome::Created(_) => created += 1, + AddThreadAttachmentOutcome::Existing(_) => {} + } + } + assert_eq!(created, 1); + Ok(()) +}