Require transition policy replay

src/session.rs

@@ -48,6 +48,8 @@ struct PolicyAppliedPayload {
 struct TransitionStartedPayload {
     schema_version: u16,
     commitment_id: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    transition_policy_id: Option<String>,
     reason: String,
     return_target: String,
     expected_duration_secs: u64,
@@ -57,6 +59,15 @@ struct TransitionStartedPayload {
     commitment_to: CommitmentState,
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+struct PendingTransitionPolicy {
+    commitment_id: String,
+    policy_id: Option<String>,
+    runtime_state: RuntimeState,
+    enforcement_level: EnforcementLevel,
+    allowed_context: AllowedContext,
+}
+
 #[derive(Debug)]
 pub struct SessionController {
     runtime_state: RuntimeState,
@@ -227,6 +238,7 @@ impl SessionController {
         let transition_payload = TransitionStartedPayload {
             schema_version: POLICY_SCHEMA_VERSION,
             commitment_id: next_commitment.id.clone(),
+            transition_policy_id: Some(transition_policy.id.clone()),
             reason: reason.into().trim().to_string(),
             return_target: return_target.into().trim().to_string(),
             expected_duration_secs: expected_duration.as_secs(),
@@ -276,8 +288,16 @@ fn hydrate_session(
     let mut runtime_state = RuntimeState::Locked;
     let mut active_commitment: Option<Commitment> = None;
     let mut active_policy: Option<PolicySnapshot> = None;
+    let mut pending_transition_policy: Option<PendingTransitionPolicy> = None;
 
     for record in event_log.records()? {
+        if pending_transition_policy.is_some() && record.event_type != EventType::PolicyApplied {
+            return Err(anyhow!(
+                "pending transition policy expected PolicyApplied at sequence {}",
+                record.sequence
+            ));
+        }
+
         match record.event_type {
             EventType::RuntimeTransition => {
                 let payload: RuntimeTransitionPayload =
@@ -353,6 +373,17 @@ fn hydrate_session(
             EventType::PolicyApplied => {
                 let payload: PolicyAppliedPayload = parse_payload(&record, "policy applied")?;
                 ensure_schema_version(payload.schema_version, record.sequence)?;
+                if let Some(expected) = pending_transition_policy.take() {
+                    apply_pending_transition_policy(
+                        &record,
+                        payload,
+                        expected,
+                        &mut runtime_state,
+                        &mut active_commitment,
+                        &mut active_policy,
+                    )?;
+                    continue;
+                }
                 ensure_record_commitment(&record, &payload.commitment_id)?;
                 let commitment = active_commitment.as_ref().ok_or_else(|| {
                     anyhow!(
@@ -472,29 +503,16 @@ fn hydrate_session(
                 }
                 let mut commitment = active.clone();
                 commitment.state = composite.commitment_state;
-                let transition_policy = PolicySnapshot {
+                commitment.transition_policy_id = payload.transition_policy_id.clone();
-                    id: commitment
+                pending_transition_policy = Some(PendingTransitionPolicy {
-                        .transition_policy_id
-                        .clone()
-                        .unwrap_or_else(|| current_policy.id.clone()),
                     commitment_id: commitment.id.clone(),
-                    schema_version: POLICY_SCHEMA_VERSION,
+                    policy_id: payload.transition_policy_id,
-                    created_at_unix_secs: record.timestamp_unix_secs,
                     runtime_state: composite.runtime_state,
                     enforcement_level: current_policy.enforcement_level,
                     allowed_context: current_policy.allowed_context.clone(),
-                    required_monitors: Vec::new(),
+                });
-                    violation_actions: Vec::new(),
-                    expires_at_unix_secs: None,
-                    generated_by_agent_version: env!("CARGO_PKG_VERSION").to_string(),
-                };
-                transition_policy.validate().with_context(|| {
-                    format!("validate transition policy at sequence {}", record.sequence)
-                })?;
-                commitment.transition_policy_id = Some(transition_policy.id.clone());
                 runtime_state = composite.runtime_state;
                 active_commitment = Some(commitment);
-                active_policy = Some(transition_policy);
             }
             ref unsupported => {
                 return Err(anyhow!(
@@ -506,6 +524,11 @@ fn hydrate_session(
         }
     }
 
+    if pending_transition_policy.is_some() {
+        return Err(anyhow!(
+            "session replay ended with pending transition policy"
+        ));
+    }
     if active_commitment.is_some() && active_policy.is_none() {
         return Err(anyhow!(
             "session replay ended with active commitment but no active policy"
@@ -525,6 +548,85 @@ fn hydrate_session(
     Ok((runtime_state, active_commitment, active_policy))
 }
 
+fn apply_pending_transition_policy(
+    record: &EventRecord,
+    payload: PolicyAppliedPayload,
+    expected: PendingTransitionPolicy,
+    runtime_state: &mut RuntimeState,
+    active_commitment: &mut Option<Commitment>,
+    active_policy: &mut Option<PolicySnapshot>,
+) -> Result<()> {
+    ensure_record_commitment(record, &payload.commitment_id)?;
+    if payload.commitment_id != expected.commitment_id {
+        return Err(anyhow!(
+            "transition policy commitment mismatch at sequence {}",
+            record.sequence
+        ));
+    }
+    if record.runtime_state != expected.runtime_state
+        || payload.runtime_state != expected.runtime_state
+    {
+        return Err(anyhow!(
+            "transition policy runtime mismatch at sequence {}",
+            record.sequence
+        ));
+    }
+    if payload.enforcement_level != expected.enforcement_level {
+        return Err(anyhow!(
+            "transition policy enforcement mismatch at sequence {}",
+            record.sequence
+        ));
+    }
+    if payload.allowed_context != expected.allowed_context {
+        return Err(anyhow!(
+            "transition policy allowed context mismatch at sequence {}",
+            record.sequence
+        ));
+    }
+    if let Some(expected_policy_id) = expected.policy_id.as_ref() {
+        if payload.policy_id != *expected_policy_id {
+            return Err(anyhow!(
+                "transition policy id mismatch at sequence {}",
+                record.sequence
+            ));
+        }
+    }
+
+    let policy = PolicySnapshot {
+        id: payload.policy_id,
+        commitment_id: payload.commitment_id,
+        schema_version: POLICY_SCHEMA_VERSION,
+        created_at_unix_secs: record.timestamp_unix_secs,
+        runtime_state: payload.runtime_state,
+        enforcement_level: payload.enforcement_level,
+        allowed_context: payload.allowed_context,
+        required_monitors: Vec::new(),
+        violation_actions: Vec::new(),
+        expires_at_unix_secs: None,
+        generated_by_agent_version: env!("CARGO_PKG_VERSION").to_string(),
+    };
+    policy
+        .validate()
+        .with_context(|| format!("validate transition policy at sequence {}", record.sequence))?;
+    if let Some(commitment) = active_commitment.as_mut() {
+        if commitment.id != policy.commitment_id {
+            return Err(anyhow!(
+                "transition policy active commitment mismatch at sequence {}",
+                record.sequence
+            ));
+        }
+        commitment.transition_policy_id = Some(policy.id.clone());
+    } else {
+        return Err(anyhow!(
+            "transition policy references unknown commitment at sequence {}",
+            record.sequence
+        ));
+    }
+    *runtime_state = expected.runtime_state;
+    *active_policy = Some(policy);
+    Ok(())
+}
+
 fn parse_payload<T>(record: &EventRecord, description: &str) -> Result<T>
 where
     T: for<'de> Deserialize<'de>,
@@ -646,6 +748,7 @@ mod tests {
         serde_json::to_value(TransitionStartedPayload {
             schema_version: POLICY_SCHEMA_VERSION,
             commitment_id: "commitment-test".to_string(),
+            transition_policy_id: Some("policy-transition-test".to_string()),
             reason: "Need to inspect failures".to_string(),
             return_target: "Return to session controller".to_string(),
             expected_duration_secs: 300,
@@ -697,6 +800,23 @@ mod tests {
         .unwrap();
     }
 
+    fn append_active_session_with_transition_and_policy_payload(
+        path: &PathBuf,
+        transition_payload: serde_json::Value,
+        policy_payload: serde_json::Value,
+    ) {
+        let commitment_id = "commitment-test";
+        append_active_session_with_transition_payload(path, transition_payload);
+        let mut log = EventLog::open(path).unwrap();
+        log.append(
+            EventType::PolicyApplied,
+            RuntimeState::Transition,
+            Some(commitment_id.to_string()),
+            policy_payload,
+        )
+        .unwrap();
+    }
+
     fn error_chain_contains(err: &anyhow::Error, expected: &str) -> bool {
         err.chain()
             .any(|cause| cause.to_string().contains(expected))
@@ -1090,6 +1210,87 @@ mod tests {
         fs::remove_file(path).unwrap();
     }
 
+    #[test]
+    fn rejects_transition_started_without_policy_on_replay() {
+        let path = temp_log_path("transition-without-policy");
+        append_active_session_with_transition_payload(&path, transition_started_payload());
+
+        let err = SessionController::new(&path)
+            .expect_err("transition started without persisted policy must be rejected");
+
+        assert!(error_chain_contains(&err, "pending transition policy"));
+        fs::remove_file(path).unwrap();
+    }
+
+    #[test]
+    fn rejects_transition_policy_enforcement_mismatch_on_replay() {
+        let path = temp_log_path("transition-policy-enforcement-mismatch");
+        let commitment_id = "commitment-test";
+        let mut payload = transition_policy_payload(commitment_id);
+        payload["enforcement_level"] = json!("block");
+        append_active_session_with_transition_and_policy_payload(
+            &path,
+            transition_started_payload(),
+            payload,
+        );
+
+        let err = SessionController::new(&path)
+            .expect_err("transition policy with impossible enforcement must be rejected");
+
+        assert!(error_chain_contains(&err, "transition policy enforcement"));
+        fs::remove_file(path).unwrap();
+    }
+
+    #[test]
+    fn rejects_non_policy_record_after_transition_started_on_replay() {
+        let path = temp_log_path("transition-interrupted-pair");
+        append_active_session_with_transition_payload(&path, transition_started_payload());
+        {
+            let mut log = EventLog::open(&path).unwrap();
+            log.append(
+                EventType::RuntimeTransition,
+                RuntimeState::Planning,
+                None,
+                json!({
+                    "schema_version": POLICY_SCHEMA_VERSION,
+                    "action": "enter_planning",
+                    "from": "transition",
+                    "to": "planning",
+                }),
+            )
+            .unwrap();
+        }
+
+        let err =
+            SessionController::new(&path).expect_err("transition policy pair must be adjacent");
+
+        assert!(error_chain_contains(&err, "expected PolicyApplied"));
+        fs::remove_file(path).unwrap();
+    }
+
+    #[test]
+    fn hydrates_transition_from_valid_replay_pair() {
+        let path = temp_log_path("transition-valid-pair");
+        append_active_session_with_transition_and_policy_payload(
+            &path,
+            transition_started_payload(),
+            transition_policy_payload("commitment-test"),
+        );
+
+        let reopened = SessionController::new(&path).unwrap();
+
+        assert_eq!(reopened.runtime_state(), RuntimeState::Transition);
+        assert_eq!(
+            reopened.active_commitment().unwrap().state,
+            CommitmentState::Paused
+        );
+        assert_eq!(
+            reopened.active_policy().unwrap().runtime_state,
+            RuntimeState::Transition
+        );
+        fs::remove_file(path).unwrap();
+    }
+
     #[test]
     fn rejects_transition_policy_runtime_mismatch_on_replay() {
         let path = temp_log_path("transition-policy-runtime-mismatch");
@@ -1109,7 +1310,10 @@ mod tests {
         let err = SessionController::new(&path)
             .expect_err("transition policy runtime mismatch must be rejected");
 
-        assert!(error_chain_contains(&err, "policy runtime target mismatch"));
+        assert!(error_chain_contains(
+            &err,
+            "transition policy runtime mismatch"
+        ));
         fs::remove_file(path).unwrap();
     }
 
@@ -1147,6 +1351,10 @@ mod tests {
             transition.payload_json["commitment_id"],
             transition.commitment_id.as_deref().unwrap()
         );
+        assert_eq!(
+            transition.payload_json["transition_policy_id"],
+            transition_policy.payload_json["policy_id"]
+        );
         assert_eq!(
             transition.payload_json["reason"],
             "Need to inspect failures"