Add transactional thread attachment mutations to the state runtime (#43949)

## What changed

- Replace the `ThreadArtifact` model and related exports with attachment terminology.
- Add `StateRuntime::add_thread_attachment` and `remove_thread_attachment` using SQLite transactions. Repeated additions for the same thread, attachment type, and identity key return the existing record without changing its payload or creation time.
- Enforce limits of 100 attachments per thread, 64 KiB per serialized payload, and 256 bytes each for nonblank attachment types and identity keys. Removal returns the deleted record or `NotFound` and frees capacity immediately. Both mutations reject unknown threads.

## Testing

Add tests for idempotency, thread isolation, removal outcomes, capacity reuse, invalid inputs, unknown threads, and concurrent additions creating exactly one record.

GitOrigin-RevId: a68fe5076832f071524f1b4fc87d08b96ac270b7
This commit is contained in:
joeytrasatti-openai
2026-09-08 23:55:24 +00:00
committed by copyberry
parent 5b682c9875
commit 589874be81
7 changed files with 493 additions and 55 deletions

View File

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

View File

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

View File

@@ -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<ThreadArtifact>,
/// Opaque cursor for the next page, or `None` when the selection is exhausted.
pub next_cursor: Option<String>,
}

View File

@@ -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<ThreadAttachment>,
/// Opaque cursor for the next page, or `None` when the selection is exhausted.
pub next_cursor: Option<String>,
}

View File

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

View File

@@ -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<AddThreadAttachmentOutcome> {
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<RemoveThreadAttachmentOutcome> {
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<ThreadAttachment> {
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;

View File

@@ -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<StateRuntime>, PathBuf, Vec<ThreadId>)> {
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(())
}