From c3b0c3d3189f56341033274bf82af5de6bd9c03e Mon Sep 17 00:00:00 2001 From: "dispatch-developer[bot]" <306909890+dispatch-developer[bot]@users.noreply.github.com> Date: Mon, 17 Aug 2026 07:29:28 +0000 Subject: [PATCH 1/2] replay: add bounded agent turn projection --- .env.example | 4 + README.md | 47 +++++++- src/api.rs | 35 +++++- src/auth.rs | 124 ++++++++++++++++++- src/capability.rs | 4 + src/config.rs | 45 +++++++ src/error.rs | 8 ++ src/model.rs | 8 ++ src/store.rs | 298 +++++++++++++++++++++++++++++++++++++++++++++- 9 files changed, 560 insertions(+), 13 deletions(-) diff --git a/.env.example b/.env.example index d5310fa..875e4c2 100644 --- a/.env.example +++ b/.env.example @@ -9,6 +9,10 @@ AUTH_MODE=trusted_headers # AUTH_AUDIENCE=threadmark-api # AUTH_JWKS_URL=https://auth.example.com/.well-known/jwks.json AUTH_MAX_OWNER_TOKEN_SECONDS=300 +AUTH_MAX_DELEGATED_TOKEN_SECONDS=600 +AGENT_REPLAY_MAX_ITEMS=200 +AGENT_REPLAY_MAX_BYTES=1048576 +AGENT_REPLAY_STRIP_TOP_LEVEL_FIELDS=id FILE_MAX_MB=32 S3_ENDPOINT=http://localhost:9010 S3_PUBLIC_URL=http://localhost:9010 diff --git a/README.md b/README.md index 5c4ef44..6f741f5 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,7 @@ This is an experiment and its API is not stable. The current slice establishes: - Idempotent item batches and turn creation. - Cursor-based item reads. - Open Responses replay projection with optional top-level `id` removal. +- Snapshot-consistent, size-bounded text replay for delegated agent turns. - Agent-scoped continuation records and optional private checkpoint state. Editing, branching, retention, event delivery, production authentication, and @@ -140,13 +141,54 @@ mode. Production uses `AUTH_MODE=jwt` with `AUTH_ISSUER`, `AUTH_AUDIENCE`, and an HTTPS `AUTH_JWKS_URL`. JWT mode accepts Ed25519 `at+jwt` owner-session tokens and derives tenant, principal, and endpoint permissions exclusively from verified -claims. Delegated-agent tokens remain rejected until their resource-bound write -invariants are implemented. +claims. It also accepts delegated-agent tokens only for the agent replay +operation described below. Delegated writes remain disabled. An agent called by Parley can receive a short-lived token scoped to the same tenant, principal, conversation, turn, and agent deployment. That authorization layer is deliberately separate from the ledger model. +### Bounded agent replay + +`POST /v1/conversations/{conversation_id}/turns/{turn_id}/agent-replay` is the +initial Bonsai replay integration. It accepts no request body and requires a +`delegated_agent` JWT with `transcript:read` and exact `tenant`, `principal`, +`conversation_id`, `turn_id`, and `agent_ref` bounds. Wrong actor or resource +bounds return a non-enumerating error. The turn must have been created by +`POST /v1/turn-starts`; its recorded `last_seq` is the immutable replay cursor. + +The operation opens a PostgreSQL repeatable-read, read-only transaction before +resolving ownership, the atomic-start boundary, and ordered items. It verifies +the complete triggering batch still exists in that snapshot. A concurrent +truncate is therefore observed wholly before or wholly after its commit; a +snapshot missing the boundary returns `replay_snapshot_unavailable` and never +returns a cursor for absent turn-start input. + +The first integration supports only these historical message shapes: + +- `type=message`, `role=user`, with a nonempty array of `input_text` parts + containing only string `text` plus the `type` discriminator; +- `type=message`, `role=assistant`, with a nonempty array of `output_text` parts + containing string `text`, the `type` discriminator, and optional + `annotations`. + +Other item types, roles, non-text parts, mixed content, and media return +`unsupported_agent_replay_item`. In particular, file and image parts are never +forwarded, and any canonical file URI anywhere in an item is rejected, so +unresolved `threadmark://` resources cannot reach a model provider through this +endpoint. Accepted top-level item fields are otherwise preserved. Only +top-level fields named by `AGENT_REPLAY_STRIP_TOP_LEVEL_FIELDS` (comma-separated, +default `id`) are removed; nested fields are untouched. + +`AGENT_REPLAY_MAX_ITEMS` (default `200`) and `AGENT_REPLAY_MAX_BYTES` (default +`1048576`) are hard inclusive limits. The byte limit is the exact compact JSON +serialization of the returned `input` array after configured field removal. +Exceeding either limit returns HTTP `413` with +`error.code=context_limit_exceeded` before a projection is returned. + +The existing owner endpoint, `POST /v1/conversations/{id}/replay`, is unchanged: +it remains an opaque, multimodal projection with caller-selected file delivery. + ## Example flow Create a conversation: @@ -236,6 +278,7 @@ GET /v1/continuations/resp_abc?agent_ref=research-agent%2Fprod | `GET` | `/v1/conversations/{id}/items` | Read ordered items after a sequence cursor | | `POST` | `/v1/conversations/{id}/items` | Atomically append an idempotent item batch | | `POST` | `/v1/conversations/{id}/replay` | Build an Open Responses input array | +| `POST` | `/v1/conversations/{conversation_id}/turns/{turn_id}/agent-replay` | Build bounded text input for a delegated agent turn | | `POST` | `/v1/conversations/{id}/turns` | Create an idempotent turn | | `PATCH` | `/v1/turns/{id}` | Update turn state and outcome | | `POST` | `/v1/conversations/{id}/continuations` | Record an agent checkpoint | diff --git a/src/api.rs b/src/api.rs index 595fe50..aec61cc 100644 --- a/src/api.rs +++ b/src/api.rs @@ -17,11 +17,12 @@ use crate::{ error::{ApiError, ApiResult}, files, model::{ - Actor, AppendItems, AppendResult, Continuation, ContinuationQuery, Conversation, - CreateContinuation, CreateConversation, CreateDownload, CreateTurn, DownloadDelivery, - DownloadGrant, FileResponse, Item, ListConversationsQuery, ListItemsQuery, - RegenerateResult, ReplayRequest, ReplayResult, StartTurn, StartTurnResult, StrictJson, - TruncateConversation, Turn, UpdateConversation, UpdateTurn, validate_json_number_tokens, + Actor, AgentReplayResult, AppendItems, AppendResult, Continuation, ContinuationQuery, + Conversation, CreateContinuation, CreateConversation, CreateDownload, CreateTurn, + DownloadDelivery, DownloadGrant, FileResponse, Item, ListConversationsQuery, + ListItemsQuery, RegenerateResult, ReplayRequest, ReplayResult, StartTurn, StartTurnResult, + StrictJson, TruncateConversation, Turn, UpdateConversation, UpdateTurn, + validate_json_number_tokens, }, object_store::ObjectStore, store, uploads, @@ -54,6 +55,10 @@ pub fn router(state: AppState) -> Router { get(list_items).post(append_items), ) .route("/v1/conversations/{id}/replay", post(replay)) + .route( + "/v1/conversations/{conversation_id}/turns/{turn_id}/agent-replay", + post(agent_replay), + ) .route( "/v1/conversations/{id}/turns", get(list_turns).post(create_turn), @@ -233,6 +238,26 @@ async fn replay( Ok(Json(store::replay(&state, &auth, &id, request).await?)) } +async fn agent_replay( + State(state): State, + auth: AuthContext, + Path((conversation_id, turn_id)): Path<(String, String)>, + body: Bytes, +) -> ApiResult> { + auth.require(Permission::AgentReplay)?; + if !body.is_empty() { + return Err(ApiError::BadRequest( + "agent replay does not accept a request body".into(), + )); + } + let agent_ref = auth + .require_agent_replay_scope(&conversation_id, &turn_id)? + .to_owned(); + Ok(Json( + store::agent_replay(&state, &auth, &conversation_id, &turn_id, &agent_ref).await?, + )) +} + async fn create_turn( State(state): State, auth: AuthContext, diff --git a/src/auth.rs b/src/auth.rs index 48a24b8..7a18fac 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -32,6 +32,7 @@ struct JwtVerifier { issuer: String, audience: String, max_owner_seconds: u64, + max_delegated_seconds: u64, keys: HashMap, } @@ -40,9 +41,18 @@ pub struct AuthContext { pub actor: Actor, pub client_id: String, agent_ref: Option, + conversation_id: Option, + turn_id: Option, + token_kind: TokenKind, permissions: HashSet, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum TokenKind { + OwnerSession, + DelegatedAgent, +} + #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] pub enum Permission { ConversationList, @@ -53,6 +63,7 @@ pub enum Permission { ConversationTruncate, ConversationRegenerate, TranscriptRead, + AgentReplay, TranscriptAppend, TurnCreate, TurnRead, @@ -102,6 +113,8 @@ struct Claims { principal: String, permissions: Vec, agent_ref: Option, + conversation_id: Option, + turn_id: Option, } #[derive(Debug, Deserialize)] @@ -146,6 +159,7 @@ impl Authenticator { issuer, audience, config.auth_max_owner_token_seconds, + config.auth_max_delegated_token_seconds, jwks, )?) } @@ -167,6 +181,7 @@ impl JwtVerifier { issuer: String, audience: String, max_owner_seconds: u64, + max_delegated_seconds: u64, jwks: JwkSet, ) -> anyhow::Result { ensure!(!jwks.keys.is_empty(), "JWKS contains no keys"); @@ -191,6 +206,7 @@ impl JwtVerifier { issuer, audience, max_owner_seconds, + max_delegated_seconds, keys, }) } @@ -219,17 +235,25 @@ impl JwtVerifier { impl Claims { fn context(self, verifier: &JwtVerifier) -> Option { let now = SystemTime::now().duration_since(UNIX_EPOCH).ok()?.as_secs(); + let token_kind = match self.token_kind.as_str() { + "owner_session" => TokenKind::OwnerSession, + "delegated_agent" => TokenKind::DelegatedAgent, + _ => return None, + }; let valid_audience = match &self.aud { Audience::One(value) => value == &verifier.audience, Audience::Many(values) => values.len() == 1 && values[0] == verifier.audience, }; if self.iss != verifier.issuer || !valid_audience - || self.token_kind != "owner_session" || self.exp <= self.iat || self.exp <= self.nbf || self.iat > now.saturating_add(30) - || self.exp.saturating_sub(self.iat) > verifier.max_owner_seconds + || self.exp.saturating_sub(self.iat) + > match token_kind { + TokenKind::OwnerSession => verifier.max_owner_seconds, + TokenKind::DelegatedAgent => verifier.max_delegated_seconds, + } || !valid_id(&self.sub) || !valid_id(&self.client_id) || self.client_id == "threadmark:trusted-headers" @@ -240,14 +264,34 @@ impl Claims { .agent_ref .as_deref() .is_some_and(|value| !valid_id(value)) + || self + .conversation_id + .as_deref() + .is_some_and(|value| !valid_id(value)) + || self + .turn_id + .as_deref() + .is_some_and(|value| !valid_id(value)) { return None; } - let permissions = self + let mut permissions = self .permissions .iter() .map(|value| Permission::parse(value)) .collect::>>()?; + if token_kind == TokenKind::DelegatedAgent + && (self.conversation_id.is_none() + || self.turn_id.is_none() + || self.agent_ref.is_none() + || permissions.len() != 1 + || !permissions.contains(&Permission::TranscriptRead)) + { + return None; + } + if token_kind == TokenKind::DelegatedAgent { + permissions = [Permission::AgentReplay].into_iter().collect(); + } Some(AuthContext { actor: Actor { tenant_id: self.tenant, @@ -255,6 +299,9 @@ impl Claims { }, client_id: self.client_id, agent_ref: self.agent_ref, + conversation_id: self.conversation_id, + turn_id: self.turn_id, + token_kind, permissions, }) } @@ -280,6 +327,22 @@ impl AuthContext { _ => Ok(()), } } + + pub fn require_agent_replay_scope( + &self, + conversation_id: &str, + turn_id: &str, + ) -> Result<&str, ApiError> { + if self.token_kind != TokenKind::DelegatedAgent + || self.conversation_id.as_deref() != Some(conversation_id) + || self.turn_id.as_deref() != Some(turn_id) + { + return Err(ApiError::NotFound("Agent replay not found.".into())); + } + self.agent_ref + .as_deref() + .ok_or_else(|| ApiError::NotFound("Agent replay not found.".into())) + } } impl std::ops::Deref for AuthContext { @@ -353,6 +416,9 @@ fn trusted_headers(headers: &HeaderMap) -> Result { }, client_id: "threadmark:trusted-headers".into(), agent_ref: None, + conversation_id: None, + turn_id: None, + token_kind: TokenKind::OwnerSession, permissions: OWNER_PERMISSIONS.into_iter().collect(), }) } @@ -385,6 +451,7 @@ mod tests { "https://issuer.example".into(), "threadmark-api".into(), 300, + 600, jwks, ) .unwrap() @@ -494,6 +561,57 @@ mod tests { )); } + #[test] + fn delegated_replay_token_enforces_all_resource_bounds() { + let mut claims = claims(); + claims["token_kind"] = json!("delegated_agent"); + claims["permissions"] = json!(["transcript:read"]); + claims["conversation_id"] = json!("conv_1"); + claims["turn_id"] = json!("turn_1"); + claims["agent_ref"] = json!("bonsai/prod"); + let context = verifier().authenticate(&headers(&claims)).unwrap(); + assert_eq!( + context + .require_agent_replay_scope("conv_1", "turn_1") + .unwrap(), + "bonsai/prod" + ); + assert!(context.require(Permission::AgentReplay).is_ok()); + assert!(matches!( + context.require(Permission::TranscriptRead), + Err(ApiError::Forbidden) + )); + assert!(matches!( + context.require_agent_replay_scope("conv_other", "turn_1"), + Err(ApiError::NotFound(_)) + )); + assert!(matches!( + context.require_agent_replay_scope("conv_1", "turn_other"), + Err(ApiError::NotFound(_)) + )); + } + + #[test] + fn delegated_tokens_reject_missing_scope_and_owner_permissions() { + let mut missing_turn = claims(); + missing_turn["token_kind"] = json!("delegated_agent"); + missing_turn["permissions"] = json!(["transcript:read"]); + missing_turn["conversation_id"] = json!("conv_1"); + missing_turn["agent_ref"] = json!("bonsai/prod"); + assert!(matches!( + verifier().authenticate(&headers(&missing_turn)), + Err(ApiError::Unauthorized) + )); + + let mut broad = missing_turn; + broad["turn_id"] = json!("turn_1"); + broad["permissions"] = json!(["transcript:read", "conversation:read"]); + assert!(matches!( + verifier().authenticate(&headers(&broad)), + Err(ApiError::Unauthorized) + )); + } + #[test] fn rejects_excessive_lifetime_and_multi_audience_tokens() { let mut long_lived = claims(); diff --git a/src/capability.rs b/src/capability.rs index 2a2a209..5534f89 100644 --- a/src/capability.rs +++ b/src/capability.rs @@ -120,6 +120,10 @@ mod tests { auth_audience: None, auth_jwks_url: None, auth_max_owner_token_seconds: 300, + auth_max_delegated_token_seconds: 600, + agent_replay_max_items: 200, + agent_replay_max_bytes: 1024 * 1024, + agent_replay_strip_top_level_fields: vec!["id".into()], })) } diff --git a/src/config.rs b/src/config.rs index 535bd69..5dfac82 100644 --- a/src/config.rs +++ b/src/config.rs @@ -27,6 +27,10 @@ pub struct Inner { pub auth_audience: Option, pub auth_jwks_url: Option, pub auth_max_owner_token_seconds: u64, + pub auth_max_delegated_token_seconds: u64, + pub agent_replay_max_items: usize, + pub agent_replay_max_bytes: usize, + pub agent_replay_strip_top_level_fields: Vec, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -100,6 +104,43 @@ impl Config { auth_max_owner_token_seconds > 0, "AUTH_MAX_OWNER_TOKEN_SECONDS must be greater than zero" ); + let auth_max_delegated_token_seconds = + parse("AUTH_MAX_DELEGATED_TOKEN_SECONDS", "600")?; + ensure!( + auth_max_delegated_token_seconds > 0, + "AUTH_MAX_DELEGATED_TOKEN_SECONDS must be greater than zero" + ); + let agent_replay_max_items = usize::try_from(parse("AGENT_REPLAY_MAX_ITEMS", "200")?) + .context("AGENT_REPLAY_MAX_ITEMS is too large")?; + ensure!( + agent_replay_max_items > 0, + "AGENT_REPLAY_MAX_ITEMS must be greater than zero" + ); + let agent_replay_max_bytes = usize::try_from(parse("AGENT_REPLAY_MAX_BYTES", "1048576")?) + .context("AGENT_REPLAY_MAX_BYTES is too large")?; + ensure!( + agent_replay_max_bytes > 0, + "AGENT_REPLAY_MAX_BYTES must be greater than zero" + ); + let agent_replay_strip_top_level_fields = std::env::var( + "AGENT_REPLAY_STRIP_TOP_LEVEL_FIELDS", + ) + .unwrap_or_else(|_| "id".into()) + .split(',') + .map(str::trim) + .filter(|field| !field.is_empty()) + .map(str::to_owned) + .collect::>(); + let mut unique_fields = agent_replay_strip_top_level_fields.clone(); + unique_fields.sort(); + unique_fields.dedup(); + ensure!( + unique_fields.len() == agent_replay_strip_top_level_fields.len() + && agent_replay_strip_top_level_fields + .iter() + .all(|field| field.chars().count() <= 200), + "AGENT_REPLAY_STRIP_TOP_LEVEL_FIELDS must contain unique field names of at most 200 characters" + ); let direct_upload_enabled = std::env::var("DIRECT_UPLOAD_ENABLED") .unwrap_or_else(|_| "false".into()) .parse() @@ -154,6 +195,10 @@ impl Config { auth_audience, auth_jwks_url, auth_max_owner_token_seconds, + auth_max_delegated_token_seconds, + agent_replay_max_items, + agent_replay_max_bytes, + agent_replay_strip_top_level_fields, }))) } } diff --git a/src/error.rs b/src/error.rs index 37aa137..59da4f1 100644 --- a/src/error.rs +++ b/src/error.rs @@ -13,6 +13,8 @@ pub enum ApiError { Forbidden, #[error("{0}")] BadRequest(String), + #[error("{message}")] + CodedBadRequest { code: &'static str, message: String }, #[error("{0}")] NotFound(String), #[error("{0}")] @@ -25,6 +27,8 @@ pub enum ApiError { CodedUnavailable { code: &'static str, message: String }, #[error("{0}")] PayloadTooLarge(String), + #[error("{message}")] + CodedPayloadTooLarge { code: &'static str, message: String }, #[error("object storage operation failed")] ObjectStore(#[source] anyhow::Error), #[error("database operation failed")] @@ -37,12 +41,16 @@ impl IntoResponse for ApiError { Self::Unauthorized => (StatusCode::UNAUTHORIZED, "unauthorized"), Self::Forbidden => (StatusCode::FORBIDDEN, "forbidden"), Self::BadRequest(_) => (StatusCode::BAD_REQUEST, "invalid_request"), + Self::CodedBadRequest { code, .. } => (StatusCode::BAD_REQUEST, *code), Self::NotFound(_) => (StatusCode::NOT_FOUND, "not_found"), Self::Conflict(_) => (StatusCode::CONFLICT, "conflict"), Self::CodedConflict { code, .. } => (StatusCode::CONFLICT, *code), Self::CodedGone { code, .. } => (StatusCode::GONE, *code), Self::CodedUnavailable { code, .. } => (StatusCode::SERVICE_UNAVAILABLE, *code), Self::PayloadTooLarge(_) => (StatusCode::PAYLOAD_TOO_LARGE, "payload_too_large"), + Self::CodedPayloadTooLarge { code, .. } => { + (StatusCode::PAYLOAD_TOO_LARGE, *code) + } Self::ObjectStore(error) => { tracing::error!(?error, "object storage request failed"); (StatusCode::BAD_GATEWAY, "object_store_error") diff --git a/src/model.rs b/src/model.rs index b689b3a..45e2540 100644 --- a/src/model.rs +++ b/src/model.rs @@ -336,6 +336,14 @@ pub struct ReplayResult { pub input: Vec, } +#[derive(Debug, Serialize)] +pub struct AgentReplayResult { + pub conversation_id: String, + pub turn_id: String, + pub through_seq: i64, + pub input: Vec, +} + #[derive(Debug, Serialize, FromRow)] pub struct Turn { pub id: String, diff --git a/src/store.rs b/src/store.rs index 2173340..8b8dd0a 100644 --- a/src/store.rs +++ b/src/store.rs @@ -11,9 +11,9 @@ use crate::{ files, ids::new_id, model::{ - Actor, AppendItems, AppendResult, Continuation, Conversation, CreateContinuation, - CreateConversation, CreateTurn, FileDelivery, Item, ReplayRequest, ReplayResult, StartTurn, - StartTurnResult, Turn, UpdateConversation, UpdateTurn, + Actor, AgentReplayResult, AppendItems, AppendResult, Continuation, Conversation, + CreateContinuation, CreateConversation, CreateTurn, FileDelivery, Item, ReplayRequest, + ReplayResult, StartTurn, StartTurnResult, Turn, UpdateConversation, UpdateTurn, }, }; @@ -854,6 +854,181 @@ pub async fn replay( }) } +pub async fn agent_replay( + state: &AppState, + actor: &Actor, + conversation_id: &str, + turn_id: &str, + agent_ref: &str, +) -> ApiResult { + let mut tx = state.pool.begin().await?; + // The turn boundary, conversation ownership, and selected items must be one + // observation even while an owner truncates the transcript. + sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY") + .execute(&mut *tx) + .await?; + let (first_seq, through_seq) = sqlx::query_as::<_, (i64, i64)>( + "SELECT turn_start.first_seq, turn_start.last_seq + FROM conversations conversation + JOIN turns turn ON turn.conversation_id = conversation.id + JOIN turn_starts turn_start ON turn_start.turn_id = turn.id + AND turn_start.conversation_id = conversation.id + WHERE conversation.id = $1 AND turn.id = $2 AND turn.agent_ref = $3 + AND conversation.tenant_id = $4 AND conversation.owner_ref = $5", + ) + .bind(conversation_id) + .bind(turn_id) + .bind(agent_ref) + .bind(&actor.tenant_id) + .bind(&actor.principal_id) + .fetch_optional(&mut *tx) + .await? + .ok_or_else(|| ApiError::NotFound("Agent replay not found.".into()))?; + + let (item_count, snapshot_complete) = sqlx::query_as::<_, (i64, bool)>( + "SELECT count(*)::bigint, + EXISTS(SELECT 1 FROM conversation_items + WHERE conversation_id = $1 AND seq = $2) + AND (SELECT count(*) FROM turn_start_items item + JOIN turn_starts start ON start.id = item.turn_start_id + WHERE start.turn_id = $3) = $2 - $4 + 1 + AND NOT EXISTS( + SELECT 1 FROM turn_start_items item + JOIN turn_starts start ON start.id = item.turn_start_id + LEFT JOIN conversation_items transcript + ON transcript.id = item.item_id + AND transcript.conversation_id = $1 + AND transcript.seq = item.seq + AND transcript.turn_id = $3 + WHERE start.turn_id = $3 AND transcript.id IS NULL) + FROM conversation_items + WHERE conversation_id = $1 AND seq <= $2", + ) + .bind(conversation_id) + .bind(through_seq) + .bind(turn_id) + .bind(first_seq) + .fetch_one(&mut *tx) + .await?; + if !snapshot_complete { + return Err(ApiError::CodedConflict { + code: "replay_snapshot_unavailable", + message: "The turn input boundary is not present in this transcript snapshot.".into(), + }); + } + if usize::try_from(item_count).unwrap_or(usize::MAX) > state.config.agent_replay_max_items { + return Err(context_limit("item count")); + } + let rows = sqlx::query_as::<_, (i64, Value)>( + "SELECT seq, payload FROM conversation_items + WHERE conversation_id = $1 AND seq <= $2 ORDER BY seq ASC", + ) + .bind(conversation_id) + .bind(through_seq) + .fetch_all(&mut *tx) + .await?; + let input = build_agent_projection( + rows, + &state.config.agent_replay_strip_top_level_fields, + state.config.agent_replay_max_items, + state.config.agent_replay_max_bytes, + )?; + tx.commit().await?; + Ok(AgentReplayResult { + conversation_id: conversation_id.into(), + turn_id: turn_id.into(), + through_seq, + input, + }) +} + +fn context_limit(limit: &str) -> ApiError { + ApiError::CodedPayloadTooLarge { + code: "context_limit_exceeded", + message: format!("Agent replay exceeds the configured {limit} limit."), + } +} + +fn build_agent_projection( + rows: Vec<(i64, Value)>, + strip_top_level_fields: &[String], + max_items: usize, + max_bytes: usize, +) -> ApiResult> { + if rows.len() > max_items { + return Err(context_limit("item count")); + } + let mut input = Vec::with_capacity(rows.len()); + for (_, mut item) in rows { + validate_agent_text_item(&item)?; + let object = item.as_object_mut().expect("validated message object"); + for field in strip_top_level_fields { + object.remove(field); + } + input.push(item); + } + let serialized_bytes = serde_json::to_vec(&input) + .map_err(|error| ApiError::BadRequest(format!("could not serialize replay: {error}")))? + .len(); + if serialized_bytes > max_bytes { + return Err(context_limit("serialized byte")); + } + Ok(input) +} + +fn validate_agent_text_item(item: &Value) -> ApiResult<()> { + if contains_threadmark_uri(item) { + return Err(unsupported_agent_replay_item()); + } + let object = item.as_object().ok_or_else(unsupported_agent_replay_item)?; + if object.get("type").and_then(Value::as_str) != Some("message") { + return Err(unsupported_agent_replay_item()); + } + let expected_part = match object.get("role").and_then(Value::as_str) { + Some("user") => "input_text", + Some("assistant") => "output_text", + _ => return Err(unsupported_agent_replay_item()), + }; + let content = object + .get("content") + .and_then(Value::as_array) + .filter(|content| !content.is_empty()) + .ok_or_else(unsupported_agent_replay_item)?; + if content.iter().any(|part| { + let Some(part) = part.as_object() else { + return true; + }; + let allowed_fields: &[&str] = if expected_part == "input_text" { + &["type", "text"] + } else { + &["type", "text", "annotations"] + }; + part.keys().any(|field| !allowed_fields.contains(&field.as_str())) + || part.get("type").and_then(Value::as_str) != Some(expected_part) + || part.get("text").and_then(Value::as_str).is_none() + }) { + return Err(unsupported_agent_replay_item()); + } + Ok(()) +} + +fn contains_threadmark_uri(value: &Value) -> bool { + match value { + Value::String(value) => files::parse_uri(value).is_some(), + Value::Array(values) => values.iter().any(contains_threadmark_uri), + Value::Object(values) => values.values().any(contains_threadmark_uri), + _ => false, + } +} + +fn unsupported_agent_replay_item() -> ApiError { + ApiError::CodedBadRequest { + code: "unsupported_agent_replay_item", + message: "Agent replay supports only user input_text and assistant output_text messages." + .into(), + } +} + async fn hydrate_file_references( state: &AppState, actor: &Actor, @@ -1354,4 +1529,121 @@ mod tests { vec!["file_document".to_owned(), "file_image".to_owned()] ); } + + #[test] + fn agent_replay_accepts_only_the_text_message_contract() { + let rows = vec![ + ( + 1, + json!({ + "id": "provider-user-id", + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }), + ), + ( + 2, + json!({ + "id": "provider-output-id", + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "hi", "annotations": []}] + }), + ), + ]; + let projected = build_agent_projection(rows, &["id".into()], 2, usize::MAX).unwrap(); + assert!(projected.iter().all(|item| item.get("id").is_none())); + assert_eq!(projected[0]["role"], "user"); + assert_eq!(projected[1]["role"], "assistant"); + assert_eq!(projected[1]["content"][0]["annotations"], json!([])); + } + + #[test] + fn agent_replay_strips_only_configured_top_level_fields() { + let projected = build_agent_projection( + vec![( + 1, + json!({ + "id": "keep-me", + "metadata": {"id": "nested", "private": true}, + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }), + )], + &["metadata".into()], + 1, + usize::MAX, + ) + .unwrap(); + assert_eq!(projected[0]["id"], "keep-me"); + assert!(projected[0].get("metadata").is_none()); + } + + #[test] + fn agent_replay_item_limit_is_inclusive() { + let item = json!({ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }); + assert!(build_agent_projection(vec![(1, item.clone())], &[], 1, usize::MAX).is_ok()); + assert!(matches!( + build_agent_projection(vec![(1, item)], &[], 0, usize::MAX), + Err(ApiError::CodedPayloadTooLarge { code: "context_limit_exceeded", .. }) + )); + } + + #[test] + fn agent_replay_serialized_byte_limit_is_inclusive() { + let item = json!({ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }); + let expected = vec![item.clone()]; + let exact = serde_json::to_vec(&expected).unwrap().len(); + assert!(build_agent_projection(vec![(1, item.clone())], &[], 1, exact).is_ok()); + assert!(matches!( + build_agent_projection(vec![(1, item)], &[], 1, exact - 1), + Err(ApiError::CodedPayloadTooLarge { code: "context_limit_exceeded", .. }) + )); + } + + #[test] + fn agent_replay_rejects_media_and_role_part_mismatches() { + for item in [ + json!({ + "type": "message", + "role": "user", + "content": [{"type": "input_image", "image_url": "threadmark://files/file_1"}] + }), + json!({ + "type": "message", + "role": "assistant", + "content": [{"type": "input_text", "text": "wrong direction"}] + }), + json!({ + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": "hello", + "image_url": "data:image/png;base64,AAAA" + }] + }), + json!({ + "type": "message", + "role": "user", + "metadata": {"source": "threadmark://files/file_1"}, + "content": [{"type": "input_text", "text": "hello"}] + }), + ] { + assert!(matches!( + build_agent_projection(vec![(1, item)], &[], 1, usize::MAX), + Err(ApiError::CodedBadRequest { code: "unsupported_agent_replay_item", .. }) + )); + } + } } From be78fe76b28da93bfdc6e0ed44dfe3f49851bafb Mon Sep 17 00:00:00 2001 From: "dispatch-developer[bot]" <306909890+dispatch-developer[bot]@users.noreply.github.com> Date: Mon, 17 Aug 2026 15:31:30 +0000 Subject: [PATCH 2/2] style: format Rust sources --- src/config.rs | 20 +++++++++----------- src/error.rs | 4 +--- src/store.rs | 18 ++++++++++++++---- 3 files changed, 24 insertions(+), 18 deletions(-) diff --git a/src/config.rs b/src/config.rs index 5dfac82..7c93c3e 100644 --- a/src/config.rs +++ b/src/config.rs @@ -104,8 +104,7 @@ impl Config { auth_max_owner_token_seconds > 0, "AUTH_MAX_OWNER_TOKEN_SECONDS must be greater than zero" ); - let auth_max_delegated_token_seconds = - parse("AUTH_MAX_DELEGATED_TOKEN_SECONDS", "600")?; + let auth_max_delegated_token_seconds = parse("AUTH_MAX_DELEGATED_TOKEN_SECONDS", "600")?; ensure!( auth_max_delegated_token_seconds > 0, "AUTH_MAX_DELEGATED_TOKEN_SECONDS must be greater than zero" @@ -122,15 +121,14 @@ impl Config { agent_replay_max_bytes > 0, "AGENT_REPLAY_MAX_BYTES must be greater than zero" ); - let agent_replay_strip_top_level_fields = std::env::var( - "AGENT_REPLAY_STRIP_TOP_LEVEL_FIELDS", - ) - .unwrap_or_else(|_| "id".into()) - .split(',') - .map(str::trim) - .filter(|field| !field.is_empty()) - .map(str::to_owned) - .collect::>(); + let agent_replay_strip_top_level_fields = + std::env::var("AGENT_REPLAY_STRIP_TOP_LEVEL_FIELDS") + .unwrap_or_else(|_| "id".into()) + .split(',') + .map(str::trim) + .filter(|field| !field.is_empty()) + .map(str::to_owned) + .collect::>(); let mut unique_fields = agent_replay_strip_top_level_fields.clone(); unique_fields.sort(); unique_fields.dedup(); diff --git a/src/error.rs b/src/error.rs index 59da4f1..6613720 100644 --- a/src/error.rs +++ b/src/error.rs @@ -48,9 +48,7 @@ impl IntoResponse for ApiError { Self::CodedGone { code, .. } => (StatusCode::GONE, *code), Self::CodedUnavailable { code, .. } => (StatusCode::SERVICE_UNAVAILABLE, *code), Self::PayloadTooLarge(_) => (StatusCode::PAYLOAD_TOO_LARGE, "payload_too_large"), - Self::CodedPayloadTooLarge { code, .. } => { - (StatusCode::PAYLOAD_TOO_LARGE, *code) - } + Self::CodedPayloadTooLarge { code, .. } => (StatusCode::PAYLOAD_TOO_LARGE, *code), Self::ObjectStore(error) => { tracing::error!(?error, "object storage request failed"); (StatusCode::BAD_GATEWAY, "object_store_error") diff --git a/src/store.rs b/src/store.rs index 8b8dd0a..347484c 100644 --- a/src/store.rs +++ b/src/store.rs @@ -1003,7 +1003,8 @@ fn validate_agent_text_item(item: &Value) -> ApiResult<()> { } else { &["type", "text", "annotations"] }; - part.keys().any(|field| !allowed_fields.contains(&field.as_str())) + part.keys() + .any(|field| !allowed_fields.contains(&field.as_str())) || part.get("type").and_then(Value::as_str) != Some(expected_part) || part.get("text").and_then(Value::as_str).is_none() }) { @@ -1591,7 +1592,10 @@ mod tests { assert!(build_agent_projection(vec![(1, item.clone())], &[], 1, usize::MAX).is_ok()); assert!(matches!( build_agent_projection(vec![(1, item)], &[], 0, usize::MAX), - Err(ApiError::CodedPayloadTooLarge { code: "context_limit_exceeded", .. }) + Err(ApiError::CodedPayloadTooLarge { + code: "context_limit_exceeded", + .. + }) )); } @@ -1607,7 +1611,10 @@ mod tests { assert!(build_agent_projection(vec![(1, item.clone())], &[], 1, exact).is_ok()); assert!(matches!( build_agent_projection(vec![(1, item)], &[], 1, exact - 1), - Err(ApiError::CodedPayloadTooLarge { code: "context_limit_exceeded", .. }) + Err(ApiError::CodedPayloadTooLarge { + code: "context_limit_exceeded", + .. + }) )); } @@ -1642,7 +1649,10 @@ mod tests { ] { assert!(matches!( build_agent_projection(vec![(1, item)], &[], 1, usize::MAX), - Err(ApiError::CodedBadRequest { code: "unsupported_agent_replay_item", .. }) + Err(ApiError::CodedBadRequest { + code: "unsupported_agent_replay_item", + .. + }) )); } }