(MOT-4107) fix(harness): recover calls safely after engine restarts

harness/src/clients/engine.rs

@@ -9,6 +9,12 @@ use iii_sdk::protocol::TriggerRequest;
 use iii_sdk::IIIClient;
 use serde_json::{json, Value};
 
+/// The message a dispatch interrupted by an engine restart resolves with —
+/// it becomes the synthesized error `function_result` that closes the call
+/// in the transcript (iii-hq/workers#507).
+pub const ENGINE_RESTART_INTERRUPTED: &str =
+    "execution interrupted by engine restart; the call ran at most once and its result is unknown";
+
 /// A registry function descriptor (id + optional schemas/description). Only
 /// the fields the harness reads are typed; the rest pass through as `extra`.
 #[derive(Debug, Clone)]
@@ -46,25 +52,51 @@ impl EngineClient {
 
     /// Dispatch an arbitrary iii function and return its raw result. This is
     /// the target invocation of an unwrapped `agent_trigger` call.
+    ///
+    /// An engine restart while the call is in flight fails the dispatch with
+    /// `engine_restart` instead of waiting out the full timeout: the engine's
+    /// in-memory invocation routing died with it, so the result can never be
+    /// delivered, and the caller must close the interrupted call promptly to
+    /// keep the session usable (iii-hq/workers#507).
     pub async fn dispatch(
         &self,
         function_id: &str,
         payload: Value,
     ) -> Result<Value, DispatchError> {
-        self.iii
-            .trigger(TriggerRequest {
-                function_id: function_id.to_string(),
-                payload,
-                action: None,
-                timeout_ms: Some(self.timeout_ms),
-            })
-            .await
-            .map_err(|e| {
-                let raw = e.to_string();
-                let mut parsed = parse_dispatch_error_message(&raw);
-                parsed.message = format!("{function_id}: {}", parsed.message);
-                parsed
-            })
+        // Capture the epoch immediately before the invocation instead of
+        // comparing against process-global state. A global baseline becomes
+        // stale when the engine restarts while this worker is idle and would
+        // falsely interrupt the first slow call made afterwards.
+        let baseline = engine_epoch_ms(&self.iii).await;
+        let call = self.iii.trigger(TriggerRequest {
+            function_id: function_id.to_string(),
+            payload,
+            action: None,
+            timeout_ms: Some(self.timeout_ms),
+        });
+        let result = match baseline {
+            Some(baseline) => {
+                tokio::select! {
+                    result = call => result,
+                    () = engine_link_interrupted(&self.iii, baseline) => {
+                        return Err(DispatchError {
+                            code: Some("engine_restart".to_string()),
+                            message: format!("{function_id}: {ENGINE_RESTART_INTERRUPTED}"),
+                        });
+                    },
+                }
+            }
+            // If the preflight epoch cannot be read, do not guess. The normal
+            // SDK timeout/error is safer than declaring an invocation
+            // interrupted from an untrusted baseline and re-running it.
+            None => call.await,
+        };
+        result.map_err(|e| {
+            let raw = e.to_string();
+            let mut parsed = parse_dispatch_error_message(&raw);
+            parsed.message = format!("{function_id}: {}", parsed.message);
+            parsed
+        })
     }
 
     /// List registry function descriptors (best-effort; empty on failure).
@@ -130,6 +162,66 @@ impl EngineClient {
     }
 }
 
+/// How often an in-flight dispatch samples the engine epoch, and how long
+/// one sample may take. Sampling is cheap relative to the dispatched calls
+/// it guards; dispatches shorter than the first interval never even probe.
+const ENGINE_EPOCH_PROBE_INTERVAL_MS: u64 = 1_000;
+const ENGINE_EPOCH_PROBE_TIMEOUT_MS: u64 = 3_000;
+
+/// The engine's boot identity: the earliest `connected_at_ms` among its
+/// in-process (`runtime == "engine"`) workers, which attach once at engine
+/// startup. A restarted engine reports a new epoch — the signal that every
+/// invocation in flight across the restart lost its result routing. `None`
+/// when the engine cannot be reached (outage in progress) or the response
+/// shape is unrecognized.
+async fn engine_epoch_ms(iii: &IIIClient) -> Option<u64> {
+    let response = iii
+        .trigger(TriggerRequest {
+            function_id: "engine::workers::list".to_string(),
+            payload: json!({}),
+            action: None,
+            timeout_ms: Some(ENGINE_EPOCH_PROBE_TIMEOUT_MS),
+        })
+        .await
+        .ok()?;
+    parse_engine_epoch(&response)
+}
+
+fn parse_engine_epoch(response: &Value) -> Option<u64> {
+    response
+        .get("workers")?
+        .as_array()?
+        .iter()
+        .filter(|worker| worker.get("runtime").and_then(Value::as_str) == Some("engine"))
+        .filter_map(|worker| worker.get("connected_at_ms").and_then(Value::as_u64))
+        .min()
+}
+
+/// Pend until the engine under an in-flight dispatch is observed to have
+/// RESTARTED (its boot epoch changed). Neither the SDK connection state nor
+/// plain liveness probes can see a fast restart: the reconnect loop reports
+/// `Connected` through its silent 2s retry sleep, and outbound messages
+/// buffered during the outage are answered by the NEW engine as if nothing
+/// happened (both verified against a SIGKILLed-and-respawned engine). An
+/// epoch read that succeeds with a changed value doubles as proof the
+/// engine is answering again, so the caller can immediately persist the
+/// synthesized "interrupted" result without tripping over the same outage.
+async fn engine_link_interrupted(iii: &IIIClient, baseline: u64) {
+    loop {
+        tokio::time::sleep(std::time::Duration::from_millis(
+            ENGINE_EPOCH_PROBE_INTERVAL_MS,
+        ))
+        .await;
+        // An unreadable epoch (outage in progress) never trips by itself:
+        // the first successful sample afterwards decides.
+        if let Some(epoch) = engine_epoch_ms(iii).await {
+            if epoch != baseline {
+                return;
+            }
+        }
+    }
+}
+
 fn parse_dispatch_error_message(raw: &str) -> DispatchError {
     if let Some(parsed) = parse_dispatch_error_value(raw) {
         return parsed;
@@ -253,6 +345,20 @@ fn descriptor_of(v: &Value) -> Option<FunctionDescriptor> {
 mod tests {
     use super::*;
 
+    #[test]
+    fn engine_epoch_uses_the_oldest_in_process_engine_worker() {
+        let response = json!({
+            "workers": [
+                { "runtime": "rust", "connected_at_ms": 1 },
+                { "runtime": "engine", "connected_at_ms": 30 },
+                { "runtime": "engine", "connected_at_ms": 20 },
+                { "runtime": "engine", "connected_at_ms": "invalid" }
+            ]
+        });
+        assert_eq!(parse_engine_epoch(&response), Some(20));
+        assert_eq!(parse_engine_epoch(&json!({ "workers": [] })), None);
+    }
+
     #[test]
     fn dispatch_error_parser_extracts_plain_json_code_and_message() {
         let err = parse_dispatch_error_message(

harness/src/functions/turn.rs

@@ -74,13 +74,62 @@ pub async fn handle(deps: &Deps, payload: TurnStepPayload) -> Result<TurnStepRes
     .await
 }
 
+/// Bounded in-place retry budget for step errors that mean a dependency is
+/// not (yet) back after an engine restart, not that the turn is broken: a
+/// restored/redelivered step can beat `context::assemble` (or another
+/// required worker) re-registering against the restarted engine, and failing
+/// the turn terminally on that boot race poisoned an otherwise recoverable
+/// session (iii-hq/workers#507). Steps are at-least-once by design, so
+/// re-running one wholesale is safe.
+const TRANSIENT_STEP_RETRIES: u32 = 10;
+const TRANSIENT_STEP_BACKOFF_MS: u64 = 500;
+
+/// Whether a step error looks like a transient dependency outage (an engine
+/// restart's registration race or a not-yet-reconnected SDK) rather than a
+/// real failure. Enqueue exhaustion is excluded: `enqueue_step` already
+/// retries internally, and re-running a step whose record already advanced
+/// would ack as stale and silently wedge the turn instead of failing it.
+fn is_transient_step_error(error: &HarnessError) -> bool {
+    let message = match error {
+        HarnessError::Dependency(m) => m.to_ascii_lowercase(),
+        // Retrying a read from the top of the step is safe. Do not classify
+        // state writes broadly: a lost acknowledgement after a successful
+        // write must not cause the whole step to repeat side effects.
+        HarnessError::State(m) if m.to_ascii_lowercase().starts_with("state::get ") => {
+            m.to_ascii_lowercase()
+        }
+        _ => return false,
+    };
+    if message.contains("enqueue harness::turn") {
+        return false;
+    }
+    message.contains("function_not_found") || message.contains("not connected")
+}
+
 async fn run(deps: &Deps, payload: TurnStepPayload) -> Result<TurnStepResult, HarnessError> {
     let (session_id, turn_id) = (payload.session_id.clone(), payload.turn_id.clone());
-    let result = match turn_loop::run_step(deps, payload).await {
-        Ok(result) => result,
-        Err(e) => {
-            tracing::error!(session_id = %session_id, turn_id = %turn_id, error = %e, "turn step failed; finalising turn as failed");
-            turn_loop::fail_turn(deps, &session_id, &turn_id, &e.to_string()).await
+    let mut transient_attempts = 0u32;
+    let result = loop {
+        match turn_loop::run_step(deps, payload.clone()).await {
+            Ok(result) => break result,
+            Err(e)
+                if transient_attempts < TRANSIENT_STEP_RETRIES && is_transient_step_error(&e) =>
+            {
+                transient_attempts += 1;
+                tracing::warn!(
+                    session_id = %session_id,
+                    turn_id = %turn_id,
+                    error = %e,
+                    attempt = transient_attempts,
+                    "turn step hit a transient dependency outage (engine restart boot race); retrying in place"
+                );
+                tokio::time::sleep(std::time::Duration::from_millis(TRANSIENT_STEP_BACKOFF_MS))
+                    .await;
+            }
+            Err(e) => {
+                tracing::error!(session_id = %session_id, turn_id = %turn_id, error = %e, "turn step failed; finalising turn as failed");
+                break turn_loop::fail_turn(deps, &session_id, &turn_id, &e.to_string()).await;
+            }
         }
     };
     record_step_status(&result);
@@ -106,3 +155,39 @@ fn record_step_status(result: &TurnStepResult) {
         span.set_status(Status::error("harness turn failed"));
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn boot_race_errors_are_transient_and_real_failures_are_not() {
+        // The issue-507 boot race: a restored step reaching a dependency the
+        // restarted engine has not re-registered yet.
+        assert!(is_transient_step_error(&HarnessError::Dependency(
+            "context::assemble: remote error (function_not_found): Function not found".into()
+        )));
+        assert!(is_transient_step_error(&HarnessError::State(
+            "state::get harness_turn/s_1: iii is not connected".into()
+        )));
+        assert!(!is_transient_step_error(&HarnessError::State(
+            "state::set harness_turn/s_1: iii is not connected".into()
+        )));
+        // Case-insensitive on the engine's code spelling.
+        assert!(is_transient_step_error(&HarnessError::Dependency(
+            "session::messages: remote error (FUNCTION_NOT_FOUND): unknown".into()
+        )));
+        // Enqueue exhaustion must fail the turn, not retry into a stale-ack
+        // wedge (enqueue_step already retried internally).
+        assert!(!is_transient_step_error(&HarnessError::Dependency(
+            "enqueue harness::turn: remote error (function_not_found): Function not found".into()
+        )));
+        // Ordinary failures stay terminal.
+        assert!(!is_transient_step_error(&HarnessError::Dependency(
+            "router::chat: provider rejected the request".into()
+        )));
+        assert!(!is_transient_step_error(&HarnessError::Internal(
+            "function_not_found mentioned in an internal error".into()
+        )));
+    }
+}

harness/src/queue.rs

@@ -69,7 +69,8 @@ fn turn_queue_definition() -> Value {
             "concurrency": 10,
             "max_retries": 3,
             "backoff_ms": 1_000,
-            "poll_interval_ms": 100
+            "poll_interval_ms": 100,
+            "redeliver_on_engine_restart": true
         }
     })
 }
@@ -90,7 +91,8 @@ mod tests {
                     "concurrency": 10,
                     "max_retries": 3,
                     "backoff_ms": 1_000,
-                    "poll_interval_ms": 100
+                    "poll_interval_ms": 100,
+                    "redeliver_on_engine_restart": true
                 }
             })
         );

harness/src/turn_loop.rs

@@ -117,6 +117,15 @@ pub struct TurnStepResult {
     pub skipped: bool,
 }
 
+/// Bounded in-place retry for the enqueue that follows a persisted step
+/// advance: `queue::enqueue` can be briefly unregistered while the queue
+/// worker replays its registrations against a restarted engine, and failing
+/// the turn in that window (or worse, wedging it Running with the step
+/// persisted but never enqueued) turns a recoverable restart into a dead
+/// session (iii-hq/workers#507).
+const ENQUEUE_ATTEMPTS: u32 = 5;
+const ENQUEUE_RETRY_BACKOFF_MS: u64 = 500;
+
 /// Enqueue the next durable loop step onto the dedicated `harness-turn` queue.
 pub async fn enqueue_step(
     iii: &IIIClient,
@@ -131,17 +140,40 @@ pub async fn enqueue_step(
     if let Some(preview) = message_preview {
         payload["message_preview"] = json!(preview);
     }
-    iii.trigger(TriggerRequest {
-        function_id: "harness::turn".to_string(),
-        payload,
-        action: Some(TriggerAction::Enqueue {
-            queue: TURN_QUEUE.to_string(),
-        }),
-        timeout_ms: None,
-    })
-    .await
-    .map(|_| ())
-    .map_err(|e| HarnessError::Dependency(format!("enqueue harness::turn: {e}")))
+    let mut last_error = String::new();
+    for attempt in 1..=ENQUEUE_ATTEMPTS {
+        match iii
+            .trigger(TriggerRequest {
+                function_id: "harness::turn".to_string(),
+                payload: payload.clone(),
+                action: Some(TriggerAction::Enqueue {
+                    queue: TURN_QUEUE.to_string(),
+                }),
+                timeout_ms: None,
+            })
+            .await
+        {
+            Ok(_) => return Ok(()),
+            Err(e) => {
+                last_error = e.to_string();
+                if attempt < ENQUEUE_ATTEMPTS {
+                    tracing::warn!(
+                        session_id = %session_id,
+                        turn_id = %turn_id,
+                        step,
+                        attempt,
+                        error = %last_error,
+                        "enqueue harness::turn failed; retrying"
+                    );
+                    tokio::time::sleep(std::time::Duration::from_millis(ENQUEUE_RETRY_BACKOFF_MS))
+                        .await;
+                }
+            }
+        }
+    }
+    Err(HarnessError::Dependency(format!(
+        "enqueue harness::turn: {last_error}"
+    )))
 }
 
 fn origin(turn_id: &str) -> Value {
@@ -190,6 +222,10 @@ pub async fn run_step(
         match crate::state::get_turn(&deps.iii, &payload.session_id, cfg.session_timeout_ms).await?
         {
             Some(r) => r,
+            // The turn record is the authoritative recovery snapshot. A
+            // transcript alone cannot recover budgets, parent linkage, output
+            // contracts, or dispatch policy safely, so an absent record stays
+            // a stale delivery and is acknowledged without fabricating state.
             None => return Ok(skipped(&payload.session_id)),
         };