From 7362e1d6f55e576b92ce1ce08ff901a69cb63e4a Mon Sep 17 00:00:00 2001 From: serrrfirat Date: Wed, 1 Jul 2026 10:38:21 +0300 Subject: [PATCH 1/2] feat(observability): trace agent loop executor stages --- Cargo.lock | 1 + crates/ironclaw_agent_loop/Cargo.toml | 1 + crates/ironclaw_agent_loop/src/executor.rs | 10 + .../src/executor/canonical.rs | 284 +++++++++++------- .../src/executor/latency.rs | 75 +++++ 5 files changed, 266 insertions(+), 105 deletions(-) create mode 100644 crates/ironclaw_agent_loop/src/executor/latency.rs diff --git a/Cargo.lock b/Cargo.lock index 789c5195af5..f5774fa0398 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4172,6 +4172,7 @@ dependencies = [ "ironclaw_agent_loop", "ironclaw_common", "ironclaw_host_api", + "ironclaw_observability", "ironclaw_turns", "serde", "serde_jcs", diff --git a/crates/ironclaw_agent_loop/Cargo.toml b/crates/ironclaw_agent_loop/Cargo.toml index cd9cb48e865..4075cf6b8eb 100644 --- a/crates/ironclaw_agent_loop/Cargo.toml +++ b/crates/ironclaw_agent_loop/Cargo.toml @@ -19,6 +19,7 @@ async-trait = "0.1" blake3 = "1" ironclaw_common = { path = "../ironclaw_common", version = "0.4.2" } ironclaw_host_api = { path = "../ironclaw_host_api", version = "0.1.0" } +ironclaw_observability = { path = "../ironclaw_observability" } ironclaw_turns = { path = "../ironclaw_turns", version = "0.1.0" } serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/crates/ironclaw_agent_loop/src/executor.rs b/crates/ironclaw_agent_loop/src/executor.rs index 6fdf9140a69..68c102e2eaa 100644 --- a/crates/ironclaw_agent_loop/src/executor.rs +++ b/crates/ironclaw_agent_loop/src/executor.rs @@ -11,6 +11,7 @@ mod checkpoint; mod exit_helpers; mod gates; mod input; +mod latency; mod loop_exit; mod mapping; mod model; @@ -190,6 +191,15 @@ enum TurnCompletedStep { Exit(LoopExit), } +impl TurnCompletedStep { + fn iteration(&self) -> u32 { + match self { + Self::Continue { state, .. } => state.iteration, + Self::Exit(_) => 0, + } + } +} + #[derive(Debug, Default)] struct PendingInputAck { tokens: Vec, diff --git a/crates/ironclaw_agent_loop/src/executor/canonical.rs b/crates/ironclaw_agent_loop/src/executor/canonical.rs index 77bb494cc9b..d2d88325b1c 100644 --- a/crates/ironclaw_agent_loop/src/executor/canonical.rs +++ b/crates/ironclaw_agent_loop/src/executor/canonical.rs @@ -12,7 +12,7 @@ use super::{ DrainInput, ExecutorStage, ExitInput, InputStep, ModelInput, ModelStep, PendingInputAck, PromptInput, PromptStep, ReplyAdmissionInput, ReplyAdmissionStep, StageContext, StopInput, StopObservationInput, StopObservationStep, StopStep, TurnCompletedStep, - UserFacingInputDrainMode, + UserFacingInputDrainMode, latency, }; impl DefaultExecutorPipeline { @@ -26,23 +26,43 @@ impl DefaultExecutorPipeline { let ctx = StageContext { planner, host }; let mut pending_input_ack = PendingInputAck::default(); + macro_rules! trace_stage { + ($operation:expr, $iteration:expr, $future:expr) => {{ + let iteration = $iteration; + let started_at = latency::started_at(); + let result = $future.await; + latency::result( + $operation, + host.run_context(), + iteration, + started_at, + &result, + ); + result + }}; + } + loop { - state = match CheckpointStage.cancel_if_requested(ctx, state).await? { + state = match trace_stage!( + "cancel_check", + state.iteration, + CheckpointStage.cancel_if_requested(ctx, state) + )? { CancelCheck::Continue(state) => *state, CancelCheck::Exit(exit) => return Ok(exit), }; - match self - .budget - .process( + match trace_stage!( + "budget", + state.iteration, + self.budget.process( ctx, BudgetInput { state, pending_input_ack: std::mem::take(&mut pending_input_ack), }, ) - .await? - { + )? { BudgetStep::Continue { state: next, pending_input_ack: ack, @@ -53,6 +73,7 @@ impl DefaultExecutorPipeline { BudgetStep::Exit(exit) => return Ok(exit), } + let progress_started_at = latency::started_at(); CheckpointStage .emit_progress( ctx, @@ -61,10 +82,17 @@ impl DefaultExecutorPipeline { }, ) .await; + latency::operation_ok( + "emit_iteration_started", + host.run_context(), + state.iteration, + progress_started_at, + ); - match self - .input - .process( + match trace_stage!( + "input_drain_steering", + state.iteration, + self.input.process( ctx, DrainInput { state, @@ -72,8 +100,7 @@ impl DefaultExecutorPipeline { mode: UserFacingInputDrainMode::Steering, }, ) - .await? - { + )? { InputStep::Continue { state: next, pending_input_ack: ack, @@ -85,17 +112,17 @@ impl DefaultExecutorPipeline { InputStep::Exit(exit) => return Ok(exit), } - match self - .prompt - .process( + match trace_stage!( + "prompt", + state.iteration, + self.prompt.process( ctx, PromptInput { state, pending_input_ack: std::mem::take(&mut pending_input_ack), }, ) - .await? - { + )? { PromptStep::Exit(exit) => return Ok(exit), PromptStep::Prepared(prompt) => { @@ -103,18 +130,21 @@ impl DefaultExecutorPipeline { state = prompt.state; pending_input_ack = prompt.pending_input_ack; - state = CheckpointStage - .process( + state = trace_stage!( + "checkpoint_before_model", + state.iteration, + CheckpointStage.process( ctx, CheckpointInput { state, kind: CheckpointKind::BeforeModel, }, ) - .await? - .state; + )? + .state; if prompt.rendered_repeated_call_warning { state.stop_state.mark_repeated_call_warning_rendered(); + let note_started_at = latency::started_at(); CheckpointStage .emit_progress( ctx, @@ -127,12 +157,23 @@ impl DefaultExecutorPipeline { })?, ) .await; + latency::operation_ok( + "emit_repeated_call_warning", + host.run_context(), + state.iteration, + note_started_at, + ); } - pending_input_ack.ack(host).await?; + trace_stage!( + "ack_pending_input_before_model", + state.iteration, + pending_input_ack.ack(host) + )?; - let model_response = match self - .model - .process( + let model_response = match trace_stage!( + "model", + state.iteration, + self.model.process( ctx, ModelInput { state, @@ -141,8 +182,7 @@ impl DefaultExecutorPipeline { capability_view: prompt.capability_view, }, ) - .await? - { + )? { ModelStep::Response(next, response) => { state = *next; response @@ -164,23 +204,24 @@ impl DefaultExecutorPipeline { let response_usage = model_response.usage; let completed = match model_response.output { ParentLoopOutput::AssistantReply(reply) => { - match self - .reply_admission - .process(ctx, ReplyAdmissionInput { state, reply }) - .await? - { - ReplyAdmissionStep::Accept { state, reply } => { - self.assistant_reply - .process( - ctx, - AssistantReplyInput { - state: *state, - reply, - usage: response_usage, - }, - ) - .await? - } + match trace_stage!( + "reply_admission", + state.iteration, + self.reply_admission + .process(ctx, ReplyAdmissionInput { state, reply }) + )? { + ReplyAdmissionStep::Accept { state, reply } => trace_stage!( + "assistant_reply", + state.iteration, + self.assistant_reply.process( + ctx, + AssistantReplyInput { + state: *state, + reply, + usage: response_usage, + }, + ) + )?, ReplyAdmissionStep::Reject { state } => { TurnCompletedStep::Continue { state, @@ -189,21 +230,25 @@ impl DefaultExecutorPipeline { } } } - ParentLoopOutput::CapabilityCalls(calls) => { - self.capabilities - .process( - ctx, - CapabilityInput { - state, - surface: prompt.surface, - calls, - }, - ) - .await? - } + ParentLoopOutput::CapabilityCalls(calls) => trace_stage!( + "capabilities", + state.iteration, + self.capabilities.process( + ctx, + CapabilityInput { + state, + surface: prompt.surface, + calls, + }, + ) + )?, }; - let completed = self.post_capability.process(ctx, completed).await?; + let completed = trace_stage!( + "post_capability", + completed.iteration(), + self.post_capability.process(ctx, completed) + )?; let (next_state, summary) = match completed { TurnCompletedStep::Continue { state, summary } => (*state, summary), @@ -211,17 +256,17 @@ impl DefaultExecutorPipeline { }; let completed_kind = summary.kind; - let (mut next_state, summary) = match self - .stop - .observe( + let (mut next_state, summary) = match trace_stage!( + "stop_observe", + next_state.iteration, + self.stop.observe( ctx, StopObservationInput { state: next_state, summary, }, ) - .await? - { + )? { StopObservationStep::Continue { state, summary } => (*state, summary), StopObservationStep::Exit(exit) => return Ok(exit), }; @@ -231,9 +276,10 @@ impl DefaultExecutorPipeline { iteration = next_state.iteration, "agent loop checking follow-up input after reply-only turn end" ); - match self - .input - .process( + match trace_stage!( + "input_drain_follow_up", + next_state.iteration, + self.input.process( ctx, DrainInput { state: next_state, @@ -241,8 +287,7 @@ impl DefaultExecutorPipeline { mode: UserFacingInputDrainMode::FollowUp, }, ) - .await? - { + )? { InputStep::Continue { state: next, pending_input_ack: ack, @@ -264,9 +309,10 @@ impl DefaultExecutorPipeline { } } - match self - .stop - .decide( + match trace_stage!( + "stop_decide", + next_state.iteration, + self.stop.decide( ctx, StopInput { state: next_state, @@ -274,15 +320,18 @@ impl DefaultExecutorPipeline { pending_input_ack: std::mem::take(&mut pending_input_ack), }, ) - .await? - { + )? { StopStep::Stop { state, kind, pending_input_ack: mut ack, } => { - let exit = self.exit.process(ctx, ExitInput { state, kind }).await?; - ack.ack(host).await?; + let exit = trace_stage!( + "exit", + state.iteration, + self.exit.process(ctx, ExitInput { state, kind }) + )?; + trace_stage!("ack_pending_input_before_exit", 0, ack.ack(host))?; return Ok(exit); } StopStep::Continue { @@ -303,10 +352,15 @@ impl DefaultExecutorPipeline { | PromptStep::ResumeExternalTool(resume) => { let resume = *resume; pending_input_ack = resume.pending_input_ack; - pending_input_ack.ack(host).await?; - let completed = self - .capabilities - .process( + trace_stage!( + "ack_pending_input_before_resume_capability", + resume.state.iteration, + pending_input_ack.ack(host) + )?; + let completed = trace_stage!( + "capabilities_resume", + resume.state.iteration, + self.capabilities.process( ctx, CapabilityInput { state: resume.state, @@ -314,33 +368,38 @@ impl DefaultExecutorPipeline { calls: vec![resume.call], }, ) - .await?; + )?; - let completed = self.post_capability.process(ctx, completed).await?; + let completed = trace_stage!( + "post_capability_resume", + completed.iteration(), + self.post_capability.process(ctx, completed) + )?; let (next_state, summary) = match completed { TurnCompletedStep::Continue { state, summary } => (*state, summary), TurnCompletedStep::Exit(exit) => return Ok(exit), }; - let (next_state, summary) = match self - .stop - .observe( + let (next_state, summary) = match trace_stage!( + "stop_observe_resume", + next_state.iteration, + self.stop.observe( ctx, StopObservationInput { state: next_state, summary, }, ) - .await? - { + )? { StopObservationStep::Continue { state, summary } => (*state, summary), StopObservationStep::Exit(exit) => return Ok(exit), }; - match self - .stop - .decide( + match trace_stage!( + "stop_decide_resume", + next_state.iteration, + self.stop.decide( ctx, StopInput { state: next_state, @@ -348,15 +407,18 @@ impl DefaultExecutorPipeline { pending_input_ack: std::mem::take(&mut pending_input_ack), }, ) - .await? - { + )? { StopStep::Stop { state, kind, pending_input_ack: mut ack, } => { - let exit = self.exit.process(ctx, ExitInput { state, kind }).await?; - ack.ack(host).await?; + let exit = trace_stage!( + "exit_resume", + state.iteration, + self.exit.process(ctx, ExitInput { state, kind }) + )?; + trace_stage!("ack_pending_input_before_exit_resume", 0, ack.ack(host))?; return Ok(exit); } StopStep::Continue { @@ -388,20 +450,24 @@ impl DefaultExecutorPipeline { pending_input_ack = ack; // Deliver the ack before stop.observe, mirroring the timing // of the Prepared path (line ~133: ack before ModelStage). - pending_input_ack.ack(host).await?; + trace_stage!( + "ack_pending_input_skip_model", + skipped_state.iteration, + pending_input_ack.ack(host) + )?; let summary = crate::strategies::TurnSummary::compaction_only(); - let (mut next_state, summary) = match self - .stop - .observe( + let (mut next_state, summary) = match trace_stage!( + "stop_observe_skip_model", + skipped_state.iteration, + self.stop.observe( ctx, StopObservationInput { state: skipped_state, summary, }, ) - .await? - { + )? { StopObservationStep::Continue { state, summary } => (*state, summary), StopObservationStep::Exit(exit) => return Ok(exit), }; @@ -409,9 +475,10 @@ impl DefaultExecutorPipeline { // Follow-up drain is skipped: CompactionOnly != ReplyOnly, // so the condition is naturally false here. - match self - .stop - .decide( + match trace_stage!( + "stop_decide_skip_model", + next_state.iteration, + self.stop.decide( ctx, StopInput { state: next_state, @@ -419,15 +486,22 @@ impl DefaultExecutorPipeline { pending_input_ack: std::mem::take(&mut pending_input_ack), }, ) - .await? - { + )? { StopStep::Stop { state, kind, pending_input_ack: mut ack, } => { - let exit = self.exit.process(ctx, ExitInput { state, kind }).await?; - ack.ack(host).await?; + let exit = trace_stage!( + "exit_skip_model", + state.iteration, + self.exit.process(ctx, ExitInput { state, kind }) + )?; + trace_stage!( + "ack_pending_input_before_exit_skip_model", + 0, + ack.ack(host) + )?; return Ok(exit); } StopStep::Continue { diff --git a/crates/ironclaw_agent_loop/src/executor/latency.rs b/crates/ironclaw_agent_loop/src/executor/latency.rs new file mode 100644 index 00000000000..93e70ef12b4 --- /dev/null +++ b/crates/ironclaw_agent_loop/src/executor/latency.rs @@ -0,0 +1,75 @@ +use std::time::Instant; + +use ironclaw_observability::live_latency_started_at; +use ironclaw_turns::run_profile::LoopRunContext; + +pub(super) fn started_at() -> Option { + live_latency_started_at() +} + +pub(super) fn operation_ok( + operation: &'static str, + context: &LoopRunContext, + iteration: u32, + started_at: Option, +) { + let Some(started_at) = started_at else { + return; + }; + + ironclaw_observability::live_latency_trace_ok!( + "agent_loop_executor", + operation, + Some(started_at), + tenant_id = %context.scope.tenant_id, + agent_id = context.scope.agent_id.as_ref().map(|id| id.as_str()).unwrap_or(""), + project_id = context.scope.project_id.as_ref().map(|id| id.as_str()).unwrap_or(""), + thread_id = %context.thread_id, + owner_user_id = context.scope.explicit_owner_user_id().map(|id| id.as_str()).unwrap_or(""), + run_id = %context.run_id, + turn_id = %context.turn_id, + iteration, + "agent loop executor operation completed", + ); +} + +pub(super) fn operation_error( + operation: &'static str, + context: &LoopRunContext, + iteration: u32, + started_at: Option, + _error: &E, +) { + let Some(started_at) = started_at else { + return; + }; + + ironclaw_observability::live_latency_trace_error!( + "agent_loop_executor", + operation, + Some(started_at), + "executor_error", + tenant_id = %context.scope.tenant_id, + agent_id = context.scope.agent_id.as_ref().map(|id| id.as_str()).unwrap_or(""), + project_id = context.scope.project_id.as_ref().map(|id| id.as_str()).unwrap_or(""), + thread_id = %context.thread_id, + owner_user_id = context.scope.explicit_owner_user_id().map(|id| id.as_str()).unwrap_or(""), + run_id = %context.run_id, + turn_id = %context.turn_id, + iteration, + "agent loop executor operation failed", + ); +} + +pub(super) fn result( + operation: &'static str, + context: &LoopRunContext, + iteration: u32, + started_at: Option, + result: &Result, +) { + match result { + Ok(_) => operation_ok(operation, context, iteration, started_at), + Err(error) => operation_error(operation, context, iteration, started_at, error), + } +} From 1880fd8594759fc62e73cfbe638a2f07d52c9c18 Mon Sep 17 00:00:00 2001 From: serrrfirat Date: Wed, 1 Jul 2026 10:50:42 +0300 Subject: [PATCH 2/2] refactor(agent-loop): address executor trace review comments --- crates/ironclaw_agent_loop/src/executor.rs | 4 +- .../src/executor/canonical.rs | 187 ++++++++++-------- .../src/executor/latency.rs | 18 ++ 3 files changed, 126 insertions(+), 83 deletions(-) diff --git a/crates/ironclaw_agent_loop/src/executor.rs b/crates/ironclaw_agent_loop/src/executor.rs index 68c102e2eaa..3fca7f57dba 100644 --- a/crates/ironclaw_agent_loop/src/executor.rs +++ b/crates/ironclaw_agent_loop/src/executor.rs @@ -192,10 +192,10 @@ enum TurnCompletedStep { } impl TurnCompletedStep { - fn iteration(&self) -> u32 { + fn iteration_or(&self, fallback: u32) -> u32 { match self { Self::Continue { state, .. } => state.iteration, - Self::Exit(_) => 0, + Self::Exit(_) => fallback, } } } diff --git a/crates/ironclaw_agent_loop/src/executor/canonical.rs b/crates/ironclaw_agent_loop/src/executor/canonical.rs index d2d88325b1c..9a6d846d0de 100644 --- a/crates/ironclaw_agent_loop/src/executor/canonical.rs +++ b/crates/ironclaw_agent_loop/src/executor/canonical.rs @@ -26,34 +26,20 @@ impl DefaultExecutorPipeline { let ctx = StageContext { planner, host }; let mut pending_input_ack = PendingInputAck::default(); - macro_rules! trace_stage { - ($operation:expr, $iteration:expr, $future:expr) => {{ - let iteration = $iteration; - let started_at = latency::started_at(); - let result = $future.await; - latency::result( - $operation, - host.run_context(), - iteration, - started_at, - &result, - ); - result - }}; - } - loop { - state = match trace_stage!( + state = match latency::stage!( "cancel_check", + host.run_context(), state.iteration, - CheckpointStage.cancel_if_requested(ctx, state) + CheckpointStage.cancel_if_requested(ctx, state), )? { CancelCheck::Continue(state) => *state, CancelCheck::Exit(exit) => return Ok(exit), }; - match trace_stage!( + match latency::stage!( "budget", + host.run_context(), state.iteration, self.budget.process( ctx, @@ -61,7 +47,7 @@ impl DefaultExecutorPipeline { state, pending_input_ack: std::mem::take(&mut pending_input_ack), }, - ) + ), )? { BudgetStep::Continue { state: next, @@ -89,8 +75,9 @@ impl DefaultExecutorPipeline { progress_started_at, ); - match trace_stage!( + match latency::stage!( "input_drain_steering", + host.run_context(), state.iteration, self.input.process( ctx, @@ -99,7 +86,7 @@ impl DefaultExecutorPipeline { pending_input_ack: std::mem::take(&mut pending_input_ack), mode: UserFacingInputDrainMode::Steering, }, - ) + ), )? { InputStep::Continue { state: next, @@ -112,8 +99,9 @@ impl DefaultExecutorPipeline { InputStep::Exit(exit) => return Ok(exit), } - match trace_stage!( + match latency::stage!( "prompt", + host.run_context(), state.iteration, self.prompt.process( ctx, @@ -121,7 +109,7 @@ impl DefaultExecutorPipeline { state, pending_input_ack: std::mem::take(&mut pending_input_ack), }, - ) + ), )? { PromptStep::Exit(exit) => return Ok(exit), @@ -130,8 +118,9 @@ impl DefaultExecutorPipeline { state = prompt.state; pending_input_ack = prompt.pending_input_ack; - state = trace_stage!( + state = latency::stage!( "checkpoint_before_model", + host.run_context(), state.iteration, CheckpointStage.process( ctx, @@ -139,7 +128,7 @@ impl DefaultExecutorPipeline { state, kind: CheckpointKind::BeforeModel, }, - ) + ), )? .state; if prompt.rendered_repeated_call_warning { @@ -164,14 +153,16 @@ impl DefaultExecutorPipeline { note_started_at, ); } - trace_stage!( + latency::stage!( "ack_pending_input_before_model", + host.run_context(), state.iteration, - pending_input_ack.ack(host) + pending_input_ack.ack(host), )?; - let model_response = match trace_stage!( + let model_response = match latency::stage!( "model", + host.run_context(), state.iteration, self.model.process( ctx, @@ -181,7 +172,7 @@ impl DefaultExecutorPipeline { surface_version: prompt.surface.version.clone(), capability_view: prompt.capability_view, }, - ) + ), )? { ModelStep::Response(next, response) => { state = *next; @@ -202,17 +193,20 @@ impl DefaultExecutorPipeline { // turns in a row. `None` is "unknown" and must NOT count as // a zero-output turn against the detector. let response_usage = model_response.usage; + let turn_iteration = state.iteration; let completed = match model_response.output { ParentLoopOutput::AssistantReply(reply) => { - match trace_stage!( + match latency::stage!( "reply_admission", - state.iteration, + host.run_context(), + turn_iteration, self.reply_admission - .process(ctx, ReplyAdmissionInput { state, reply }) + .process(ctx, ReplyAdmissionInput { state, reply }), )? { - ReplyAdmissionStep::Accept { state, reply } => trace_stage!( + ReplyAdmissionStep::Accept { state, reply } => latency::stage!( "assistant_reply", - state.iteration, + host.run_context(), + turn_iteration, self.assistant_reply.process( ctx, AssistantReplyInput { @@ -220,7 +214,7 @@ impl DefaultExecutorPipeline { reply, usage: response_usage, }, - ) + ), )?, ReplyAdmissionStep::Reject { state } => { TurnCompletedStep::Continue { @@ -230,9 +224,10 @@ impl DefaultExecutorPipeline { } } } - ParentLoopOutput::CapabilityCalls(calls) => trace_stage!( + ParentLoopOutput::CapabilityCalls(calls) => latency::stage!( "capabilities", - state.iteration, + host.run_context(), + turn_iteration, self.capabilities.process( ctx, CapabilityInput { @@ -240,14 +235,15 @@ impl DefaultExecutorPipeline { surface: prompt.surface, calls, }, - ) + ), )?, }; - let completed = trace_stage!( + let completed = latency::stage!( "post_capability", - completed.iteration(), - self.post_capability.process(ctx, completed) + host.run_context(), + completed.iteration_or(turn_iteration), + self.post_capability.process(ctx, completed), )?; let (next_state, summary) = match completed { @@ -256,8 +252,9 @@ impl DefaultExecutorPipeline { }; let completed_kind = summary.kind; - let (mut next_state, summary) = match trace_stage!( + let (mut next_state, summary) = match latency::stage!( "stop_observe", + host.run_context(), next_state.iteration, self.stop.observe( ctx, @@ -265,7 +262,7 @@ impl DefaultExecutorPipeline { state: next_state, summary, }, - ) + ), )? { StopObservationStep::Continue { state, summary } => (*state, summary), StopObservationStep::Exit(exit) => return Ok(exit), @@ -276,8 +273,9 @@ impl DefaultExecutorPipeline { iteration = next_state.iteration, "agent loop checking follow-up input after reply-only turn end" ); - match trace_stage!( + match latency::stage!( "input_drain_follow_up", + host.run_context(), next_state.iteration, self.input.process( ctx, @@ -286,7 +284,7 @@ impl DefaultExecutorPipeline { pending_input_ack: std::mem::take(&mut pending_input_ack), mode: UserFacingInputDrainMode::FollowUp, }, - ) + ), )? { InputStep::Continue { state: next, @@ -309,8 +307,9 @@ impl DefaultExecutorPipeline { } } - match trace_stage!( + match latency::stage!( "stop_decide", + host.run_context(), next_state.iteration, self.stop.decide( ctx, @@ -319,19 +318,26 @@ impl DefaultExecutorPipeline { summary, pending_input_ack: std::mem::take(&mut pending_input_ack), }, - ) + ), )? { StopStep::Stop { state, kind, pending_input_ack: mut ack, } => { - let exit = trace_stage!( + let exit_iteration = state.iteration; + let exit = latency::stage!( "exit", - state.iteration, - self.exit.process(ctx, ExitInput { state, kind }) + host.run_context(), + exit_iteration, + self.exit.process(ctx, ExitInput { state, kind }), + )?; + latency::stage!( + "ack_pending_input_before_exit", + host.run_context(), + exit_iteration, + ack.ack(host), )?; - trace_stage!("ack_pending_input_before_exit", 0, ack.ack(host))?; return Ok(exit); } StopStep::Continue { @@ -351,15 +357,18 @@ impl DefaultExecutorPipeline { | PromptStep::ResumeAuth(resume) | PromptStep::ResumeExternalTool(resume) => { let resume = *resume; + let resume_iteration = resume.state.iteration; pending_input_ack = resume.pending_input_ack; - trace_stage!( + latency::stage!( "ack_pending_input_before_resume_capability", - resume.state.iteration, - pending_input_ack.ack(host) + host.run_context(), + resume_iteration, + pending_input_ack.ack(host), )?; - let completed = trace_stage!( + let completed = latency::stage!( "capabilities_resume", - resume.state.iteration, + host.run_context(), + resume_iteration, self.capabilities.process( ctx, CapabilityInput { @@ -367,13 +376,14 @@ impl DefaultExecutorPipeline { surface: resume.surface, calls: vec![resume.call], }, - ) + ), )?; - let completed = trace_stage!( + let completed = latency::stage!( "post_capability_resume", - completed.iteration(), - self.post_capability.process(ctx, completed) + host.run_context(), + completed.iteration_or(resume_iteration), + self.post_capability.process(ctx, completed), )?; let (next_state, summary) = match completed { @@ -381,8 +391,9 @@ impl DefaultExecutorPipeline { TurnCompletedStep::Exit(exit) => return Ok(exit), }; - let (next_state, summary) = match trace_stage!( + let (next_state, summary) = match latency::stage!( "stop_observe_resume", + host.run_context(), next_state.iteration, self.stop.observe( ctx, @@ -390,14 +401,15 @@ impl DefaultExecutorPipeline { state: next_state, summary, }, - ) + ), )? { StopObservationStep::Continue { state, summary } => (*state, summary), StopObservationStep::Exit(exit) => return Ok(exit), }; - match trace_stage!( + match latency::stage!( "stop_decide_resume", + host.run_context(), next_state.iteration, self.stop.decide( ctx, @@ -406,19 +418,26 @@ impl DefaultExecutorPipeline { summary, pending_input_ack: std::mem::take(&mut pending_input_ack), }, - ) + ), )? { StopStep::Stop { state, kind, pending_input_ack: mut ack, } => { - let exit = trace_stage!( + let exit_iteration = state.iteration; + let exit = latency::stage!( "exit_resume", - state.iteration, - self.exit.process(ctx, ExitInput { state, kind }) + host.run_context(), + exit_iteration, + self.exit.process(ctx, ExitInput { state, kind }), + )?; + latency::stage!( + "ack_pending_input_before_exit_resume", + host.run_context(), + exit_iteration, + ack.ack(host), )?; - trace_stage!("ack_pending_input_before_exit_resume", 0, ack.ack(host))?; return Ok(exit); } StopStep::Continue { @@ -450,15 +469,17 @@ impl DefaultExecutorPipeline { pending_input_ack = ack; // Deliver the ack before stop.observe, mirroring the timing // of the Prepared path (line ~133: ack before ModelStage). - trace_stage!( + latency::stage!( "ack_pending_input_skip_model", + host.run_context(), skipped_state.iteration, - pending_input_ack.ack(host) + pending_input_ack.ack(host), )?; let summary = crate::strategies::TurnSummary::compaction_only(); - let (mut next_state, summary) = match trace_stage!( + let (mut next_state, summary) = match latency::stage!( "stop_observe_skip_model", + host.run_context(), skipped_state.iteration, self.stop.observe( ctx, @@ -466,7 +487,7 @@ impl DefaultExecutorPipeline { state: skipped_state, summary, }, - ) + ), )? { StopObservationStep::Continue { state, summary } => (*state, summary), StopObservationStep::Exit(exit) => return Ok(exit), @@ -475,8 +496,9 @@ impl DefaultExecutorPipeline { // Follow-up drain is skipped: CompactionOnly != ReplyOnly, // so the condition is naturally false here. - match trace_stage!( + match latency::stage!( "stop_decide_skip_model", + host.run_context(), next_state.iteration, self.stop.decide( ctx, @@ -485,22 +507,25 @@ impl DefaultExecutorPipeline { summary, pending_input_ack: std::mem::take(&mut pending_input_ack), }, - ) + ), )? { StopStep::Stop { state, kind, pending_input_ack: mut ack, } => { - let exit = trace_stage!( + let exit_iteration = state.iteration; + let exit = latency::stage!( "exit_skip_model", - state.iteration, - self.exit.process(ctx, ExitInput { state, kind }) + host.run_context(), + exit_iteration, + self.exit.process(ctx, ExitInput { state, kind }), )?; - trace_stage!( + latency::stage!( "ack_pending_input_before_exit_skip_model", - 0, - ack.ack(host) + host.run_context(), + exit_iteration, + ack.ack(host), )?; return Ok(exit); } diff --git a/crates/ironclaw_agent_loop/src/executor/latency.rs b/crates/ironclaw_agent_loop/src/executor/latency.rs index 93e70ef12b4..5f6084d4ddc 100644 --- a/crates/ironclaw_agent_loop/src/executor/latency.rs +++ b/crates/ironclaw_agent_loop/src/executor/latency.rs @@ -73,3 +73,21 @@ pub(super) fn result( Err(error) => operation_error(operation, context, iteration, started_at, error), } } + +macro_rules! stage { + ($operation:expr, $context:expr, $iteration:expr, $future:expr $(,)?) => {{ + let iteration = $iteration; + let stage_started_at = $crate::executor::latency::started_at(); + let result = $future.await; + $crate::executor::latency::result( + $operation, + $context, + iteration, + stage_started_at, + &result, + ); + result + }}; +} + +pub(super) use stage;