Persist security risk scores in rollout history (#38363)

## What changed

- Add a `SecurityRiskScore` rollout item containing a category and numeric score.
- Persist the item in both thread history modes while excluding it from model context, user-visible thread history, search text, forks, and reconstructed conversation history.
- Re-export the score type from the extension API.

## Testing

- Cover serialization, persistence and loading, thread history projection, session reconstruction, append planning, and memory filtering.

GitOrigin-RevId: 1926fe366aeaa75052708a6da589f45a38eefb52
This commit is contained in:
jif
2026-08-13 12:04:29 +00:00
committed by copyberry
parent 9ed0047a61
commit 72fa74fbc9
32 changed files with 165 additions and 4 deletions
@@ -405,6 +405,7 @@ impl ThreadHistoryBuilder {
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::SessionMeta(_) => {}
}
}
@@ -82,6 +82,7 @@ pub fn project_rollout_line(line: &RolloutLine) -> ThreadHistoryChangeSet {
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => ThreadHistoryChangeSet::default(),
}
}
@@ -10,6 +10,7 @@ use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::TurnCompleteEvent;
use codex_protocol::protocol::TurnStartedEvent;
use codex_protocol::security_risk::SecurityRiskScore;
use codex_protocol::user_input::UserInput;
use codex_rollout::CompactedItem;
use codex_rollout::RolloutItem;
@@ -200,9 +201,14 @@ fn ignores_legacy_abort_without_turn_id_and_context_only_records() {
previous_window_id: None,
window_id: None,
}));
let security_risk = project(RolloutItem::SecurityRiskScore(SecurityRiskScore {
category: "action_risk".to_string(),
score: 0.92,
}));
assert!(aborted.is_empty());
assert!(compacted.is_empty());
assert!(security_risk.is_empty());
}
#[test]
+3 -1
View File
@@ -71,7 +71,8 @@ fn keep_forked_rollout_item(item: &RolloutItem, preserve_reference_context_item:
| ResponseItem::Other => false,
},
RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. } => false,
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::SecurityRiskScore(_) => false,
// Full-history forks preserve the cached prompt prefix and can keep diffing
// from the parent's durable baseline. Truncated forks drop part of that prompt,
// so they must rebuild context on their first child turn.
@@ -794,6 +795,7 @@ impl AgentControl {
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::WorldState(_) => true,
RolloutItem::SecurityRiskScore(_) => false,
}
});
// Full forks reuse the parent's reference context instead of rebuilding it. If that
+1
View File
@@ -250,6 +250,7 @@ async fn persisted_originator(thread: &CodexThread) -> String {
| RolloutItem::EventMsg(_)
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::TurnContext(_) => None,
})
.expect("session metadata should be persisted")
@@ -282,6 +282,7 @@ impl Session {
}
RolloutItem::EventMsg(_)
| RolloutItem::SessionMeta(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::InterAgentCommunicationMetadata { .. } => {}
}
@@ -371,6 +372,7 @@ impl Session {
RolloutItem::EventMsg(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::SessionMeta(_) => {}
}
}
@@ -409,6 +411,7 @@ impl Session {
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::TurnContext(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => {
unreachable!("only world-state replay items are collected")
}
@@ -16,6 +16,7 @@ use codex_protocol::protocol::SessionContextWindow;
use codex_protocol::protocol::SessionMeta;
use codex_protocol::protocol::SessionMetaLine;
use codex_protocol::protocol::WorldStateItem;
use codex_protocol::security_risk::SecurityRiskScore;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::path::PathBuf;
@@ -146,6 +147,32 @@ async fn record_initial_history_reconstructs_typed_inter_agent_message() {
);
}
#[tokio::test]
async fn record_initial_history_ignores_security_risk_scores() {
let (session, _turn_context) = make_session_and_context().await;
let user_item = user_message("visible user input");
let security_risk = SecurityRiskScore {
category: "credential_access".to_string(),
score: 0.92,
};
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ThreadId::default(),
history: Arc::new(vec![
RolloutItem::ResponseItem(ResponseItemEnvelope::new(user_item.clone())),
RolloutItem::SecurityRiskScore(security_risk),
]),
rollout_path: Some(PathBuf::from("/tmp/resume.jsonl")),
}))
.await;
assert_eq!(
raw_history_items(&session.state.lock().await.clone_history()),
vec![user_item]
);
}
#[tokio::test]
async fn record_initial_history_restores_world_state_baseline() {
let (session, turn_context) = make_session_and_context().await;
+2
View File
@@ -3256,6 +3256,7 @@ async fn start_new_context_window_assigns_and_persists_item_ids() {
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => None,
});
assert_eq!(
@@ -3322,6 +3323,7 @@ async fn record_initial_history_assigns_and_persists_id_for_forked_response_item
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => None,
});
let persisted_item = persisted_item.expect("forked response item should be persisted");
+1
View File
@@ -15,6 +15,7 @@ pub use capabilities::ResponseItemInjectionFuture;
pub use capabilities::ResponseItemInjector;
pub use codex_context_fragments::ContextualUserFragment;
pub use codex_protocol::models::ResponseItem;
pub use codex_protocol::security_risk::SecurityRiskScore;
pub use codex_tools::ConversationHistory;
pub use codex_tools::ExtensionTurnItem;
pub use codex_tools::FunctionCallError;
@@ -241,6 +241,7 @@ fn source_model_items(items: &[RolloutItem]) -> Option<Vec<SourceModelItem<'_>>>
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::WorldState(_) => return None,
RolloutItem::EventMsg(_) => {}
}
@@ -252,7 +253,9 @@ fn history_model_items(items: &[RolloutItem]) -> Option<Vec<&ResponseItem>> {
let mut model_items = Vec::new();
for item in items {
match item {
RolloutItem::SessionMeta(_) | RolloutItem::InterAgentCommunicationMetadata { .. } => {}
RolloutItem::SessionMeta(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::SecurityRiskScore(_) => {}
RolloutItem::ResponseItem(response_item) => model_items.push(&response_item.item),
RolloutItem::EventMsg(
EventMsg::ContextCompacted(_) | EventMsg::ThreadRolledBack(_),
@@ -7,6 +7,7 @@ use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::ContextCompactedEvent;
use codex_protocol::protocol::ThreadRolledBackEvent;
use codex_protocol::security_risk::SecurityRiskScore;
use pretty_assertions::assert_eq;
#[test]
@@ -62,9 +63,17 @@ fn requires_a_strict_nonempty_model_prefix() {
event.started_at = Some(9_999);
}
}
let security_risk = RolloutItem::SecurityRiskScore(SecurityRiskScore {
category: "action_risk".to_string(),
score: 0.92,
});
metadata_changed.push(security_risk.clone());
assert!(model_transcripts_match(&history, &metadata_changed));
assert!(!model_transcripts_match(&source, &history));
assert!(plan_append(&source, &metadata_changed).is_some());
let mut source_with_security_risk = source.clone();
source_with_security_risk.push(security_risk);
assert!(plan_append(&source_with_security_risk, &metadata_changed).is_none());
for event in [
EventMsg::ContextCompacted(ContextCompactedEvent),
+3
View File
@@ -22,6 +22,7 @@ use codex_protocol::protocol::ThreadHistoryMode;
use codex_protocol::protocol::ThreadSource;
use codex_protocol::protocol::TurnContextItem;
use codex_protocol::protocol::WorldStateItem;
use codex_protocol::security_risk::SecurityRiskScore;
use schemars::JsonSchema;
use serde::Deserialize;
use serde::Deserializer;
@@ -98,6 +99,7 @@ pub enum RolloutItem {
Compacted(CompactedItem),
TurnContext(TurnContextItem),
WorldState(WorldStateItem),
SecurityRiskScore(SecurityRiskScore),
EventMsg(EventMsg),
}
@@ -410,6 +412,7 @@ fn multi_agent_version_from_items(
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => None,
})
})
+10
View File
@@ -7,6 +7,7 @@ use super::InterAgentCommunication;
use super::ResponseItem;
use super::ResponseItemEnvelope;
use super::RolloutItem;
use super::SecurityRiskScore;
use super::SessionMetaLine;
use super::TurnContextItem;
use super::WorldStateItem;
@@ -41,6 +42,9 @@ pub(super) enum RolloutItemWire<'a> {
WorldState {
payload: Cow<'a, WorldStateItem>,
},
SecurityRiskScore {
payload: Cow<'a, SecurityRiskScore>,
},
EventMsg {
payload: Cow<'a, EventMsg>,
},
@@ -75,6 +79,9 @@ impl<'a> From<&'a RolloutItem> for RolloutItemWire<'a> {
RolloutItem::WorldState(payload) => Self::WorldState {
payload: Cow::Borrowed(payload),
},
RolloutItem::SecurityRiskScore(payload) => Self::SecurityRiskScore {
payload: Cow::Borrowed(payload),
},
RolloutItem::EventMsg(payload) => Self::EventMsg {
payload: Cow::Borrowed(payload),
},
@@ -103,6 +110,9 @@ impl From<RolloutItemWire<'_>> for RolloutItem {
RolloutItemWire::Compacted { payload } => Self::Compacted(payload.into_owned()),
RolloutItemWire::TurnContext { payload } => Self::TurnContext(payload.into_owned()),
RolloutItemWire::WorldState { payload } => Self::WorldState(payload.into_owned()),
RolloutItemWire::SecurityRiskScore { payload } => {
Self::SecurityRiskScore(payload.into_owned())
}
RolloutItemWire::EventMsg { payload } => Self::EventMsg(payload.into_owned()),
}
}
+5 -1
View File
@@ -352,6 +352,10 @@ fn rollout_item_variants_preserve_existing_payload_shapes() -> Result<()> {
"type": "world_state",
"payload": { "full": true, "state": { "cwd": "/tmp" } },
}),
json!({
"type": "security_risk_score",
"payload": { "category": "action_risk", "score": 0.92 },
}),
json!({
"type": "event_msg",
"payload": { "type": "warning", "message": "heads up" },
@@ -370,7 +374,7 @@ fn rollout_item_variants_preserve_existing_payload_shapes() -> Result<()> {
fn rollout_item_schema_matches_tagged_payload_and_sibling_metadata() -> Result<()> {
let schema = serde_json::to_value(schemars::schema_for!(RolloutItem))?;
let variants = schema["oneOf"].as_array().expect("rollout variants");
assert_eq!(variants.len(), 8);
assert_eq!(variants.len(), 9);
for variant in variants {
let required = variant["required"].as_array().expect("required fields");
+6
View File
@@ -417,6 +417,7 @@ mod job {
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => None,
})
.collect::<Vec<_>>();
@@ -672,6 +673,7 @@ mod tests {
use super::*;
use codex_protocol::AgentPath;
use codex_protocol::protocol::InterAgentCommunication;
use codex_protocol::security_risk::SecurityRiskScore;
use pretty_assertions::assert_eq;
#[test]
@@ -722,6 +724,10 @@ mod tests {
let serialized = job::serialize_filtered_rollout_response_items(&[
RolloutItem::ResponseItem(mixed_contextual_message.into()),
RolloutItem::ResponseItem(skill_message.into()),
RolloutItem::SecurityRiskScore(SecurityRiskScore {
category: "action_risk".to_string(),
score: 0.92,
}),
RolloutItem::ResponseItem(subagent_message.clone().into()),
])
.expect("serialize");
+1
View File
@@ -35,6 +35,7 @@ pub mod protocol;
pub mod request_permissions;
pub mod request_user_input;
pub mod review_format;
pub mod security_risk;
pub mod shell_environment;
pub mod turn_input;
pub mod user_input;
+13
View File
@@ -0,0 +1,13 @@
use schemars::JsonSchema;
use serde::Deserialize;
use serde::Serialize;
/// A durable, thread-owned security risk classifier score.
///
/// These records belong to rollout history only and must not enter model-visible
/// conversation context or user-visible thread item projections.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
pub struct SecurityRiskScore {
pub category: String,
pub score: f64,
}
+3 -1
View File
@@ -1170,7 +1170,7 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result<HeadTai
RolloutItem::TurnContext(_) => {
// Not included in `head`; skip.
}
RolloutItem::WorldState(_) => {
RolloutItem::WorldState(_) | RolloutItem::SecurityRiskScore(_) => {
// Not included in `head`; skip.
}
RolloutItem::Compacted(_) => {
@@ -1243,6 +1243,7 @@ pub async fn read_head_for_summary(path: &Path) -> io::Result<Vec<serde_json::Va
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => {}
}
}
@@ -1295,6 +1296,7 @@ pub async fn read_session_meta_line(path: &Path) -> io::Result<SessionMetaLine>
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => {}
}
}
+2
View File
@@ -75,6 +75,7 @@ pub fn builder_from_items(
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => None,
}) && let Some(builder) = builder_from_session_meta(session_meta, rollout_path)
{
@@ -143,6 +144,7 @@ pub async fn extract_metadata_from_rollout(
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => None,
}),
parse_errors,
+1
View File
@@ -153,6 +153,7 @@ impl ModelContextScan {
RolloutItem::EventMsg(_)
| RolloutItem::SessionMeta(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::WorldState(_) => {}
}
@@ -236,6 +236,7 @@ fn rollout_item_type(item: &RolloutItem) -> String {
RolloutItem::Compacted(_) => "compacted".to_string(),
RolloutItem::TurnContext(_) => "turn_context".to_string(),
RolloutItem::WorldState(_) => "world_state".to_string(),
RolloutItem::SecurityRiskScore(_) => "security_risk_score".to_string(),
RolloutItem::EventMsg(EventMsg::ItemCompleted(event)) => {
format!("event.item_completed.{}", turn_item_type(&event.item))
}
+1
View File
@@ -16,6 +16,7 @@ pub fn is_persisted_rollout_item(item: &RolloutItem, history_mode: ThreadHistory
RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::SessionMeta(_) => true,
}
}
+1
View File
@@ -2067,6 +2067,7 @@ async fn resume_candidate_matches_cwd(
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => None,
})
{
+52
View File
@@ -20,6 +20,7 @@ use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::ThreadHistoryMode;
use codex_protocol::protocol::TurnContextItem;
use codex_protocol::protocol::UserMessageEvent;
use codex_protocol::security_risk::SecurityRiskScore;
use codex_utils_absolute_path::test_support::PathExt;
use pretty_assertions::assert_eq;
use std::fs;
@@ -431,6 +432,57 @@ async fn load_rollout_items_preserves_legacy_guardian_assessment_lines() -> std:
Ok(())
}
#[tokio::test]
async fn load_rollout_items_preserves_security_risk_scores() -> std::io::Result<()> {
let home = TempDir::new().expect("temp dir");
let rollout_path = home.path().join("rollout.jsonl");
let thread_id = ThreadId::new();
let security_risk = SecurityRiskScore {
category: "action_risk".to_string(),
score: 0.76,
};
let security_risk_item = RolloutItem::SecurityRiskScore(security_risk.clone());
for history_mode in [ThreadHistoryMode::Legacy, ThreadHistoryMode::Paginated] {
assert!(crate::is_persisted_rollout_item(
&security_risk_item,
history_mode
));
}
let mut file = File::create(&rollout_path)?;
for (ordinal, item) in [
paginated_session_meta_item(thread_id, home.path()),
security_risk_item,
]
.into_iter()
.enumerate()
{
let line = RolloutLine {
timestamp: "2026-07-09T00:00:00Z".to_string(),
ordinal: Some(ordinal as u64),
item,
};
writeln!(
file,
"{}",
serde_json::to_string(&line).map_err(std::io::Error::other)?
)?;
}
let (items, loaded_thread_id, parse_errors) =
RolloutRecorder::load_rollout_items(&rollout_path).await?;
assert_eq!(loaded_thread_id, Some(thread_id));
assert_eq!(parse_errors, 0);
assert_eq!(items.len(), 2);
let RolloutItem::SecurityRiskScore(persisted_security_risk) = &items[1] else {
panic!("expected security risk score rollout item");
};
assert_eq!(persisted_security_risk, &security_risk);
Ok(())
}
#[tokio::test]
async fn load_rollout_items_filters_legacy_ghost_snapshots_from_compaction_history()
-> std::io::Result<()> {
+1
View File
@@ -292,6 +292,7 @@ fn conversation_text_from_item(item: &RolloutItem) -> Option<String> {
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::WorldState(_) => None,
}
}
+2
View File
@@ -26,6 +26,7 @@ pub fn apply_rollout_item(
| RolloutItem::InterAgentCommunicationMetadata { .. } => {}
RolloutItem::Compacted(_) => {}
RolloutItem::WorldState(_) => {}
RolloutItem::SecurityRiskScore(_) => {}
}
if metadata.model_provider.is_empty() {
metadata.model_provider = default_provider.to_string();
@@ -52,6 +53,7 @@ pub fn rollout_item_affects_thread_metadata(item: &RolloutItem) -> bool {
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::WorldState(_) => false,
}
}
+1
View File
@@ -1293,6 +1293,7 @@ pub(super) fn extract_memory_mode(items: &[RolloutItem]) -> Option<String> {
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => None,
})
}
@@ -292,6 +292,7 @@ impl LegacyRolloutCanonicalizer {
}
item @ (RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::TurnContext(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::WorldState(_)) => {
self.write_item(writer, &timestamp, item).await?;
}
@@ -220,6 +220,7 @@ impl RollbackPlanner {
}
}
RolloutItem::WorldState(_) => {}
RolloutItem::SecurityRiskScore(_) => self.record_boundaries[index] = None,
}
Ok(())
@@ -75,6 +75,7 @@ impl ModelReplayPlanner {
| RolloutItem::ResponseItem(_)
| RolloutItem::EventMsg(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::WorldState(_) => return,
};
self.records.push(record);
@@ -304,6 +304,7 @@ impl ThreadMetadataSync {
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::WorldState(_) => {}
}
}