Add paginated reverse lookup for thread attachments (#50083)

## What changed

Add `StateRuntime::list_thread_attachment_threads` to find threads by exact `attachment_type` and `identity_key`, including threads without previews. Return thread IDs and archive state, with filters for all, archived, or non-archived threads.

Order results by thread ID and bind pagination cursors to the attachment identity and archive filter. Add a database index on `(attachment_type, identity_key, thread_id)` to support the lookup.

## Testing

Add coverage for exact identity matching, archive filters, pagination, rejection of cursors reused with different lookup parameters, and membership updates after attachment removal or thread deletion.

GitOrigin-RevId: 21f32eb0cb1845a82d740d67b374080b6e12eb03

codex-rs/state/migrations/0059_thread_attachment_reverse_lookup.sql

@@ -0,0 +1,2 @@
+CREATE INDEX idx_thread_attachments_identity_thread
+    ON thread_attachments(attachment_type, identity_key, thread_id);

codex-rs/state/src/lib.rs

@@ -62,6 +62,9 @@ pub use model::Stage1JobClaimOutcome;
 pub use model::Stage1Output;
 pub use model::Stage1StartupClaimParams;
 pub use model::ThreadAttachment;
+pub use model::ThreadAttachmentArchiveFilter;
+pub use model::ThreadAttachmentOwner;
+pub use model::ThreadAttachmentOwnerPage;
 pub use model::ThreadAttachmentPage;
 pub use model::ThreadGoal;
 pub use model::ThreadGoalStatus;

codex-rs/state/src/model/mod.rs

@@ -32,6 +32,9 @@ pub use rollout_migration_state::RolloutMigrationState;
 pub use thread_attachment::AddThreadAttachmentOutcome;
 pub use thread_attachment::RemoveThreadAttachmentOutcome;
 pub use thread_attachment::ThreadAttachment;
+pub use thread_attachment::ThreadAttachmentArchiveFilter;
+pub use thread_attachment::ThreadAttachmentOwner;
+pub use thread_attachment::ThreadAttachmentOwnerPage;
 pub use thread_attachment::ThreadAttachmentPage;
 pub use thread_goal::ThreadGoal;
 pub use thread_goal::ThreadGoalStatus;

codex-rs/state/src/model/thread_attachment.rs

@@ -46,3 +46,25 @@ pub struct ThreadAttachmentPage {
     /// Opaque cursor for the next page, or `None` when the selection is exhausted.
     pub next_cursor: Option<String>,
 }
+
+/// Selects attachment owners by their persisted archive state.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ThreadAttachmentArchiveFilter {
+    All,
+    NonArchived,
+    Archived,
+}
+
+/// A thread whose current attachment membership contains the requested resource.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ThreadAttachmentOwner {
+    pub thread_id: ThreadId,
+    pub archived: bool,
+}
+
+/// One page of attachment owners, ordered by thread ID.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ThreadAttachmentOwnerPage {
+    pub threads: Vec<ThreadAttachmentOwner>,
+    pub next_cursor: Option<String>,
+}

codex-rs/state/src/runtime/thread_attachments.rs

@@ -9,10 +9,15 @@ use crate::MAX_THREAD_ATTACHMENT_TYPE_BYTES;
 use crate::MAX_THREAD_ATTACHMENTS_PER_THREAD;
 use crate::RemoveThreadAttachmentOutcome;
 use crate::ThreadAttachment;
+use crate::ThreadAttachmentArchiveFilter;
+use crate::ThreadAttachmentOwner;
+use crate::ThreadAttachmentOwnerPage;
 use crate::ThreadAttachmentPage;
 use anyhow::Context;
 use chrono::Utc;
 use codex_protocol::ThreadId;
+use serde::Deserialize;
+use serde::Serialize;
 use serde_json::Value;
 use sqlx::QueryBuilder;
 use sqlx::Row;
@@ -174,6 +179,89 @@ impl StateRuntime {
         Ok(outcome)
     }
 
+    /// Lists current owners of an exact attachment identity, including threads without previews.
+    /// `archive_filter` selects archived, non-archived, or all owners.
+    pub async fn list_thread_attachment_threads(
+        &self,
+        attachment_type: &str,
+        identity_key: &str,
+        archive_filter: ThreadAttachmentArchiveFilter,
+        cursor: Option<&str>,
+        limit: usize,
+    ) -> anyhow::Result<ThreadAttachmentOwnerPage> {
+        validate_attachment_identity(attachment_type, identity_key)?;
+        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 archived = match archive_filter {
+            ThreadAttachmentArchiveFilter::All => None,
+            ThreadAttachmentArchiveFilter::NonArchived => Some(false),
+            ThreadAttachmentArchiveFilter::Archived => Some(true),
+        };
+        let anchor = cursor
+            .map(serde_json::from_str::<AttachmentThreadsCursor>)
+            .transpose()
+            .context("invalid thread attachment request: invalid pagination cursor")?;
+        if let Some(anchor) = &anchor
+            && (anchor.attachment_type != attachment_type
+                || anchor.identity_key != identity_key
+                || anchor.archived != archived)
+        {
+            anyhow::bail!("invalid thread attachment request: invalid pagination cursor");
+        }
+        let mut query = QueryBuilder::<Sqlite>::new(
+            "SELECT thread_attachments.thread_id, threads.archived FROM thread_attachments \
+             JOIN threads ON threads.id = thread_attachments.thread_id \
+             WHERE thread_attachments.attachment_type = ",
+        );
+        query.push_bind(attachment_type);
+        query.push(" AND thread_attachments.identity_key = ");
+        query.push_bind(identity_key);
+        if let Some(archived) = archived {
+            query.push(" AND threads.archived = ");
+            query.push_bind(archived);
+        }
+        if let Some(anchor) = anchor {
+            query.push(" AND thread_attachments.thread_id > ");
+            query.push_bind(anchor.thread_id.to_string());
+        }
+        query.push(" ORDER BY thread_attachments.thread_id ASC LIMIT ");
+        query.push_bind(i64::try_from(limit + 1)?);
+        let rows = query.build().fetch_all(self.pool.as_ref()).await?;
+        let mut threads = rows
+            .iter()
+            .map(|row| -> anyhow::Result<ThreadAttachmentOwner> {
+                Ok(ThreadAttachmentOwner {
+                    thread_id: ThreadId::from_string(row.try_get("thread_id")?)
+                        .context("invalid persisted thread attachment owner")?,
+                    archived: row.try_get("archived")?,
+                })
+            })
+            .collect::<anyhow::Result<Vec<_>>>()?;
+        let next_cursor = if threads.len() > limit {
+            threads.pop();
+            threads
+                .last()
+                .map(|thread| {
+                    serde_json::to_string(&AttachmentThreadsCursor {
+                        attachment_type: attachment_type.to_owned(),
+                        identity_key: identity_key.to_owned(),
+                        archived,
+                        thread_id: thread.thread_id,
+                    })
+                })
+                .transpose()?
+        } else {
+            None
+        };
+        Ok(ThreadAttachmentOwnerPage {
+            threads,
+            next_cursor,
+        })
+    }
+
     /// List one bounded page of attachments for one thread in stable keyset order.
     pub async fn list_thread_attachments(
         &self,
@@ -230,6 +318,14 @@ impl StateRuntime {
     }
 }
 
+#[derive(Serialize, Deserialize)]
+struct AttachmentThreadsCursor {
+    attachment_type: String,
+    identity_key: String,
+    archived: Option<bool>,
+    thread_id: ThreadId,
+}
+
 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");

codex-rs/state/src/runtime/thread_attachments_tests.rs

@@ -7,6 +7,9 @@ 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::ThreadAttachmentArchiveFilter;
+use crate::ThreadAttachmentOwner;
+use crate::ThreadAttachmentOwnerPage;
 use crate::runtime::test_support::test_thread_metadata;
 use crate::runtime::test_support::unique_temp_dir;
 use anyhow::Result;
@@ -454,3 +457,132 @@ async fn concurrent_attachment_attachments_preserve_one_deterministic_identity()
     );
     Ok(())
 }
+
+#[tokio::test]
+async fn attachment_reverse_lookup_filters_pages_and_tracks_membership() -> Result<()> {
+    let (runtime, codex_home, mut thread_ids) = runtime_with_threads(/*count*/ 3).await?;
+    thread_ids.sort_by_key(ThreadId::to_string);
+    // Reference-backed forks may have attachments before their own first user message.
+    sqlx::query("UPDATE threads SET preview = '', first_user_message = ''")
+        .execute(runtime.pool.as_ref())
+        .await?;
+    for thread_id in &thread_ids[..2] {
+        runtime
+            .add_thread_attachment(*thread_id, "worktree", "repo|branch", &json!({}))
+            .await?;
+    }
+    runtime
+        .add_thread_attachment(thread_ids[2], "pull_request", "repo|branch", &json!({}))
+        .await?;
+    runtime
+        .add_thread_attachment(thread_ids[2], "worktree", "other", &json!({}))
+        .await?;
+    runtime
+        .mark_archived(
+            thread_ids[0],
+            &codex_home.join("archived.jsonl"),
+            Utc::now(),
+        )
+        .await?;
+    let archived_thread = ThreadAttachmentOwner {
+        thread_id: thread_ids[0],
+        archived: true,
+    };
+    let active_thread = ThreadAttachmentOwner {
+        thread_id: thread_ids[1],
+        archived: false,
+    };
+    let first_page = runtime
+        .list_thread_attachment_threads(
+            "worktree",
+            "repo|branch",
+            ThreadAttachmentArchiveFilter::All,
+            /*cursor*/ None,
+            /*limit*/ 1,
+        )
+        .await?;
+    assert_eq!(first_page.threads, vec![archived_thread.clone()]);
+    let cursor = first_page
+        .next_cursor
+        .expect("second owner requires continuation");
+    assert_eq!(
+        runtime
+            .list_thread_attachment_threads(
+                "worktree",
+                "repo|branch",
+                ThreadAttachmentArchiveFilter::All,
+                Some(&cursor),
+                /*limit*/ 1
+            )
+            .await?,
+        ThreadAttachmentOwnerPage {
+            threads: vec![active_thread.clone()],
+            next_cursor: None
+        }
+    );
+    for (attachment_type, identity_key, archive_filter) in [
+        (
+            "pull_request",
+            "repo|branch",
+            ThreadAttachmentArchiveFilter::All,
+        ),
+        ("worktree", "other", ThreadAttachmentArchiveFilter::All),
+        (
+            "worktree",
+            "repo|branch",
+            ThreadAttachmentArchiveFilter::NonArchived,
+        ),
+    ] {
+        let error = runtime
+            .list_thread_attachment_threads(
+                attachment_type,
+                identity_key,
+                archive_filter,
+                Some(&cursor),
+                /*limit*/ 1,
+            )
+            .await
+            .expect_err("cursor is bound to the complete lookup");
+        assert!(error.to_string().contains("invalid pagination cursor"));
+    }
+    for (archive_filter, expected) in [
+        (ThreadAttachmentArchiveFilter::Archived, archived_thread),
+        (ThreadAttachmentArchiveFilter::NonArchived, active_thread),
+    ] {
+        assert_eq!(
+            runtime
+                .list_thread_attachment_threads(
+                    "worktree",
+                    "repo|branch",
+                    archive_filter,
+                    /*cursor*/ None,
+                    /*limit*/ 1
+                )
+                .await?,
+            ThreadAttachmentOwnerPage {
+                threads: vec![expected],
+                next_cursor: None
+            }
+        );
+    }
+    runtime
+        .remove_thread_attachment(thread_ids[0], "worktree", "repo|branch")
+        .await?;
+    runtime.delete_thread(thread_ids[1]).await?;
+    assert_eq!(
+        runtime
+            .list_thread_attachment_threads(
+                "worktree",
+                "repo|branch",
+                ThreadAttachmentArchiveFilter::All,
+                /*cursor*/ None,
+                /*limit*/ 10
+            )
+            .await?,
+        ThreadAttachmentOwnerPage {
+            threads: Vec::new(),
+            next_cursor: None
+        }
+    );
+    Ok(())
+}