feat(plugin): add timestamped scope guards

crates/core/src/plugin/dynamic/worker.rs

@@ -14,6 +14,7 @@ use std::sync::{Arc, Condvar, Mutex};
 use std::task::{Context, Poll};
 use std::time::Duration;
 
+use chrono::{DateTime, Utc};
 use futures_util::FutureExt;
 use nemo_relay_worker_proto::v1::plugin_worker_client::PluginWorkerClient;
 use nemo_relay_worker_proto::v1::relay_host_runtime_server::{
@@ -3352,6 +3353,7 @@ impl RelayHostRuntime for WorkerHostRuntimeService {
                     .data_opt(optional_envelope_to_json(request.data)?)
                     .metadata_opt(optional_envelope_to_json(request.metadata)?)
                     .input_opt(optional_envelope_to_json(request.input)?)
+                    .timestamp_opt(optional_worker_timestamp(request.timestamp_unix_micros)?)
                     .build(),
             )
         });
@@ -3393,6 +3395,10 @@ impl RelayHostRuntime for WorkerHostRuntimeService {
         let request = request.into_inner();
         self.state
             .authorize(&request.activation_id, &request.auth_token)?;
+        let timestamp = match optional_worker_timestamp(request.timestamp_unix_micros) {
+            Ok(timestamp) => timestamp,
+            Err(err) => return Ok(Response::new(host_ack(Err(err)))),
+        };
         let handle = self
             .state
             .scope_handles
@@ -3408,6 +3414,7 @@ impl RelayHostRuntime for WorkerHostRuntimeService {
                     .handle_uuid(&handle.handle.uuid)
                     .output_opt(output)
                     .metadata_opt(metadata)
+                    .timestamp_opt(timestamp)
                     .build(),
             )
         };
@@ -3931,6 +3938,18 @@ fn optional_envelope_to_json(value: Option<JsonEnvelope>) -> FlowResult<Option<J
         .transpose()
 }
 
+fn optional_worker_timestamp(value: Option<i64>) -> FlowResult<Option<DateTime<Utc>>> {
+    value
+        .map(|timestamp| {
+            DateTime::<Utc>::from_timestamp_micros(timestamp).ok_or_else(|| {
+                FlowError::InvalidArgument(
+                    "timestamp unix microseconds are outside supported range".into(),
+                )
+            })
+        })
+        .transpose()
+}
+
 fn optional_typed_envelope<T: serde::de::DeserializeOwned>(
     value: Option<JsonEnvelope>,
     field: &str,

crates/core/tests/unit/dynamic_worker_tests.rs

@@ -7,7 +7,7 @@ use std::sync::{Arc, Mutex};
 #[cfg(unix)]
 use std::os::unix::fs::PermissionsExt;
 
-use crate::api::event::{BaseEvent, MarkEvent};
+use crate::api::event::{BaseEvent, Event, MarkEvent, ScopeCategory};
 use crate::api::optimization::{
     LlmOptimizationRecorder, record_llm_optimization_contribution, scope_llm_optimization_recorder,
 };
@@ -1722,6 +1722,7 @@ async fn dropping_callback_future_cancels_worker_and_cleans_host_state() {
                 data: None,
                 metadata: None,
                 input: None,
+                timestamp_unix_micros: None,
             }))
             .await
             .expect("worker scope should push")
@@ -2513,6 +2514,7 @@ async fn host_runtime_service_covers_auth_scope_and_ack_errors() {
             }),
             metadata: None,
             input: None,
+            timestamp_unix_micros: None,
         }))
         .await
         .expect("invalid JSON should be structured")
@@ -2524,13 +2526,127 @@ async fn host_runtime_service_covers_auth_scope_and_ack_errors() {
             .contains("invalid JSON")
     );
 
+    let historical_events = Arc::new(Mutex::new(Vec::<Event>::new()));
+    let captured_historical_events = Arc::clone(&historical_events);
+    crate::api::subscriber::register_subscriber(
+        "worker-historical-scope-timestamps",
+        Arc::new(move |event| {
+            if event.name() == "historical-scope" {
+                captured_historical_events
+                    .lock()
+                    .expect("historical events lock")
+                    .push(event.clone());
+            }
+        }),
+    )
+    .expect("historical timestamp subscriber should register");
+
+    let historical_push = service
+        .push_scope(Request::new(PushScopeRequest {
+            activation_id: ACTIVATION_ID.into(),
+            auth_token: AUTH_TOKEN.into(),
+            scope: None,
+            name: "historical-scope".into(),
+            scope_type: ProtoScopeType::Custom as i32,
+            data: None,
+            metadata: None,
+            input: None,
+            timestamp_unix_micros: Some(-2),
+        }))
+        .await
+        .expect("historical scope should push")
+        .into_inner();
+    assert!(historical_push.error.is_none());
+    let historical_handle_id = historical_push.scope_handle_id;
+    assert_eq!(
+        state
+            .scope_handles
+            .lock()
+            .expect("scope handles lock")
+            .get(&historical_handle_id)
+            .expect("historical scope handle")
+            .handle
+            .started_at
+            .timestamp_micros(),
+        -2
+    );
+
+    let invalid_pop = service
+        .pop_scope(Request::new(PopScopeRequest {
+            activation_id: ACTIVATION_ID.into(),
+            auth_token: AUTH_TOKEN.into(),
+            scope_handle_id: historical_handle_id.clone(),
+            output: None,
+            metadata: None,
+            timestamp_unix_micros: Some(i64::MAX),
+        }))
+        .await
+        .expect("invalid timestamp should return a host ack")
+        .into_inner();
+    assert!(!invalid_pop.ok);
+    assert!(
+        invalid_pop
+            .error
+            .expect("invalid timestamp error")
+            .message
+            .contains("outside supported range")
+    );
+    assert!(
+        state
+            .scope_handles
+            .lock()
+            .expect("scope handles lock")
+            .contains_key(&historical_handle_id),
+        "an invalid timestamp must not consume the pop handle"
+    );
+
+    let historical_pop = service
+        .pop_scope(Request::new(PopScopeRequest {
+            activation_id: ACTIVATION_ID.into(),
+            auth_token: AUTH_TOKEN.into(),
+            scope_handle_id: historical_handle_id.clone(),
+            output: None,
+            metadata: None,
+            timestamp_unix_micros: Some(0),
+        }))
+        .await
+        .expect("epoch timestamp should pop")
+        .into_inner();
+    assert!(historical_pop.ok, "{:?}", historical_pop.error);
+    assert!(
+        !state
+            .scope_handles
+            .lock()
+            .expect("scope handles lock")
+            .contains_key(&historical_handle_id)
+    );
+    crate::api::subscriber::flush_subscribers().expect("historical timestamp events should flush");
+    assert!(
+        crate::api::subscriber::deregister_subscriber("worker-historical-scope-timestamps")
+            .expect("historical timestamp subscriber should deregister")
+    );
+    {
+        let historical_events = historical_events.lock().expect("historical events lock");
+        let historical_start = historical_events
+            .iter()
+            .find(|event| event.scope_category() == Some(ScopeCategory::Start))
+            .expect("historical start event");
+        let historical_end = historical_events
+            .iter()
+            .find(|event| event.scope_category() == Some(ScopeCategory::End))
+            .expect("historical end event");
+        assert_eq!(historical_start.timestamp().timestamp_micros(), -2);
+        assert_eq!(historical_end.timestamp().timestamp_micros(), 0);
+    }
+
     let pop_error = service
         .pop_scope(Request::new(PopScopeRequest {
             activation_id: ACTIVATION_ID.into(),
             auth_token: AUTH_TOKEN.into(),
             scope_handle_id: "missing-scope".into(),
             output: None,
             metadata: None,
+            timestamp_unix_micros: None,
         }))
         .await
         .expect_err("missing scope handle should fail");
@@ -2774,6 +2890,7 @@ async fn host_runtime_service_reports_poisoned_internal_locks() {
             data: None,
             metadata: None,
             input: None,
+            timestamp_unix_micros: None,
         }))
         .await
         .expect_err("poisoned scope handle lock should fail");
@@ -2786,6 +2903,7 @@ async fn host_runtime_service_reports_poisoned_internal_locks() {
             scope_handle_id: "missing".into(),
             output: None,
             metadata: None,
+            timestamp_unix_micros: None,
         }))
         .await
         .expect_err("poisoned scope handle lock should fail");

crates/plugin/src/lib.rs

@@ -18,6 +18,7 @@ use std::marker::{PhantomData, PhantomPinned};
 use std::panic::{AssertUnwindSafe, catch_unwind};
 use std::ptr;
 use std::sync::{Arc, Mutex};
+use std::time::{SystemTime, UNIX_EPOCH};
 
 pub use nemo_relay_types::Json;
 pub use nemo_relay_types::api::event::{
@@ -2094,6 +2095,36 @@ impl PluginRuntime {
         })
     }
 
+    /// Opens a scope and records `started_at` on its start event.
+    ///
+    /// This is typed SDK access to the timestamp slot already carried by the
+    /// native host's `scope_push` function. It does not introduce a distinct
+    /// scope event or change the native ABI.
+    pub fn scope_at(
+        &self,
+        name: &str,
+        scope_type: ScopeType,
+        data: Option<&Json>,
+        metadata: Option<&Json>,
+        input: Option<&Json>,
+        started_at: SystemTime,
+    ) -> Result<ScopeGuard<'_>> {
+        let timestamp = unix_micros(started_at)?;
+        let handle = push_scope_with_timestamp(
+            &self.host,
+            name,
+            scope_type.into(),
+            data,
+            metadata,
+            input,
+            Some(timestamp),
+        )?;
+        Ok(ScopeGuard {
+            runtime: self,
+            handle: Some(handle),
+        })
+    }
+
     /// Emits a mark event under the current scope.
     pub fn emit_mark(
         &self,
@@ -2215,7 +2246,8 @@ impl From<ScopeType> for NemoRelayNativeScopeType {
     }
 }
 
-/// RAII guard for a host scope opened by [`PluginRuntime::scope`].
+/// RAII guard for a host scope opened by [`PluginRuntime::scope`] or
+/// [`PluginRuntime::scope_at`].
 ///
 /// A guard may move between threads only while its scope stack is bound on the
 /// destination thread. Async middleware restores that binding around each poll;
@@ -2241,6 +2273,33 @@ impl<'a> ScopeGuard<'a> {
         self.handle.take();
         Ok(())
     }
+
+    /// Pops the scope and records `ended_at` on its end event.
+    ///
+    /// This is typed SDK access to the timestamp slot already carried by the
+    /// native host's `scope_pop` function. The handle remains owned by this
+    /// guard if the host rejects the close, so [`Drop`] can still attempt the
+    /// ordinary cleanup path.
+    pub fn close_at(
+        &mut self,
+        output: Option<&Json>,
+        metadata: Option<&Json>,
+        ended_at: SystemTime,
+    ) -> Result<()> {
+        let Some(handle) = self.handle.as_ref() else {
+            return Ok(());
+        };
+        let timestamp = unix_micros(ended_at)?;
+        pop_scope_with_timestamp(
+            &self.runtime.host,
+            handle,
+            output,
+            metadata,
+            Some(timestamp),
+        )?;
+        self.handle.take();
+        Ok(())
+    }
 }
 
 impl Drop for ScopeGuard<'_> {
@@ -2559,6 +2618,18 @@ pub fn push_scope<'a>(
     data: Option<&Json>,
     metadata: Option<&Json>,
     input: Option<&Json>,
+) -> Result<ScopeHandle<'a>> {
+    push_scope_with_timestamp(host, name, scope_type, data, metadata, input, None)
+}
+
+fn push_scope_with_timestamp<'a>(
+    host: &'a NemoRelayNativeHostApiV1,
+    name: &str,
+    scope_type: NemoRelayNativeScopeType,
+    data: Option<&Json>,
+    metadata: Option<&Json>,
+    input: Option<&Json>,
+    timestamp: Option<i64>,
 ) -> Result<ScopeHandle<'a>> {
     let name =
         HostString::new(host, name).ok_or_else(|| "failed to allocate scope name".to_string())?;
@@ -2575,7 +2646,7 @@ pub fn push_scope<'a>(
             data.as_ptr(),
             metadata.as_ptr(),
             input.as_ptr(),
-            ptr::null(),
+            timestamp.as_ref().map_or(ptr::null(), ptr::from_ref),
             &mut out,
         )
     };
@@ -2592,6 +2663,16 @@ pub fn pop_scope(
     handle: &ScopeHandle<'_>,
     output: Option<&Json>,
     metadata: Option<&Json>,
+) -> Result<()> {
+    pop_scope_with_timestamp(host, handle, output, metadata, None)
+}
+
+fn pop_scope_with_timestamp(
+    host: &NemoRelayNativeHostApiV1,
+    handle: &ScopeHandle<'_>,
+    output: Option<&Json>,
+    metadata: Option<&Json>,
+    timestamp: Option<i64>,
 ) -> Result<()> {
     let output = OptionalHostJson::new(host, output)?;
     let metadata = OptionalHostJson::new(host, metadata)?;
@@ -2600,7 +2681,7 @@ pub fn pop_scope(
             handle.as_ptr(),
             output.as_ptr(),
             metadata.as_ptr(),
-            ptr::null(),
+            timestamp.as_ref().map_or(ptr::null(), ptr::from_ref),
         )
     };
     if status == NemoRelayStatus::Ok {
@@ -2610,6 +2691,22 @@ pub fn pop_scope(
     }
 }
 
+fn unix_micros(timestamp: SystemTime) -> Result<i64> {
+    let micros = match timestamp.duration_since(UNIX_EPOCH) {
+        Ok(duration) => i128::try_from(duration.as_micros()),
+        Err(error) => {
+            let duration = error.duration();
+            i128::try_from(duration.as_micros()).map(|micros| {
+                // `Duration::as_micros` truncates toward zero, but a signed
+                // Unix timestamp must floor pre-epoch sub-microsecond values.
+                -micros - i128::from(duration.subsec_nanos() % 1_000 != 0)
+            })
+        }
+    }
+    .map_err(|_| "scope timestamp exceeds the supported range".to_string())?;
+    i64::try_from(micros).map_err(|_| "scope timestamp exceeds the supported range".to_string())
+}
+
 /// Emits a mark event under the current scope.
 pub fn emit_mark(
     host: &NemoRelayNativeHostApiV1,