diff --git a/letta/schemas/run.py b/letta/schemas/run.py index a03cd49d..c029281f 100644 --- a/letta/schemas/run.py +++ b/letta/schemas/run.py @@ -65,4 +65,5 @@ class RunUpdate(RunBase): completed_at: Optional[datetime] = Field(None, description="The timestamp when the run was completed.") stop_reason: Optional[StopReasonType] = Field(None, description="The reason why the run was stopped.") metadata: Optional[dict] = Field(None, validation_alias="metadata_", description="Additional metadata for the run.") + total_duration_ns: Optional[int] = Field(None, description="Total run duration in nanoseconds") model_config = ConfigDict(extra="ignore") # Ignores extra fields diff --git a/letta/services/run_manager.py b/letta/services/run_manager.py index 44838536..c45b97db 100644 --- a/letta/services/run_manager.py +++ b/letta/services/run_manager.py @@ -206,11 +206,16 @@ class RunManager: async with db_registry.async_session() as session: metrics = await RunMetricsModel.read_async(db_session=session, identifier=run_id, actor=actor) # Calculate runtime if run is completing - if is_terminal_update and metrics.run_start_ns: - import time + if is_terminal_update: + # Use total_duration_ns from RunUpdate if provided + # Otherwise fall back to system time + if update.total_duration_ns is not None: + metrics.run_ns = update.total_duration_ns + elif metrics.run_start_ns: + import time - current_ns = int(time.time() * 1e9) - metrics.run_ns = current_ns - metrics.run_start_ns + current_ns = int(time.time() * 1e9) + metrics.run_ns = current_ns - metrics.run_start_ns metrics.num_steps = num_steps await metrics.update_async(db_session=session, actor=actor, no_commit=True, no_refresh=True) await session.commit()