TUIに向けたprotocolの詳細調整

This commit is contained in:
2026-04-21 20:50:59 +09:00
parent de3272fdfd
commit ce59c5320e
9 changed files with 337 additions and 15 deletions
+29
View File
@@ -162,6 +162,11 @@ pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
/// can be forwarded to the user — distinct from `tracing::warn!`,
/// which is for developer-facing logs.
warning_cbs: Vec<Box<dyn Fn(&str) + Send + Sync>>,
/// Tool-result callbacks. Invoked once per completed tool call
/// after post-execution interceptors and the output byte-cap
/// truncation have been applied — i.e. on the same data that
/// enters history.
tool_result_cbs: Vec<Box<dyn Fn(&ToolResult) + Send + Sync>>,
/// Request configuration (max_tokens, temperature, etc.)
request_config: RequestConfig,
/// Whether the previous run was interrupted
@@ -302,6 +307,22 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
}
/// Register a callback invoked once per completed tool execution.
///
/// Fired after `post_tool_call` interceptors and any `content`
/// truncation from `tool_output_limits`, so the callback observes
/// exactly what is persisted to history. Intended for upper layers
/// (e.g. Pod) to forward tool results to clients.
pub fn on_tool_result(&mut self, callback: impl Fn(&ToolResult) + Send + Sync + 'static) {
self.tool_result_cbs.push(Box::new(callback));
}
fn emit_tool_result(&self, result: &ToolResult) {
for cb in &self.tool_result_cbs {
cb(result);
}
}
/// Register a turn-end callback (receives 0-based turn number).
pub fn on_turn_end(&mut self, callback: impl Fn(usize) + Send + Sync + 'static) {
self.turn_end_cbs.push(Box::new(callback));
@@ -753,6 +774,11 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
}
// Emit per-result callbacks on the post-truncation payload.
for tool_result in &results {
self.emit_tool_result(tool_result);
}
Ok(ToolExecutionResult::Completed(results))
}
@@ -1016,6 +1042,7 @@ impl<C: LlmClient> Worker<C, Mutable> {
turn_start_cbs: Vec::new(),
turn_end_cbs: Vec::new(),
warning_cbs: Vec::new(),
tool_result_cbs: Vec::new(),
request_config: RequestConfig::default(),
last_run_interrupted: false,
cancel_tx,
@@ -1270,6 +1297,7 @@ impl<C: LlmClient> Worker<C, Mutable> {
turn_start_cbs: self.turn_start_cbs,
turn_end_cbs: self.turn_end_cbs,
warning_cbs: self.warning_cbs,
tool_result_cbs: self.tool_result_cbs,
request_config: self.request_config,
last_run_interrupted: self.last_run_interrupted,
@@ -1344,6 +1372,7 @@ impl<C: LlmClient> Worker<C, Locked> {
turn_start_cbs: self.turn_start_cbs,
turn_end_cbs: self.turn_end_cbs,
warning_cbs: self.warning_cbs,
tool_result_cbs: self.tool_result_cbs,
request_config: self.request_config,
last_run_interrupted: self.last_run_interrupted,
+141
View File
@@ -6,9 +6,11 @@ mod common;
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use common::MockLlmClient;
use llm_worker::Worker;
use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent as ClientStatusEvent};
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
// =============================================================================
// Tests
@@ -149,6 +151,145 @@ async fn test_callback_turn_events() {
assert_eq!(ends[0], 0);
}
/// Stub tool returning a fixed [`ToolOutput`] for result-callback tests.
struct FixedOutputTool {
output: ToolOutput,
}
#[async_trait]
impl Tool for FixedOutputTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
Ok(self.output.clone())
}
}
fn fixed_tool(name: &'static str, output: ToolOutput) -> ToolDefinition {
Arc::new(move || {
let meta = ToolMeta::new(name).input_schema(serde_json::json!({"type":"object"}));
(
meta,
Arc::new(FixedOutputTool {
output: output.clone(),
}) as Arc<dyn Tool>,
)
})
}
/// Verify that on_tool_result fires once per executed tool with
/// summary/content/is_error matching what the tool returned.
#[tokio::test]
async fn test_callback_tool_result_events() {
let events = vec![
Event::tool_use_start(0, "call_1", "fixed"),
Event::tool_input_delta(0, "{}"),
Event::tool_use_stop(0),
Event::Status(ClientStatusEvent {
status: ResponseStatus::Completed,
}),
];
let client = MockLlmClient::new(events);
let mut worker = Worker::new(client);
worker.register_tool(fixed_tool(
"fixed",
ToolOutput {
summary: "did the thing".into(),
content: Some("full detail body".into()),
},
));
let captured: Arc<Mutex<Vec<(String, String, Option<String>, bool)>>> =
Arc::new(Mutex::new(Vec::new()));
let sink = captured.clone();
worker.on_tool_result(move |result| {
sink.lock().unwrap().push((
result.tool_use_id.clone(),
result.summary.clone(),
result.content.clone(),
result.is_error,
));
});
let _ = worker.run("call it").await;
let observed = captured.lock().unwrap();
assert_eq!(observed.len(), 1);
assert_eq!(observed[0].0, "call_1");
assert_eq!(observed[0].1, "did the thing");
assert_eq!(observed[0].2.as_deref(), Some("full detail body"));
assert!(!observed[0].3);
}
/// Stub tool that always fails, for exercising the error path through
/// `on_tool_result`.
struct ErroringTool {
message: String,
}
#[async_trait]
impl Tool for ErroringTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
Err(ToolError::ExecutionFailed(self.message.clone()))
}
}
fn erroring_tool(name: &'static str, message: &'static str) -> ToolDefinition {
Arc::new(move || {
let meta = ToolMeta::new(name).input_schema(serde_json::json!({"type":"object"}));
(
meta,
Arc::new(ErroringTool {
message: message.to_string(),
}) as Arc<dyn Tool>,
)
})
}
/// Verify on_tool_result also fires for failed executions with
/// is_error=true, and that the ToolOutput content channel stays empty.
#[tokio::test]
async fn test_callback_tool_result_error_path() {
let events = vec![
Event::tool_use_start(0, "call_err", "erroring"),
Event::tool_input_delta(0, "{}"),
Event::tool_use_stop(0),
Event::Status(ClientStatusEvent {
status: ResponseStatus::Completed,
}),
];
let client = MockLlmClient::new(events);
let mut worker = Worker::new(client);
worker.register_tool(erroring_tool("erroring", "boom"));
let captured: Arc<Mutex<Vec<(String, String, Option<String>, bool)>>> =
Arc::new(Mutex::new(Vec::new()));
let sink = captured.clone();
worker.on_tool_result(move |result| {
sink.lock().unwrap().push((
result.tool_use_id.clone(),
result.summary.clone(),
result.content.clone(),
result.is_error,
));
});
let _ = worker.run("fail it").await;
let observed = captured.lock().unwrap();
assert_eq!(observed.len(), 1);
assert_eq!(observed[0].0, "call_err");
assert!(
observed[0].1.contains("boom"),
"summary should carry the error message: {}",
observed[0].1
);
assert!(observed[0].2.is_none());
assert!(observed[0].3);
}
/// Verify that on_usage callback receives usage events
#[tokio::test]
async fn test_callback_usage_events() {