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
This commit is contained in:
joeytrasatti-openai
2026-09-09 22:54:45 +00:00
committed by copyberry
parent 2808a9c348
commit 130d6e4fba
3 changed files with 222 additions and 1 deletions

View File

@@ -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;

View File

@@ -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<ThreadAttachmentPage> {
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::<Sqlite>::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::<anyhow::Result<Vec<_>>>()?;
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::<i64>()
.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<ThreadAttachment> {
let thread_id: String = row.try_get("thread_id")?;
let payload: String = row.try_get("payload")?;

View File

@@ -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(())
}