wip: telemetry progress

This commit is contained in:
Yeachan-Heo
2026-04-01 04:30:29 +00:00
parent ac6c5d00a8
commit 5170718306
12 changed files with 1060 additions and 70 deletions

View File

@@ -11,6 +11,7 @@ glob = "0.3"
regex = "1"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
telemetry = { path = "../telemetry" }
tokio = { version = "1", features = ["io-util", "macros", "process", "rt", "rt-multi-thread", "time"] }
walkdir = "2"

View File

@@ -1,6 +1,9 @@
use std::collections::BTreeMap;
use std::fmt::{Display, Formatter};
use serde_json::{Map, Value};
use telemetry::SessionTracer;
use crate::compact::{
compact_session, estimate_session_tokens, CompactionConfig, CompactionResult,
};
@@ -97,6 +100,7 @@ pub struct ConversationRuntime<C, T> {
max_iterations: usize,
usage_tracker: UsageTracker,
hook_runner: HookRunner,
session_tracer: Option<SessionTracer>,
}
impl<C, T> ConversationRuntime<C, T>
@@ -118,7 +122,7 @@ where
tool_executor,
permission_policy,
system_prompt,
RuntimeFeatureConfig::default(),
&RuntimeFeatureConfig::default(),
)
}
@@ -129,7 +133,7 @@ where
tool_executor: T,
permission_policy: PermissionPolicy,
system_prompt: Vec<String>,
feature_config: RuntimeFeatureConfig,
feature_config: &RuntimeFeatureConfig,
) -> Self {
let usage_tracker = UsageTracker::from_session(&session);
Self {
@@ -140,7 +144,8 @@ where
system_prompt,
max_iterations: usize::MAX,
usage_tracker,
hook_runner: HookRunner::from_feature_config(&feature_config),
hook_runner: HookRunner::from_feature_config(feature_config),
session_tracer: None,
}
}
@@ -150,14 +155,22 @@ where
self
}
#[must_use]
pub fn with_session_tracer(mut self, session_tracer: SessionTracer) -> Self {
self.session_tracer = Some(session_tracer);
self
}
pub fn run_turn(
&mut self,
user_input: impl Into<String>,
mut prompter: Option<&mut dyn PermissionPrompter>,
) -> Result<TurnSummary, RuntimeError> {
let user_input = user_input.into();
self.record_turn_started(&user_input);
self.session
.messages
.push(ConversationMessage::user_text(user_input.into()));
.push(ConversationMessage::user_text(user_input));
let mut assistant_messages = Vec::new();
let mut tool_results = Vec::new();
@@ -166,16 +179,24 @@ where
loop {
iterations += 1;
if iterations > self.max_iterations {
return Err(RuntimeError::new(
let error = RuntimeError::new(
"conversation loop exceeded the maximum number of iterations",
));
);
self.record_turn_failed(iterations, &error);
return Err(error);
}
let request = ApiRequest {
system_prompt: self.system_prompt.clone(),
messages: self.session.messages.clone(),
};
let events = self.api_client.stream(request)?;
let events = match self.api_client.stream(request) {
Ok(events) => events,
Err(error) => {
self.record_turn_failed(iterations, &error);
return Err(error);
}
};
let (assistant_message, usage) = build_assistant_message(events)?;
if let Some(usage) = usage {
self.usage_tracker.record(usage);
@@ -190,6 +211,7 @@ where
_ => None,
})
.collect::<Vec<_>>();
self.record_assistant_iteration(iterations, &assistant_message, pending_tool_uses.len());
self.session.messages.push(assistant_message.clone());
assistant_messages.push(assistant_message);
@@ -199,6 +221,7 @@ where
}
for (tool_use_id, tool_name, input) in pending_tool_uses {
self.record_tool_started(iterations, &tool_name);
let permission_outcome = if let Some(prompt) = prompter.as_mut() {
self.permission_policy
.authorize(&tool_name, &input, Some(*prompt))
@@ -249,17 +272,20 @@ where
ConversationMessage::tool_result(tool_use_id, tool_name, reason, true)
}
};
self.record_tool_finished(iterations, &result_message);
self.session.messages.push(result_message.clone());
tool_results.push(result_message);
}
}
Ok(TurnSummary {
let summary = TurnSummary {
assistant_messages,
tool_results,
iterations,
usage: self.usage_tracker.cumulative_usage(),
})
};
self.record_turn_completed(&summary);
Ok(summary)
}
#[must_use]
@@ -286,6 +312,125 @@ where
pub fn into_session(self) -> Session {
self.session
}
fn record_turn_started(&self, user_input: &str) {
if let Some(tracer) = &self.session_tracer {
let mut attributes = Map::new();
attributes.insert(
"message_count_before".to_string(),
Value::from(u64::try_from(self.session.messages.len()).unwrap_or(u64::MAX)),
);
attributes.insert(
"input_chars".to_string(),
Value::from(u64::try_from(user_input.chars().count()).unwrap_or(u64::MAX)),
);
tracer.record("turn_started", attributes);
}
}
fn record_assistant_iteration(
&self,
iteration: usize,
assistant_message: &ConversationMessage,
pending_tool_count: usize,
) {
if let Some(tracer) = &self.session_tracer {
let mut attributes = Map::new();
attributes.insert(
"iteration".to_string(),
Value::from(u64::try_from(iteration).unwrap_or(u64::MAX)),
);
attributes.insert(
"block_count".to_string(),
Value::from(u64::try_from(assistant_message.blocks.len()).unwrap_or(u64::MAX)),
);
attributes.insert(
"pending_tool_count".to_string(),
Value::from(u64::try_from(pending_tool_count).unwrap_or(u64::MAX)),
);
tracer.record("assistant_iteration_completed", attributes);
}
}
fn record_tool_started(&self, iteration: usize, tool_name: &str) {
if let Some(tracer) = &self.session_tracer {
let mut attributes = Map::new();
attributes.insert(
"iteration".to_string(),
Value::from(u64::try_from(iteration).unwrap_or(u64::MAX)),
);
attributes.insert("tool_name".to_string(), Value::String(tool_name.to_string()));
tracer.record("tool_execution_started", attributes);
}
}
fn record_tool_finished(&self, iteration: usize, result_message: &ConversationMessage) {
let Some(tracer) = &self.session_tracer else {
return;
};
let Some(ContentBlock::ToolResult {
tool_name,
is_error,
output,
..
}) = result_message.blocks.first()
else {
return;
};
let mut attributes = Map::new();
attributes.insert(
"iteration".to_string(),
Value::from(u64::try_from(iteration).unwrap_or(u64::MAX)),
);
attributes.insert("tool_name".to_string(), Value::String(tool_name.clone()));
attributes.insert("is_error".to_string(), Value::Bool(*is_error));
attributes.insert(
"output_chars".to_string(),
Value::from(u64::try_from(output.chars().count()).unwrap_or(u64::MAX)),
);
tracer.record("tool_execution_finished", attributes);
}
fn record_turn_completed(&self, summary: &TurnSummary) {
if let Some(tracer) = &self.session_tracer {
let mut attributes = Map::new();
attributes.insert(
"assistant_message_count".to_string(),
Value::from(
u64::try_from(summary.assistant_messages.len()).unwrap_or(u64::MAX),
),
);
attributes.insert(
"tool_result_count".to_string(),
Value::from(u64::try_from(summary.tool_results.len()).unwrap_or(u64::MAX)),
);
attributes.insert(
"iterations".to_string(),
Value::from(u64::try_from(summary.iterations).unwrap_or(u64::MAX)),
);
attributes.insert(
"total_input_tokens".to_string(),
Value::from(summary.usage.input_tokens),
);
attributes.insert(
"total_output_tokens".to_string(),
Value::from(summary.usage.output_tokens),
);
tracer.record("turn_completed", attributes);
}
}
fn record_turn_failed(&self, iteration: usize, error: &RuntimeError) {
if let Some(tracer) = &self.session_tracer {
let mut attributes = Map::new();
attributes.insert(
"iteration".to_string(),
Value::from(u64::try_from(iteration).unwrap_or(u64::MAX)),
);
attributes.insert("error".to_string(), Value::String(error.to_string()));
tracer.record("turn_failed", attributes);
}
}
}
fn build_assistant_message(
@@ -609,7 +754,7 @@ mod tests {
}),
PermissionPolicy::new(PermissionMode::DangerFullAccess),
vec!["system".to_string()],
RuntimeFeatureConfig::default().with_hooks(RuntimeHookConfig::new(
&RuntimeFeatureConfig::default().with_hooks(RuntimeHookConfig::new(
vec![shell_snippet("printf 'blocked by hook'; exit 2")],
Vec::new(),
)),
@@ -675,7 +820,7 @@ mod tests {
StaticToolExecutor::new().register("add", |_input| Ok("4".to_string())),
PermissionPolicy::new(PermissionMode::DangerFullAccess),
vec!["system".to_string()],
RuntimeFeatureConfig::default().with_hooks(RuntimeHookConfig::new(
&RuntimeFeatureConfig::default().with_hooks(RuntimeHookConfig::new(
vec![shell_snippet("printf 'pre hook ran'")],
vec![shell_snippet("printf 'post hook ran'")],
)),
@@ -697,7 +842,7 @@ mod tests {
"post hook should preserve non-error result: {output:?}"
);
assert!(
output.contains("4"),
output.contains('4'),
"tool output missing value: {output:?}"
);
assert!(

View File

@@ -64,7 +64,7 @@ impl HookRunner {
#[must_use]
pub fn run_pre_tool_use(&self, tool_name: &str, tool_input: &str) -> HookRunResult {
self.run_commands(
Self::run_commands(
HookEvent::PreToolUse,
self.config.pre_tool_use(),
tool_name,
@@ -82,7 +82,7 @@ impl HookRunner {
tool_output: &str,
is_error: bool,
) -> HookRunResult {
self.run_commands(
Self::run_commands(
HookEvent::PostToolUse,
self.config.post_tool_use(),
tool_name,
@@ -93,7 +93,6 @@ impl HookRunner {
}
fn run_commands(
&self,
event: HookEvent,
commands: &[String],
tool_name: &str,
@@ -114,19 +113,19 @@ impl HookRunner {
"tool_result_is_error": is_error,
})
.to_string();
let invocation = HookInvocation {
event,
tool_name,
tool_input,
tool_output,
is_error,
payload: &payload,
};
let mut messages = Vec::new();
for command in commands {
match self.run_command(
command,
event,
tool_name,
tool_input,
tool_output,
is_error,
&payload,
) {
match Self::run_command(command, &invocation) {
HookCommandOutcome::Allow { message } => {
if let Some(message) = message {
messages.push(message);
@@ -149,29 +148,23 @@ impl HookRunner {
HookRunResult::allow(messages)
}
fn run_command(
&self,
command: &str,
event: HookEvent,
tool_name: &str,
tool_input: &str,
tool_output: Option<&str>,
is_error: bool,
payload: &str,
) -> HookCommandOutcome {
fn run_command(command: &str, invocation: &HookInvocation<'_>) -> HookCommandOutcome {
let mut child = shell_command(command);
child.stdin(std::process::Stdio::piped());
child.stdout(std::process::Stdio::piped());
child.stderr(std::process::Stdio::piped());
child.env("HOOK_EVENT", event.as_str());
child.env("HOOK_TOOL_NAME", tool_name);
child.env("HOOK_TOOL_INPUT", tool_input);
child.env("HOOK_TOOL_IS_ERROR", if is_error { "1" } else { "0" });
if let Some(tool_output) = tool_output {
child.env("HOOK_EVENT", invocation.event.as_str());
child.env("HOOK_TOOL_NAME", invocation.tool_name);
child.env("HOOK_TOOL_INPUT", invocation.tool_input);
child.env(
"HOOK_TOOL_IS_ERROR",
if invocation.is_error { "1" } else { "0" },
);
if let Some(tool_output) = invocation.tool_output {
child.env("HOOK_TOOL_OUTPUT", tool_output);
}
match child.output_with_stdin(payload.as_bytes()) {
match child.output_with_stdin(invocation.payload.as_bytes()) {
Ok(output) => {
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
@@ -189,8 +182,9 @@ impl HookRunner {
},
None => HookCommandOutcome::Warn {
message: format!(
"{} hook `{command}` terminated by signal while handling `{tool_name}`",
event.as_str()
"{} hook `{command}` terminated by signal while handling `{}`",
invocation.event.as_str(),
invocation.tool_name
),
},
}
@@ -198,13 +192,23 @@ impl HookRunner {
Err(error) => HookCommandOutcome::Warn {
message: format!(
"{} hook `{command}` failed to start for `{tool_name}`: {error}",
event.as_str()
invocation.event.as_str(),
tool_name = invocation.tool_name
),
},
}
}
}
struct HookInvocation<'a> {
event: HookEvent,
tool_name: &'a str,
tool_input: &'a str,
tool_output: Option<&'a str>,
is_error: bool,
payload: &'a str,
}
enum HookCommandOutcome {
Allow { message: Option<String> },
Deny { message: Option<String> },