From 130d6e4fba5612e2e6e0724e6463b510380498da Mon Sep 17 00:00:00 2001 From: joeytrasatti-openai Date: Wed, 9 Sep 2026 22:54:45 +0000 Subject: [PATCH] Add paginated thread attachment listing to the state runtime (#44330) ## What changed Add `StateRuntime::list_thread_attachments` with page sizes from 1 to 100 and stable keyset ordering by `created_at` and `id`. Return a continuation cursor when more attachments remain, and reject malformed cursors or cursors belonging to another thread. ## Testing Add coverage for thread-scoped pagination, cursor validation, and attachment persistence across reopening the database and archiving/unarchiving a thread. Verify listings reflect attachment removal and cascading deletion when a thread is deleted. GitOrigin-RevId: e07433e1a675669082cc58f86b0dc4535d365ba3 --- codex-rs/state/src/lib.rs | 3 + .../state/src/runtime/thread_attachments.rs | 78 ++++++++++ .../src/runtime/thread_attachments_tests.rs | 142 +++++++++++++++++- 3 files changed, 222 insertions(+), 1 deletion(-) diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index 5c4358ee34..9e5b341686 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -108,6 +108,9 @@ 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 attachments returned in one page. +pub const MAX_THREAD_ATTACHMENT_LIST_PAGE_SIZE: usize = 100; + /// Maximum number of active attachments retained for one thread. pub const MAX_THREAD_ATTACHMENTS_PER_THREAD: usize = 100; diff --git a/codex-rs/state/src/runtime/thread_attachments.rs b/codex-rs/state/src/runtime/thread_attachments.rs index abe2c34423..7e6cfddcb0 100644 --- a/codex-rs/state/src/runtime/thread_attachments.rs +++ b/codex-rs/state/src/runtime/thread_attachments.rs @@ -3,16 +3,20 @@ use super::StateRuntime; use crate::AddThreadAttachmentOutcome; use crate::MAX_THREAD_ATTACHMENT_IDENTITY_KEY_BYTES; +use crate::MAX_THREAD_ATTACHMENT_LIST_PAGE_SIZE; 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 crate::ThreadAttachmentPage; use anyhow::Context; use chrono::Utc; use codex_protocol::ThreadId; use serde_json::Value; +use sqlx::QueryBuilder; use sqlx::Row; +use sqlx::Sqlite; use sqlx::sqlite::SqliteRow; use uuid::Uuid; @@ -129,6 +133,61 @@ impl StateRuntime { transaction.commit().await?; Ok(outcome) } + + /// List one bounded page of attachments for one thread in stable keyset order. + pub async fn list_thread_attachments( + &self, + thread_id: ThreadId, + cursor: Option<&str>, + limit: usize, + ) -> anyhow::Result { + if !(1..=MAX_THREAD_ATTACHMENT_LIST_PAGE_SIZE).contains(&limit) { + anyhow::bail!( + "invalid thread attachment request: page limit must be between 1 and {MAX_THREAD_ATTACHMENT_LIST_PAGE_SIZE}" + ); + } + + let thread_id_string = thread_id.to_string(); + let anchor = cursor.map(parse_attachment_cursor).transpose()?; + if let Some((cursor_thread_id, _, _)) = anchor.as_ref() + && cursor_thread_id != &thread_id_string + { + anyhow::bail!("invalid thread attachment request: invalid pagination cursor"); + } + let mut query = QueryBuilder::::new( + "SELECT id, thread_id, attachment_type, identity_key, payload, created_at FROM thread_attachments WHERE thread_id = ", + ); + query.push_bind(thread_id_string); + if let Some((_, created_at, attachment_id)) = anchor { + query.push(" AND (created_at, id) > ("); + query.push_bind(created_at); + query.push(", "); + query.push_bind(attachment_id); + query.push(")"); + } + query.push(" ORDER BY created_at ASC, id ASC LIMIT "); + query.push_bind(i64::try_from(limit + 1)?); + let rows = query.build().fetch_all(self.pool.as_ref()).await?; + let mut attachments = rows + .iter() + .map(attachment_from_row) + .collect::>>()?; + let next_cursor = if attachments.len() > limit { + attachments.pop(); + attachments.last().map(|attachment| { + format!( + "{}|{}|{}", + attachment.thread_id, attachment.created_at, attachment.id + ) + }) + } else { + None + }; + Ok(ThreadAttachmentPage { + attachments, + next_cursor, + }) + } } fn validate_attachment_identity(attachment_type: &str, identity_key: &str) -> anyhow::Result<()> { @@ -153,6 +212,25 @@ fn validate_attachment_identity(attachment_type: &str, identity_key: &str) -> an Ok(()) } +fn parse_attachment_cursor(cursor: &str) -> anyhow::Result<(String, i64, String)> { + let mut segments = cursor.split('|'); + let (Some(thread_id), Some(created_at), Some(attachment_id), None) = ( + segments.next(), + segments.next(), + segments.next(), + segments.next(), + ) else { + anyhow::bail!("invalid thread attachment request: invalid pagination cursor"); + }; + if ThreadId::from_string(thread_id).is_err() || Uuid::parse_str(attachment_id).is_err() { + anyhow::bail!("invalid thread attachment request: invalid pagination cursor"); + } + let created_at = created_at + .parse::() + .context("invalid thread attachment request: invalid pagination cursor")?; + Ok((thread_id.to_string(), created_at, attachment_id.to_string())) +} + fn attachment_from_row(row: &SqliteRow) -> anyhow::Result { let thread_id: String = row.try_get("thread_id")?; let payload: String = row.try_get("payload")?; diff --git a/codex-rs/state/src/runtime/thread_attachments_tests.rs b/codex-rs/state/src/runtime/thread_attachments_tests.rs index 24afdf5c00..7bbcb6be84 100644 --- a/codex-rs/state/src/runtime/thread_attachments_tests.rs +++ b/codex-rs/state/src/runtime/thread_attachments_tests.rs @@ -10,6 +10,7 @@ use crate::RemoveThreadAttachmentOutcome; use crate::runtime::test_support::test_thread_metadata; use crate::runtime::test_support::unique_temp_dir; use anyhow::Result; +use chrono::Utc; use codex_protocol::ThreadId; use codex_utils_absolute_path::test_support::PathExt; use pretty_assertions::assert_eq; @@ -80,6 +81,13 @@ async fn attachment_attachments_are_idempotent_and_scoped_to_their_thread() -> R anyhow::bail!("the same identity on another thread should create its own attachment"); }; assert_ne!(first.id, other.id); + assert_eq!( + runtime + .list_thread_attachments(thread_ids[0], /*cursor*/ None, /*limit*/ 10) + .await? + .attachments, + vec![first] + ); Ok(()) } @@ -112,6 +120,68 @@ async fn attachment_removals_return_not_found_or_the_removed_record() -> Result< .await?, RemoveThreadAttachmentOutcome::Removed(explicit) ); + assert!( + runtime + .list_thread_attachments(thread_id, /*cursor*/ None, /*limit*/ 10) + .await? + .attachments + .is_empty() + ); + Ok(()) +} + +#[tokio::test] +async fn attachment_listing_scopes_pagination_to_one_thread() -> Result<()> { + let (runtime, _codex_home, thread_ids) = runtime_with_threads(/*count*/ 2).await?; + let mut expected = Vec::new(); + for (thread_id, identity_key) in [ + (thread_ids[0], "first"), + (thread_ids[0], "second"), + (thread_ids[0], "third"), + (thread_ids[1], "excluded"), + ] { + let created = runtime + .add_thread_attachment( + thread_id, + "pull_request", + identity_key, + &json!({ "identity": identity_key }), + ) + .await?; + if thread_id == thread_ids[0] { + let AddThreadAttachmentOutcome::Created(attachment) = created else { + anyhow::bail!("new identities should create attachments"); + }; + expected.push(attachment); + } + } + + let first_page = runtime + .list_thread_attachments(thread_ids[0], /*cursor*/ None, /*limit*/ 2) + .await?; + assert_eq!(first_page.attachments, expected[..2]); + let other_thread_error = runtime + .list_thread_attachments( + thread_ids[1], + first_page.next_cursor.as_deref(), + /*limit*/ 2, + ) + .await + .expect_err("a cursor from another thread must be rejected"); + assert!( + other_thread_error + .to_string() + .contains("invalid pagination cursor") + ); + let second_page = runtime + .list_thread_attachments( + thread_ids[0], + first_page.next_cursor.as_deref(), + /*limit*/ 2, + ) + .await?; + assert_eq!(second_page.attachments, expected[2..]); + assert_eq!(second_page.next_cursor, None); Ok(()) } @@ -172,7 +242,60 @@ async fn active_attachment_limit_is_freed_by_removal() -> Result<()> { } #[tokio::test] -async fn attachment_mutations_reject_invalid_identity_payload_and_unknown_threads() -> Result<()> { +async fn attachments_survive_restart_and_archive_but_cascade_on_thread_deletion() -> Result<()> { + let (runtime, codex_home, thread_ids) = runtime_with_threads(/*count*/ 1).await?; + let thread_id = thread_ids[0]; + let AddThreadAttachmentOutcome::Created(attachment) = runtime + .add_thread_attachment( + thread_id, + "pull_request", + "openai/codex#123", + &json!({ "url": "https://github.com/openai/codex/pull/123" }), + ) + .await? + else { + anyhow::bail!("first attachment should create an attachment"); + }; + + let reopened = StateRuntime::init( + crate::SqliteConfig::new_for_testing(codex_home.as_path().abs()), + "test-provider".to_string(), + ) + .await?; + assert_eq!( + reopened + .list_thread_attachments(thread_id, /*cursor*/ None, /*limit*/ 10) + .await? + .attachments, + vec![attachment.clone()] + ); + + let rollout_path = codex_home.join("archived.jsonl"); + reopened + .mark_archived(thread_id, &rollout_path, Utc::now()) + .await?; + reopened.mark_unarchived(thread_id, &rollout_path).await?; + assert_eq!( + reopened + .list_thread_attachments(thread_id, /*cursor*/ None, /*limit*/ 10) + .await? + .attachments, + vec![attachment] + ); + + assert_eq!(reopened.delete_thread(thread_id).await?, 1); + assert!( + reopened + .list_thread_attachments(thread_id, /*cursor*/ None, /*limit*/ 10) + .await? + .attachments + .is_empty() + ); + Ok(()) +} + +#[tokio::test] +async fn attachment_requests_reject_invalid_identity_payload_and_cursor() -> 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 [ @@ -210,6 +333,15 @@ async fn attachment_mutations_reject_invalid_identity_payload_and_unknown_thread .await .expect_err("oversized attachment payload must be rejected"); assert!(too_large.to_string().contains("payload exceeds")); + let invalid_cursor = runtime + .list_thread_attachments(thread_id, Some("not-a-cursor"), /*limit*/ 10) + .await + .expect_err("malformed cursor must be rejected"); + assert!( + invalid_cursor + .to_string() + .contains("invalid pagination cursor") + ); let missing = ThreadId::new(); let missing_error = runtime .add_thread_attachment(missing, "pull_request", "pr", &json!({})) @@ -246,5 +378,13 @@ async fn concurrent_attachment_attachments_preserve_one_deterministic_identity() } } assert_eq!(created, 1); + assert_eq!( + runtime + .list_thread_attachments(thread_id, /*cursor*/ None, /*limit*/ 10) + .await? + .attachments + .len(), + 1 + ); Ok(()) }