diff --git a/crates/harness-server/src/omp_rpc.rs b/crates/harness-server/src/omp_rpc.rs index 746046ec0..8e0361f52 100644 --- a/crates/harness-server/src/omp_rpc.rs +++ b/crates/harness-server/src/omp_rpc.rs @@ -42,6 +42,19 @@ use crate::turn::CodexTurnNormalizer; use crate::util::write_value; use crate::wire::collab_state_wire_value; use crate::{HarnessServerError, Result}; +const DEFAULT_OMP_TURN_TIMEOUT: Duration = Duration::from_secs(300); +const OMP_TURN_TIMEOUT_GRACE: Duration = Duration::from_secs(5); + +fn turn_timeout_from_trace_context(trace_context: &crate::otel::TraceContext) -> Duration { + trace_context + .metadata + .get("max_duration_ms") + .and_then(Value::as_u64) + .filter(|duration_ms| *duration_ms > 0) + .map(Duration::from_millis) + .and_then(|duration| duration.checked_add(OMP_TURN_TIMEOUT_GRACE)) + .unwrap_or(DEFAULT_OMP_TURN_TIMEOUT) +} /// The resident session ownership fence. Admission requires both fields; a /// stale generation (one that no longer matches the current owner) or a @@ -963,6 +976,8 @@ pub fn run_omp_blocks_server() -> Result<()> { write_value(&mut stdout, ¬ification_to_wire_value(¬ification)?)?; } + let turn_timeout = turn_timeout_from_trace_context(&trace_context); + let drive_result = drive_omp_turn( child, &mut event_normalizer, @@ -974,6 +989,7 @@ pub fn run_omp_blocks_server() -> Result<()> { &thread_id, &input_rx, &admitted_ownership, + turn_timeout, )?; turn_active.store(false, Ordering::SeqCst); @@ -1546,6 +1562,7 @@ fn drive_omp_turn( thread_id: &str, active_rx: &mpsc::Receiver, admitted_ownership: &Option, + turn_timeout: Duration, ) -> Result { use crate::omp::OmpHarness; use crate::traits::{HarnessServer, NormalizedEvent}; @@ -1558,12 +1575,12 @@ fn drive_omp_turn( let mut prompt_error: Option = None; // Settle window: arm only after a terminal assistant stop. let mut settle_deadline: Option = None; - let absolute_deadline = std::time::Instant::now() + Duration::from_secs(300); + let absolute_deadline = std::time::Instant::now().checked_add(turn_timeout); while !terminal { // #5: check deadlines unconditionally, regardless of frame flow. let now = std::time::Instant::now(); - if now >= absolute_deadline { + if absolute_deadline.is_some_and(|deadline| now >= deadline) { eprintln!("omp rpc: absolute turn timeout, terminating"); let abort_id = child.next_request_id(); child.send_command(&abort_command(&abort_id))?; @@ -2378,6 +2395,49 @@ done assert!(!deadline.child_reusable); } + #[test] + fn configured_max_duration_extends_the_omp_turn_deadline() { + let trace_context = crate::otel::TraceContext { + metadata: std::collections::BTreeMap::from([( + "max_duration_ms".to_owned(), + json!(2_700_000), + )]), + ..Default::default() + }; + + assert_eq!( + turn_timeout_from_trace_context(&trace_context), + Duration::from_millis(2_705_000) + ); + } + + #[test] + fn missing_invalid_or_zero_max_duration_keeps_the_default_deadline() { + assert_eq!( + turn_timeout_from_trace_context(&crate::otel::TraceContext::default()), + DEFAULT_OMP_TURN_TIMEOUT + ); + let invalid = crate::otel::TraceContext { + metadata: std::collections::BTreeMap::from([( + "max_duration_ms".to_owned(), + json!("2700000"), + )]), + ..Default::default() + }; + assert_eq!( + turn_timeout_from_trace_context(&invalid), + DEFAULT_OMP_TURN_TIMEOUT + ); + let zero = crate::otel::TraceContext { + metadata: std::collections::BTreeMap::from([("max_duration_ms".to_owned(), json!(0))]), + ..Default::default() + }; + assert_eq!( + turn_timeout_from_trace_context(&zero), + DEFAULT_OMP_TURN_TIMEOUT + ); + } + #[test] fn read_line_timeout_blank_stream_respects_absolute_deadline() { // Bridge emits ready then continuous blank lines. A naive per-line diff --git a/services/api-rs/crates/centaur-session-runtime/src/lib.rs b/services/api-rs/crates/centaur-session-runtime/src/lib.rs index ea1b13f46..f1702eb1b 100644 --- a/services/api-rs/crates/centaur-session-runtime/src/lib.rs +++ b/services/api-rs/crates/centaur-session-runtime/src/lib.rs @@ -2164,7 +2164,8 @@ impl SessionRuntime { } }; - let trace = SessionTraceContext::new(thread_key, Some(&execution_trace_span)); + let trace = SessionTraceContext::new(thread_key, Some(&execution_trace_span)) + .with_max_duration_ms(max_duration_ms); // Inject the trusted ownership fence (acquired above) so the // harness-server resident OMP host can fence stale/missing // ownership. Only OMP sessions have a generation; non-OMP skip. @@ -8270,6 +8271,9 @@ struct SessionTraceContext { /// on this ownership; a stale or missing fence is rejected. owner_id: Option, generation: Option, + /// Trusted execution duration forwarded to the harness as a local safety + /// deadline. api-rs remains authoritative and records the timeout first. + max_duration_ms: Option, } impl SessionTraceContext { @@ -8279,9 +8283,15 @@ impl SessionTraceContext { traceparent: execution_span.and_then(centaur_telemetry::traceparent_for_span), owner_id: None, generation: None, + max_duration_ms: None, } } + fn with_max_duration_ms(mut self, max_duration_ms: Option) -> Self { + self.max_duration_ms = max_duration_ms; + self + } + /// Attach the trusted session ownership fence. Called after /// `acquire_oneshot_session_ownership` succeeds so the harness-server /// can fence stale/missing ownership without trusting client input. @@ -8353,25 +8363,23 @@ fn input_line_with_session_context( map.entry("traceparent") .or_insert_with(|| Value::String(traceparent.clone())); } - // Inject the trusted session ownership fence into trace_metadata so the - // harness-server resident OMP host can fence stale/missing ownership. - // This is api-rs-injected (from the DB ownership lease), never client- - // asserted: the client never sees owner_id or generation in the input - // line — they are added here, after the line leaves the client boundary. - if let (Some(owner_id), Some(generation)) = (&trace.owner_id, trace.generation) { - // Overwrite trace_metadata unconditionally so a client-supplied - // non-object value (null, string, array) is replaced with a fresh - // object carrying only the trusted owner_id and generation. - // A malicious client cannot bypass the fence with a non-object. + // Inject trusted execution controls into trace_metadata after the line + // leaves the client boundary. Client-supplied values are overwritten. + if trace.owner_id.is_some() || trace.max_duration_ms.is_some() { let mut metadata = match map.get("trace_metadata") { Some(Value::Object(existing)) => existing.clone(), _ => serde_json::Map::new(), }; - metadata.insert("owner_id".to_owned(), Value::String(owner_id.clone())); - metadata.insert( - "generation".to_owned(), - Value::Number(serde_json::Number::from(generation)), - ); + if let (Some(owner_id), Some(generation)) = (&trace.owner_id, trace.generation) { + metadata.insert("owner_id".to_owned(), Value::String(owner_id.clone())); + metadata.insert( + "generation".to_owned(), + Value::Number(serde_json::Number::from(generation)), + ); + } + if let Some(max_duration_ms) = trace.max_duration_ms { + metadata.insert("max_duration_ms".to_owned(), json!(max_duration_ms)); + } map.insert("trace_metadata".to_owned(), Value::Object(metadata)); } merge_session_context(map, session_context_for_thread(thread_key)); @@ -10397,6 +10405,23 @@ mod tests { ); } + #[test] + fn execution_max_duration_is_injected_into_trusted_harness_metadata() { + let thread_key = ThreadKey::parse("workflow:test:repair").unwrap(); + let trace = SessionTraceContext::new(&thread_key, None) + .with_max_duration_ms(Some(2_700_000)) + .with_ownership("api-rs-owner", 7); + let input = r#"{"type":"user","text":"repair","trace_metadata":{"max_duration_ms":1}}"#; + + let value: Value = + serde_json::from_str(&input_line_with_session_context(&thread_key, &trace, input)) + .unwrap(); + + assert_eq!(value["trace_metadata"]["max_duration_ms"], 2_700_000); + assert_eq!(value["trace_metadata"]["owner_id"], "api-rs-owner"); + assert_eq!(value["trace_metadata"]["generation"], 7); + } + #[test] fn execution_metadata_preserves_idle_and_max_duration() { let metadata = @@ -11104,6 +11129,7 @@ mod tests { traceparent: Some("00-0123456789abcdef0123456789abcdef-0123456789abcdef-01".to_owned()), owner_id: None, generation: None, + max_duration_ms: None, }; let line = input_line_with_session_context(