mirror of
https://github.com/openai/codex.git
synced 2026-09-10 20:26:47 +00:00
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:
committed by
copyberry
parent
2808a9c348
commit
130d6e4fba
@@ -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;
|
||||
|
||||
|
||||
@@ -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")?;
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user