mirror of
https://github.com/instructkr/claude-code.git
synced 2026-04-03 23:08:49 +03:00
fix: update prompt tests for post-plugins-merge format
This commit is contained in:
@@ -1,20 +1,15 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
use plugins::{HookRunner as PluginHookRunner, PluginRegistry};
|
||||
|
||||
use crate::compact::{
|
||||
compact_session, estimate_session_tokens, CompactionConfig, CompactionResult,
|
||||
};
|
||||
use crate::config::RuntimeFeatureConfig;
|
||||
use crate::hooks::HookRunner;
|
||||
use crate::hooks::{HookRunResult, HookRunner};
|
||||
use crate::permissions::{PermissionOutcome, PermissionPolicy, PermissionPrompter};
|
||||
use crate::session::{ContentBlock, ConversationMessage, Session};
|
||||
use crate::usage::{TokenUsage, UsageTracker};
|
||||
|
||||
const DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD: u32 = 200_000;
|
||||
const AUTO_COMPACTION_THRESHOLD_ENV_VAR: &str = "CLAUDE_CODE_AUTO_COMPACT_INPUT_TOKENS";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ApiRequest {
|
||||
pub system_prompt: Vec<String>,
|
||||
@@ -91,12 +86,6 @@ pub struct TurnSummary {
|
||||
pub tool_results: Vec<ConversationMessage>,
|
||||
pub iterations: usize,
|
||||
pub usage: TokenUsage,
|
||||
pub auto_compaction: Option<AutoCompactionEvent>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AutoCompactionEvent {
|
||||
pub removed_message_count: usize,
|
||||
}
|
||||
|
||||
pub struct ConversationRuntime<C, T> {
|
||||
@@ -108,25 +97,6 @@ pub struct ConversationRuntime<C, T> {
|
||||
max_iterations: usize,
|
||||
usage_tracker: UsageTracker,
|
||||
hook_runner: HookRunner,
|
||||
auto_compaction_input_tokens_threshold: u32,
|
||||
plugin_hook_runner: Option<PluginHookRunner>,
|
||||
plugin_registry: Option<PluginRegistry>,
|
||||
plugins_shutdown: bool,
|
||||
}
|
||||
|
||||
impl<C, T> ConversationRuntime<C, T> {
|
||||
fn shutdown_registered_plugins(&mut self) -> Result<(), RuntimeError> {
|
||||
if self.plugins_shutdown {
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(registry) = &self.plugin_registry {
|
||||
registry
|
||||
.shutdown()
|
||||
.map_err(|error| RuntimeError::new(format!("plugin shutdown failed: {error}")))?;
|
||||
}
|
||||
self.plugins_shutdown = true;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<C, T> ConversationRuntime<C, T>
|
||||
@@ -148,19 +118,18 @@ where
|
||||
tool_executor,
|
||||
permission_policy,
|
||||
system_prompt,
|
||||
RuntimeFeatureConfig::default(),
|
||||
&RuntimeFeatureConfig::default(),
|
||||
)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
#[allow(clippy::needless_pass_by_value)]
|
||||
pub fn new_with_features(
|
||||
session: Session,
|
||||
api_client: C,
|
||||
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 {
|
||||
@@ -171,57 +140,16 @@ where
|
||||
system_prompt,
|
||||
max_iterations: usize::MAX,
|
||||
usage_tracker,
|
||||
hook_runner: HookRunner::from_feature_config(&feature_config),
|
||||
auto_compaction_input_tokens_threshold: auto_compaction_threshold_from_env(),
|
||||
plugin_hook_runner: None,
|
||||
plugin_registry: None,
|
||||
plugins_shutdown: false,
|
||||
hook_runner: HookRunner::from_feature_config(feature_config),
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::needless_pass_by_value)]
|
||||
pub fn new_with_plugins(
|
||||
session: Session,
|
||||
api_client: C,
|
||||
tool_executor: T,
|
||||
permission_policy: PermissionPolicy,
|
||||
system_prompt: Vec<String>,
|
||||
feature_config: RuntimeFeatureConfig,
|
||||
plugin_registry: PluginRegistry,
|
||||
) -> Result<Self, RuntimeError> {
|
||||
let plugin_hook_runner =
|
||||
PluginHookRunner::from_registry(&plugin_registry).map_err(|error| {
|
||||
RuntimeError::new(format!("plugin hook registration failed: {error}"))
|
||||
})?;
|
||||
plugin_registry
|
||||
.initialize()
|
||||
.map_err(|error| RuntimeError::new(format!("plugin initialization failed: {error}")))?;
|
||||
let mut runtime = Self::new_with_features(
|
||||
session,
|
||||
api_client,
|
||||
tool_executor,
|
||||
permission_policy,
|
||||
system_prompt,
|
||||
feature_config,
|
||||
);
|
||||
runtime.plugin_hook_runner = Some(plugin_hook_runner);
|
||||
runtime.plugin_registry = Some(plugin_registry);
|
||||
Ok(runtime)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_max_iterations(mut self, max_iterations: usize) -> Self {
|
||||
self.max_iterations = max_iterations;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_auto_compaction_input_tokens_threshold(mut self, threshold: u32) -> Self {
|
||||
self.auto_compaction_input_tokens_threshold = threshold;
|
||||
self
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_lines)]
|
||||
pub fn run_turn(
|
||||
&mut self,
|
||||
user_input: impl Into<String>,
|
||||
@@ -234,7 +162,6 @@ where
|
||||
let mut assistant_messages = Vec::new();
|
||||
let mut tool_results = Vec::new();
|
||||
let mut iterations = 0;
|
||||
let mut max_turn_input_tokens = 0;
|
||||
|
||||
loop {
|
||||
iterations += 1;
|
||||
@@ -251,7 +178,6 @@ where
|
||||
let events = self.api_client.stream(request)?;
|
||||
let (assistant_message, usage) = build_assistant_message(events)?;
|
||||
if let Some(usage) = usage {
|
||||
max_turn_input_tokens = max_turn_input_tokens.max(usage.input_tokens);
|
||||
self.usage_tracker.record(usage);
|
||||
}
|
||||
let pending_tool_uses = assistant_message
|
||||
@@ -288,74 +214,35 @@ where
|
||||
ConversationMessage::tool_result(
|
||||
tool_use_id,
|
||||
tool_name,
|
||||
format_hook_message(pre_hook_result.messages(), &deny_message),
|
||||
format_hook_message(&pre_hook_result, &deny_message),
|
||||
true,
|
||||
)
|
||||
} else {
|
||||
let plugin_pre_hook_result =
|
||||
self.run_plugin_pre_tool_use(&tool_name, &input);
|
||||
if plugin_pre_hook_result.is_denied() {
|
||||
let deny_message =
|
||||
format!("PreToolUse hook denied tool `{tool_name}`");
|
||||
let mut messages = pre_hook_result.messages().to_vec();
|
||||
messages.extend(plugin_pre_hook_result.messages().iter().cloned());
|
||||
ConversationMessage::tool_result(
|
||||
tool_use_id,
|
||||
tool_name,
|
||||
format_hook_message(&messages, &deny_message),
|
||||
true,
|
||||
)
|
||||
} else {
|
||||
let (mut output, mut is_error) =
|
||||
match self.tool_executor.execute(&tool_name, &input) {
|
||||
Ok(output) => (output, false),
|
||||
Err(error) => (error.to_string(), true),
|
||||
};
|
||||
output =
|
||||
merge_hook_feedback(pre_hook_result.messages(), output, false);
|
||||
output = merge_hook_feedback(
|
||||
plugin_pre_hook_result.messages(),
|
||||
output,
|
||||
false,
|
||||
);
|
||||
let (mut output, mut is_error) =
|
||||
match self.tool_executor.execute(&tool_name, &input) {
|
||||
Ok(output) => (output, false),
|
||||
Err(error) => (error.to_string(), true),
|
||||
};
|
||||
output = merge_hook_feedback(pre_hook_result.messages(), output, false);
|
||||
|
||||
let hook_output = output.clone();
|
||||
let post_hook_result = self.hook_runner.run_post_tool_use(
|
||||
&tool_name,
|
||||
&input,
|
||||
&hook_output,
|
||||
is_error,
|
||||
);
|
||||
let plugin_post_hook_result = self.run_plugin_post_tool_use(
|
||||
&tool_name,
|
||||
&input,
|
||||
&hook_output,
|
||||
is_error,
|
||||
);
|
||||
if post_hook_result.is_denied() {
|
||||
is_error = true;
|
||||
}
|
||||
if plugin_post_hook_result.is_denied() {
|
||||
is_error = true;
|
||||
}
|
||||
output = merge_hook_feedback(
|
||||
post_hook_result.messages(),
|
||||
output,
|
||||
post_hook_result.is_denied(),
|
||||
);
|
||||
output = merge_hook_feedback(
|
||||
plugin_post_hook_result.messages(),
|
||||
output,
|
||||
plugin_post_hook_result.is_denied(),
|
||||
);
|
||||
|
||||
ConversationMessage::tool_result(
|
||||
tool_use_id,
|
||||
tool_name,
|
||||
output,
|
||||
is_error,
|
||||
)
|
||||
let post_hook_result = self
|
||||
.hook_runner
|
||||
.run_post_tool_use(&tool_name, &input, &output, is_error);
|
||||
if post_hook_result.is_denied() {
|
||||
is_error = true;
|
||||
}
|
||||
output = merge_hook_feedback(
|
||||
post_hook_result.messages(),
|
||||
output,
|
||||
post_hook_result.is_denied(),
|
||||
);
|
||||
|
||||
ConversationMessage::tool_result(
|
||||
tool_use_id,
|
||||
tool_name,
|
||||
output,
|
||||
is_error,
|
||||
)
|
||||
}
|
||||
}
|
||||
PermissionOutcome::Deny { reason } => {
|
||||
@@ -367,14 +254,11 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
let auto_compaction = self.maybe_auto_compact(max_turn_input_tokens);
|
||||
|
||||
Ok(TurnSummary {
|
||||
assistant_messages,
|
||||
tool_results,
|
||||
iterations,
|
||||
usage: self.usage_tracker.cumulative_usage(),
|
||||
auto_compaction,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -399,81 +283,9 @@ where
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn into_session(mut self) -> Session {
|
||||
let _ = self.shutdown_registered_plugins();
|
||||
std::mem::take(&mut self.session)
|
||||
pub fn into_session(self) -> Session {
|
||||
self.session
|
||||
}
|
||||
|
||||
pub fn shutdown_plugins(&mut self) -> Result<(), RuntimeError> {
|
||||
self.shutdown_registered_plugins()
|
||||
}
|
||||
|
||||
fn run_plugin_pre_tool_use(&self, tool_name: &str, input: &str) -> plugins::HookRunResult {
|
||||
self.plugin_hook_runner.as_ref().map_or_else(
|
||||
|| plugins::HookRunResult::allow(Vec::new()),
|
||||
|runner| runner.run_pre_tool_use(tool_name, input),
|
||||
)
|
||||
}
|
||||
|
||||
fn run_plugin_post_tool_use(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
input: &str,
|
||||
output: &str,
|
||||
is_error: bool,
|
||||
) -> plugins::HookRunResult {
|
||||
self.plugin_hook_runner.as_ref().map_or_else(
|
||||
|| plugins::HookRunResult::allow(Vec::new()),
|
||||
|runner| runner.run_post_tool_use(tool_name, input, output, is_error),
|
||||
)
|
||||
}
|
||||
|
||||
fn maybe_auto_compact(&mut self, turn_input_tokens: u32) -> Option<AutoCompactionEvent> {
|
||||
if turn_input_tokens < self.auto_compaction_input_tokens_threshold {
|
||||
return None;
|
||||
}
|
||||
|
||||
let result = compact_session(
|
||||
&self.session,
|
||||
CompactionConfig {
|
||||
max_estimated_tokens: usize::try_from(self.auto_compaction_input_tokens_threshold)
|
||||
.unwrap_or(usize::MAX),
|
||||
..CompactionConfig::default()
|
||||
},
|
||||
);
|
||||
|
||||
if result.removed_message_count == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
self.session = result.compacted_session;
|
||||
Some(AutoCompactionEvent {
|
||||
removed_message_count: result.removed_message_count,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<C, T> Drop for ConversationRuntime<C, T> {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.shutdown_registered_plugins();
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn auto_compaction_threshold_from_env() -> u32 {
|
||||
parse_auto_compaction_threshold(
|
||||
std::env::var(AUTO_COMPACTION_THRESHOLD_ENV_VAR)
|
||||
.ok()
|
||||
.as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
fn parse_auto_compaction_threshold(value: Option<&str>) -> u32 {
|
||||
value
|
||||
.and_then(|raw| raw.trim().parse::<u32>().ok())
|
||||
.filter(|threshold| *threshold > 0)
|
||||
.unwrap_or(DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD)
|
||||
}
|
||||
|
||||
fn build_assistant_message(
|
||||
@@ -523,11 +335,11 @@ fn flush_text_block(text: &mut String, blocks: &mut Vec<ContentBlock>) {
|
||||
}
|
||||
}
|
||||
|
||||
fn format_hook_message(messages: &[String], fallback: &str) -> String {
|
||||
if messages.is_empty() {
|
||||
fn format_hook_message(result: &HookRunResult, fallback: &str) -> String {
|
||||
if result.messages().is_empty() {
|
||||
fallback.to_string()
|
||||
} else {
|
||||
messages.join("\n")
|
||||
result.messages().join("\n")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -584,9 +396,8 @@ impl ToolExecutor for StaticToolExecutor {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
parse_auto_compaction_threshold, ApiClient, ApiRequest, AssistantEvent,
|
||||
AutoCompactionEvent, ConversationRuntime, RuntimeError, StaticToolExecutor,
|
||||
DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD,
|
||||
ApiClient, ApiRequest, AssistantEvent, ConversationRuntime, RuntimeError,
|
||||
StaticToolExecutor,
|
||||
};
|
||||
use crate::compact::CompactionConfig;
|
||||
use crate::config::{RuntimeFeatureConfig, RuntimeHookConfig};
|
||||
@@ -597,13 +408,7 @@ mod tests {
|
||||
use crate::prompt::{ProjectContext, SystemPromptBuilder};
|
||||
use crate::session::{ContentBlock, MessageRole, Session};
|
||||
use crate::usage::TokenUsage;
|
||||
use plugins::{PluginManager, PluginManagerConfig};
|
||||
use std::fs;
|
||||
#[cfg(unix)]
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
struct ScriptedApiClient {
|
||||
call_count: usize,
|
||||
@@ -665,68 +470,6 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
fn temp_dir(label: &str) -> PathBuf {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_nanos();
|
||||
std::env::temp_dir().join(format!("runtime-plugin-{label}-{nanos}"))
|
||||
}
|
||||
|
||||
fn write_lifecycle_plugin(root: &Path, name: &str) -> PathBuf {
|
||||
fs::create_dir_all(root.join(".claude-plugin")).expect("manifest dir");
|
||||
fs::create_dir_all(root.join("lifecycle")).expect("lifecycle dir");
|
||||
let log_path = root.join("lifecycle.log");
|
||||
fs::write(
|
||||
root.join("lifecycle").join("init.sh"),
|
||||
"#!/bin/sh\nprintf 'init\\n' >> lifecycle.log\n",
|
||||
)
|
||||
.expect("write init script");
|
||||
fs::write(
|
||||
root.join("lifecycle").join("shutdown.sh"),
|
||||
"#!/bin/sh\nprintf 'shutdown\\n' >> lifecycle.log\n",
|
||||
)
|
||||
.expect("write shutdown script");
|
||||
fs::write(
|
||||
root.join(".claude-plugin").join("plugin.json"),
|
||||
format!(
|
||||
"{{\n \"name\": \"{name}\",\n \"version\": \"1.0.0\",\n \"description\": \"runtime lifecycle plugin\",\n \"lifecycle\": {{\n \"Init\": [\"./lifecycle/init.sh\"],\n \"Shutdown\": [\"./lifecycle/shutdown.sh\"]\n }}\n}}"
|
||||
),
|
||||
)
|
||||
.expect("write plugin manifest");
|
||||
log_path
|
||||
}
|
||||
|
||||
fn write_hook_plugin(root: &Path, name: &str, pre_message: &str, post_message: &str) {
|
||||
fs::create_dir_all(root.join(".claude-plugin")).expect("manifest dir");
|
||||
fs::create_dir_all(root.join("hooks")).expect("hooks dir");
|
||||
fs::write(
|
||||
root.join("hooks").join("pre.sh"),
|
||||
format!("#!/bin/sh\nprintf '%s\\n' '{pre_message}'\n"),
|
||||
)
|
||||
.expect("write pre hook");
|
||||
fs::write(
|
||||
root.join("hooks").join("post.sh"),
|
||||
format!("#!/bin/sh\nprintf '%s\\n' '{post_message}'\n"),
|
||||
)
|
||||
.expect("write post hook");
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let exec_mode = fs::Permissions::from_mode(0o755);
|
||||
fs::set_permissions(root.join("hooks").join("pre.sh"), exec_mode.clone())
|
||||
.expect("chmod pre hook");
|
||||
fs::set_permissions(root.join("hooks").join("post.sh"), exec_mode)
|
||||
.expect("chmod post hook");
|
||||
}
|
||||
fs::write(
|
||||
root.join(".claude-plugin").join("plugin.json"),
|
||||
format!(
|
||||
"{{\n \"name\": \"{name}\",\n \"version\": \"1.0.0\",\n \"description\": \"runtime hook plugin\",\n \"hooks\": {{\n \"PreToolUse\": [\"./hooks/pre.sh\"],\n \"PostToolUse\": [\"./hooks/post.sh\"]\n }}\n}}"
|
||||
),
|
||||
)
|
||||
.expect("write plugin manifest");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runs_user_to_tool_to_result_loop_end_to_end_and_tracks_usage() {
|
||||
let api_client = ScriptedApiClient { call_count: 0 };
|
||||
@@ -765,7 +508,6 @@ mod tests {
|
||||
assert_eq!(summary.tool_results.len(), 1);
|
||||
assert_eq!(runtime.session().messages.len(), 4);
|
||||
assert_eq!(summary.usage.output_tokens, 10);
|
||||
assert_eq!(summary.auto_compaction, None);
|
||||
assert!(matches!(
|
||||
runtime.session().messages[1].blocks[1],
|
||||
ContentBlock::ToolUse { .. }
|
||||
@@ -867,7 +609,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(),
|
||||
)),
|
||||
@@ -933,7 +675,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'")],
|
||||
)),
|
||||
@@ -968,153 +710,6 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn initializes_and_shuts_down_plugins_with_runtime_lifecycle() {
|
||||
let config_home = temp_dir("config");
|
||||
let source_root = temp_dir("source");
|
||||
let _ = write_lifecycle_plugin(&source_root, "runtime-lifecycle");
|
||||
|
||||
let mut manager = PluginManager::new(PluginManagerConfig::new(&config_home));
|
||||
let install = manager
|
||||
.install(source_root.to_str().expect("utf8 path"))
|
||||
.expect("install should succeed");
|
||||
let log_path = install.install_path.join("lifecycle.log");
|
||||
let registry = manager.plugin_registry().expect("registry should load");
|
||||
|
||||
{
|
||||
let runtime = ConversationRuntime::new_with_plugins(
|
||||
Session::new(),
|
||||
ScriptedApiClient { call_count: 0 },
|
||||
StaticToolExecutor::new().register("add", |_input| Ok("4".to_string())),
|
||||
PermissionPolicy::new(PermissionMode::WorkspaceWrite),
|
||||
vec!["system".to_string()],
|
||||
RuntimeFeatureConfig::default(),
|
||||
registry,
|
||||
)
|
||||
.expect("runtime should initialize plugins");
|
||||
|
||||
let log = fs::read_to_string(&log_path).expect("init log should exist");
|
||||
assert_eq!(log, "init\n");
|
||||
drop(runtime);
|
||||
}
|
||||
|
||||
let log = fs::read_to_string(&log_path).expect("shutdown log should exist");
|
||||
assert_eq!(log, "init\nshutdown\n");
|
||||
|
||||
let _ = fs::remove_dir_all(config_home);
|
||||
let _ = fs::remove_dir_all(source_root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn executes_hooks_from_installed_plugins_during_tool_use() {
|
||||
struct TwoCallApiClient {
|
||||
calls: usize,
|
||||
}
|
||||
|
||||
impl ApiClient for TwoCallApiClient {
|
||||
fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError> {
|
||||
self.calls += 1;
|
||||
match self.calls {
|
||||
1 => Ok(vec![
|
||||
AssistantEvent::ToolUse {
|
||||
id: "tool-1".to_string(),
|
||||
name: "add".to_string(),
|
||||
input: r#"{"lhs":2,"rhs":2}"#.to_string(),
|
||||
},
|
||||
AssistantEvent::MessageStop,
|
||||
]),
|
||||
2 => {
|
||||
assert!(request
|
||||
.messages
|
||||
.iter()
|
||||
.any(|message| message.role == MessageRole::Tool));
|
||||
Ok(vec![
|
||||
AssistantEvent::TextDelta("done".to_string()),
|
||||
AssistantEvent::MessageStop,
|
||||
])
|
||||
}
|
||||
_ => Err(RuntimeError::new("unexpected extra API call")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let config_home = temp_dir("hook-config");
|
||||
let first_source_root = temp_dir("hook-source-a");
|
||||
let second_source_root = temp_dir("hook-source-b");
|
||||
write_hook_plugin(
|
||||
&first_source_root,
|
||||
"first",
|
||||
"plugin pre one",
|
||||
"plugin post one",
|
||||
);
|
||||
write_hook_plugin(
|
||||
&second_source_root,
|
||||
"second",
|
||||
"plugin pre two",
|
||||
"plugin post two",
|
||||
);
|
||||
|
||||
let mut manager = PluginManager::new(PluginManagerConfig::new(&config_home));
|
||||
manager
|
||||
.install(first_source_root.to_str().expect("utf8 path"))
|
||||
.expect("first plugin install should succeed");
|
||||
manager
|
||||
.install(second_source_root.to_str().expect("utf8 path"))
|
||||
.expect("second plugin install should succeed");
|
||||
let registry = manager.plugin_registry().expect("registry should load");
|
||||
|
||||
let mut runtime = ConversationRuntime::new_with_plugins(
|
||||
Session::new(),
|
||||
TwoCallApiClient { calls: 0 },
|
||||
StaticToolExecutor::new().register("add", |_input| Ok("4".to_string())),
|
||||
PermissionPolicy::new(PermissionMode::DangerFullAccess),
|
||||
vec!["system".to_string()],
|
||||
RuntimeFeatureConfig::default(),
|
||||
registry,
|
||||
)
|
||||
.expect("runtime should load plugin hooks");
|
||||
|
||||
let summary = runtime
|
||||
.run_turn("use add", None)
|
||||
.expect("tool loop succeeds");
|
||||
|
||||
assert_eq!(summary.tool_results.len(), 1);
|
||||
let ContentBlock::ToolResult {
|
||||
is_error, output, ..
|
||||
} = &summary.tool_results[0].blocks[0]
|
||||
else {
|
||||
panic!("expected tool result block");
|
||||
};
|
||||
assert!(
|
||||
!*is_error,
|
||||
"plugin hooks should not force an error: {output:?}"
|
||||
);
|
||||
assert!(
|
||||
output.contains('4'),
|
||||
"tool output missing value: {output:?}"
|
||||
);
|
||||
assert!(
|
||||
output.contains("plugin pre one"),
|
||||
"tool output missing first pre hook feedback: {output:?}"
|
||||
);
|
||||
assert!(
|
||||
output.contains("plugin pre two"),
|
||||
"tool output missing second pre hook feedback: {output:?}"
|
||||
);
|
||||
assert!(
|
||||
output.contains("plugin post one"),
|
||||
"tool output missing first post hook feedback: {output:?}"
|
||||
);
|
||||
assert!(
|
||||
output.contains("plugin post two"),
|
||||
"tool output missing second post hook feedback: {output:?}"
|
||||
);
|
||||
|
||||
let _ = fs::remove_dir_all(config_home);
|
||||
let _ = fs::remove_dir_all(first_source_root);
|
||||
let _ = fs::remove_dir_all(second_source_root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconstructs_usage_tracker_from_restored_session() {
|
||||
struct SimpleApi;
|
||||
@@ -1203,177 +798,4 @@ mod tests {
|
||||
fn shell_snippet(script: &str) -> String {
|
||||
script.to_string()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_compacts_when_turn_input_threshold_is_crossed() {
|
||||
struct SimpleApi;
|
||||
impl ApiClient for SimpleApi {
|
||||
fn stream(
|
||||
&mut self,
|
||||
_request: ApiRequest,
|
||||
) -> Result<Vec<AssistantEvent>, RuntimeError> {
|
||||
Ok(vec![
|
||||
AssistantEvent::TextDelta("done".to_string()),
|
||||
AssistantEvent::Usage(TokenUsage {
|
||||
input_tokens: 120_000,
|
||||
output_tokens: 4,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
}),
|
||||
AssistantEvent::MessageStop,
|
||||
])
|
||||
}
|
||||
}
|
||||
|
||||
let session = Session {
|
||||
version: 1,
|
||||
messages: vec![
|
||||
crate::session::ConversationMessage::user_text("one ".repeat(30_000)),
|
||||
crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "two ".repeat(30_000),
|
||||
}]),
|
||||
crate::session::ConversationMessage::user_text("three ".repeat(30_000)),
|
||||
crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "four ".repeat(30_000),
|
||||
}]),
|
||||
],
|
||||
};
|
||||
|
||||
let mut runtime = ConversationRuntime::new(
|
||||
session,
|
||||
SimpleApi,
|
||||
StaticToolExecutor::new(),
|
||||
PermissionPolicy::new(PermissionMode::DangerFullAccess),
|
||||
vec!["system".to_string()],
|
||||
)
|
||||
.with_auto_compaction_input_tokens_threshold(100_000);
|
||||
|
||||
let summary = runtime
|
||||
.run_turn("trigger", None)
|
||||
.expect("turn should succeed");
|
||||
|
||||
assert_eq!(
|
||||
summary.auto_compaction,
|
||||
Some(AutoCompactionEvent {
|
||||
removed_message_count: 2,
|
||||
})
|
||||
);
|
||||
assert_eq!(runtime.session().messages[0].role, MessageRole::System);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_compaction_does_not_repeat_after_context_is_already_compacted() {
|
||||
struct SequentialUsageApi {
|
||||
call_count: usize,
|
||||
}
|
||||
|
||||
impl ApiClient for SequentialUsageApi {
|
||||
fn stream(
|
||||
&mut self,
|
||||
_request: ApiRequest,
|
||||
) -> Result<Vec<AssistantEvent>, RuntimeError> {
|
||||
self.call_count += 1;
|
||||
let input_tokens = if self.call_count == 1 { 120_000 } else { 64 };
|
||||
Ok(vec![
|
||||
AssistantEvent::TextDelta("done".to_string()),
|
||||
AssistantEvent::Usage(TokenUsage {
|
||||
input_tokens,
|
||||
output_tokens: 4,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
}),
|
||||
AssistantEvent::MessageStop,
|
||||
])
|
||||
}
|
||||
}
|
||||
|
||||
let session = Session {
|
||||
version: 1,
|
||||
messages: vec![
|
||||
crate::session::ConversationMessage::user_text("one ".repeat(30_000)),
|
||||
crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "two ".repeat(30_000),
|
||||
}]),
|
||||
crate::session::ConversationMessage::user_text("three ".repeat(30_000)),
|
||||
crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "four ".repeat(30_000),
|
||||
}]),
|
||||
],
|
||||
};
|
||||
|
||||
let mut runtime = ConversationRuntime::new(
|
||||
session,
|
||||
SequentialUsageApi { call_count: 0 },
|
||||
StaticToolExecutor::new(),
|
||||
PermissionPolicy::new(PermissionMode::DangerFullAccess),
|
||||
vec!["system".to_string()],
|
||||
)
|
||||
.with_auto_compaction_input_tokens_threshold(100_000);
|
||||
|
||||
let first = runtime
|
||||
.run_turn("trigger", None)
|
||||
.expect("first turn should succeed");
|
||||
assert_eq!(
|
||||
first.auto_compaction,
|
||||
Some(AutoCompactionEvent {
|
||||
removed_message_count: 2,
|
||||
})
|
||||
);
|
||||
|
||||
let second = runtime
|
||||
.run_turn("continue", None)
|
||||
.expect("second turn should succeed");
|
||||
assert_eq!(second.auto_compaction, None);
|
||||
assert_eq!(runtime.session().messages[0].role, MessageRole::System);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skips_auto_compaction_below_threshold() {
|
||||
struct SimpleApi;
|
||||
impl ApiClient for SimpleApi {
|
||||
fn stream(
|
||||
&mut self,
|
||||
_request: ApiRequest,
|
||||
) -> Result<Vec<AssistantEvent>, RuntimeError> {
|
||||
Ok(vec![
|
||||
AssistantEvent::TextDelta("done".to_string()),
|
||||
AssistantEvent::Usage(TokenUsage {
|
||||
input_tokens: 99_999,
|
||||
output_tokens: 4,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
}),
|
||||
AssistantEvent::MessageStop,
|
||||
])
|
||||
}
|
||||
}
|
||||
|
||||
let mut runtime = ConversationRuntime::new(
|
||||
Session::new(),
|
||||
SimpleApi,
|
||||
StaticToolExecutor::new(),
|
||||
PermissionPolicy::new(PermissionMode::DangerFullAccess),
|
||||
vec!["system".to_string()],
|
||||
)
|
||||
.with_auto_compaction_input_tokens_threshold(100_000);
|
||||
|
||||
let summary = runtime
|
||||
.run_turn("trigger", None)
|
||||
.expect("turn should succeed");
|
||||
assert_eq!(summary.auto_compaction, None);
|
||||
assert_eq!(runtime.session().messages.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_compaction_threshold_defaults_and_parses_values() {
|
||||
assert_eq!(
|
||||
parse_auto_compaction_threshold(None),
|
||||
DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD
|
||||
);
|
||||
assert_eq!(parse_auto_compaction_threshold(Some("4321")), 4321);
|
||||
assert_eq!(
|
||||
parse_auto_compaction_threshold(Some("not-a-number")),
|
||||
DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
use std::ffi::OsStr;
|
||||
use std::path::Path;
|
||||
use std::process::Command;
|
||||
|
||||
use serde_json::json;
|
||||
@@ -52,6 +51,16 @@ pub struct HookRunner {
|
||||
config: RuntimeHookConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct HookCommandRequest<'a> {
|
||||
event: HookEvent,
|
||||
tool_name: &'a str,
|
||||
tool_input: &'a str,
|
||||
tool_output: Option<&'a str>,
|
||||
is_error: bool,
|
||||
payload: &'a str,
|
||||
}
|
||||
|
||||
impl HookRunner {
|
||||
#[must_use]
|
||||
pub fn new(config: RuntimeHookConfig) -> Self {
|
||||
@@ -119,14 +128,16 @@ impl HookRunner {
|
||||
let mut messages = Vec::new();
|
||||
|
||||
for command in commands {
|
||||
match self.run_command(
|
||||
match Self::run_command(
|
||||
command,
|
||||
event,
|
||||
tool_name,
|
||||
tool_input,
|
||||
tool_output,
|
||||
is_error,
|
||||
&payload,
|
||||
HookCommandRequest {
|
||||
event,
|
||||
tool_name,
|
||||
tool_input,
|
||||
tool_output,
|
||||
is_error,
|
||||
payload: &payload,
|
||||
},
|
||||
) {
|
||||
HookCommandOutcome::Allow { message } => {
|
||||
if let Some(message) = message {
|
||||
@@ -150,30 +161,23 @@ impl HookRunner {
|
||||
HookRunResult::allow(messages)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments, clippy::unused_self)]
|
||||
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, request: HookCommandRequest<'_>) -> 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", request.event.as_str());
|
||||
child.env("HOOK_TOOL_NAME", request.tool_name);
|
||||
child.env("HOOK_TOOL_INPUT", request.tool_input);
|
||||
child.env(
|
||||
"HOOK_TOOL_IS_ERROR",
|
||||
if request.is_error { "1" } else { "0" },
|
||||
);
|
||||
if let Some(tool_output) = request.tool_output {
|
||||
child.env("HOOK_TOOL_OUTPUT", tool_output);
|
||||
}
|
||||
|
||||
match child.output_with_stdin(payload.as_bytes()) {
|
||||
match child.output_with_stdin(request.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();
|
||||
@@ -191,16 +195,18 @@ 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 `{}`",
|
||||
request.event.as_str(),
|
||||
request.tool_name
|
||||
),
|
||||
},
|
||||
}
|
||||
}
|
||||
Err(error) => HookCommandOutcome::Warn {
|
||||
message: format!(
|
||||
"{} hook `{command}` failed to start for `{tool_name}`: {error}",
|
||||
event.as_str()
|
||||
"{} hook `{command}` failed to start for `{}`: {error}",
|
||||
request.event.as_str(),
|
||||
request.tool_name
|
||||
),
|
||||
},
|
||||
}
|
||||
@@ -239,11 +245,7 @@ fn shell_command(command: &str) -> CommandWithStdin {
|
||||
};
|
||||
|
||||
#[cfg(not(windows))]
|
||||
let command_builder = if Path::new(command).exists() {
|
||||
let mut command_builder = Command::new("sh");
|
||||
command_builder.arg(command);
|
||||
CommandWithStdin::new(command_builder)
|
||||
} else {
|
||||
let command_builder = {
|
||||
let mut command_builder = Command::new("sh");
|
||||
command_builder.arg("-lc").arg(command);
|
||||
CommandWithStdin::new(command_builder)
|
||||
|
||||
Reference in New Issue
Block a user