#![allow(dead_code)]
use std::collections::HashSet;
use std::sync::Arc;
use uuid::Uuid;
use crate::types::Message;
use crate::utils::hooks::helpers::{add_arguments_to_prompt, hook_response_json_schema};
use crate::utils::hooks::hook_helpers::SYNTHETIC_OUTPUT_TOOL_NAME;
use crate::utils::hooks::session_hooks::clear_session_hooks;
const MAX_AGENT_TURNS: usize = 50;
pub enum HookResult {
Success {
hook_name: String,
hook_event: String,
tool_use_id: String,
},
Blocking {
blocking_error: String,
command: String,
},
Cancelled,
NonBlockingError {
hook_name: String,
hook_event: String,
tool_use_id: String,
stderr: String,
stdout: String,
exit_code: i32,
},
}
pub struct AgentHook {
pub prompt: String,
pub timeout: Option<u64>,
pub model: Option<String>,
}
pub async fn exec_agent_hook(
hook: &AgentHook,
hook_name: &str,
hook_event: &str,
json_input: &str,
signal: tokio::sync::watch::Receiver<bool>,
tool_use_context: Arc<crate::utils::hooks::can_use_tool::ToolUseContext>,
tool_use_id: Option<String>,
_messages: &[Message],
agent_name: Option<&str>,
) -> HookResult {
let effective_tool_use_id = tool_use_id.unwrap_or_else(|| format!("hook-{}", Uuid::new_v4()));
let transcript_path = format!("session_{}_transcript.json", tool_use_context.session_id);
let hook_start = std::time::Instant::now();
let processed_prompt = add_arguments_to_prompt(&hook.prompt, json_input);
log_for_debugging(&format!(
"Hooks: Processing agent hook with prompt: {}",
processed_prompt.chars().take(200).collect::<String>()
));
let user_message = create_user_message(&processed_prompt);
let mut agent_messages = vec![user_message];
log_for_debugging(&format!(
"Hooks: Starting agent query with {} messages",
agent_messages.len()
));
let hook_timeout_ms = hook.timeout.map_or(60_000, |t| t * 1000);
let (abort_tx, abort_rx) = tokio::sync::watch::channel(false);
let abort_tx_clone = abort_tx.clone();
let timeout_handle = tokio::spawn(async move {
tokio::time::sleep(tokio::time::Duration::from_millis(hook_timeout_ms)).await;
let _ = abort_tx_clone.send(true);
});
let model = hook.model.clone().unwrap_or_else(get_small_fast_model);
let hook_agent_id = format!("hook-agent-{}", Uuid::new_v4());
let agent_tool_use_context = Arc::new(crate::utils::hooks::can_use_tool::ToolUseContext {
session_id: format!("{}-{}", tool_use_context.session_id, hook_agent_id),
cwd: tool_use_context.cwd.clone(),
is_non_interactive_session: true,
options: Some(crate::utils::hooks::can_use_tool::ToolUseContextOptions {
tools: Some(Vec::new()), }),
});
register_structured_output_enforcement_impl(&hook_agent_id);
let mut structured_output_result: Option<serde_json::Value> = None;
let mut turn_count = 0;
let mut hit_max_turns = false;
for message in simulate_query_loop(&agent_messages, &transcript_path, &model).await {
if message.get("type") == Some(&serde_json::json!("stream_event"))
|| message.get("type") == Some(&serde_json::json!("stream_request_start"))
{
continue;
}
if message.get("type") == Some(&serde_json::json!("assistant")) {
turn_count += 1;
if turn_count >= MAX_AGENT_TURNS {
hit_max_turns = true;
log_for_debugging(&format!(
"Hooks: Agent turn {} hit max turns, aborting",
turn_count
));
let _ = abort_tx.send(true);
break;
}
}
if let Some(attachment) = message.get("attachment") {
if let Some(attachment_type) = attachment.get("type") {
if attachment_type == "structured_output" {
if let Some(data) = attachment.get("data") {
if let Ok(parsed) = serde_json::from_value::<
crate::utils::hooks::hook_helpers::HookResponse,
>(data.clone())
{
structured_output_result = Some(data.clone());
log_for_debugging(&format!(
"Hooks: Got structured output: {}",
serde_json::to_string(data).unwrap_or_default()
));
let _ = abort_tx.send(true);
break;
}
}
}
}
}
if *abort_rx.borrow() {
break;
}
}
timeout_handle.abort();
clear_session_hooks_impl(&hook_agent_id);
if structured_output_result.is_none() {
if hit_max_turns {
log_for_debugging(&format!(
"Hooks: Agent hook did not complete within {} turns",
MAX_AGENT_TURNS
));
log_event(
"tengu_agent_stop_hook_max_turns",
&serde_json::json!({
"duration_ms": hook_start.elapsed().as_millis(),
"turn_count": turn_count,
"agent_name": agent_name.unwrap_or("unknown"),
}),
);
return HookResult::Cancelled;
}
log_for_debugging("Hooks: Agent hook did not return structured output");
log_event(
"tengu_agent_stop_hook_error",
&serde_json::json!({
"duration_ms": hook_start.elapsed().as_millis(),
"turn_count": turn_count,
"error_type": 1, "agent_name": agent_name.unwrap_or("unknown"),
}),
);
return HookResult::Cancelled;
}
let result = structured_output_result.unwrap();
if let Some(ok) = result.get("ok").and_then(|v| v.as_bool()) {
if !ok {
let reason = result
.get("reason")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
log_for_debugging(&format!(
"Hooks: Agent hook condition was not met: {}",
reason
));
return HookResult::Blocking {
blocking_error: format!("Agent hook condition was not met: {}", reason),
command: hook.prompt.clone(),
};
}
log_for_debugging("Hooks: Agent hook condition was met");
log_event(
"tengu_agent_stop_hook_success",
&serde_json::json!({
"duration_ms": hook_start.elapsed().as_millis(),
"turn_count": turn_count,
"agent_name": agent_name.unwrap_or("unknown"),
}),
);
return HookResult::Success {
hook_name: hook_name.to_string(),
hook_event: hook_event.to_string(),
tool_use_id: effective_tool_use_id,
};
}
HookResult::Cancelled
}
fn create_user_message(content: &str) -> serde_json::Value {
serde_json::json!({
"type": "user",
"message": {
"content": content
}
})
}
fn get_small_fast_model() -> String {
"claude-3-haiku-20240307".to_string()
}
async fn simulate_query_loop(
messages: &[serde_json::Value],
transcript_path: &str,
model: &str,
) -> Vec<serde_json::Value> {
use crate::utils::hooks::hook_helpers::hook_response_schema;
let prompt = messages
.iter()
.filter_map(|m| {
Some(
m.get("message")
.and_then(|msg| msg.get("content"))
.or_else(|| m.get("content"))?
.as_str()?
.to_string(),
)
})
.collect::<Vec<String>>()
.join("\n");
let transcript_content = tokio::fs::read_to_string(transcript_path)
.await
.unwrap_or_default();
let system_prompt = format!(
"You are verifying a stop condition in Claude Code. Your task is to verify that \
the agent completed the given plan.\n\nConversation transcript:{}\n\n\
Use the transcript above to analyze the conversation history.\
Return your verification result as JSON.",
if transcript_content.is_empty() {
" (not available)".to_string()
} else {
format!("\n---\n{}\n---", transcript_content.chars().take(50000).collect::<String>())
}
);
let user_msg = serde_json::json!({
"role": "user",
"content": prompt
});
let query_messages = vec![user_msg];
let base_url = std::env::var("AI_API_BASE_URL")
.unwrap_or_else(|_| "https://api.anthropic.com".to_string());
let api_key = std::env::var("AI_AUTH_TOKEN")
.or_else(|_| std::env::var("ANTHROPIC_API_KEY"))
.or_else(|_| std::env::var("ANTHROPIC_AUTH_TOKEN"))
.ok();
if api_key.is_none() {
log_for_debugging("Hooks: No API key available, skipping agent query");
return Vec::new();
}
let api_key = api_key.unwrap();
let url = format!("{}/v1/messages", base_url);
let request_body = serde_json::json!({
"model": model,
"max_tokens": 4096,
"system": [{"type": "text", "text": system_prompt}],
"messages": query_messages,
"temperature": 0.0,
"output": {
"type": "json_schema",
"name": "hook_response",
"schema": hook_response_schema(),
"strict": true
}
});
let client = reqwest::Client::new();
let mut req_builder = client.post(&url)
.json(&request_body)
.header("Content-Type", "application/json");
if base_url.contains("anthropic.com") {
req_builder = req_builder
.header("x-api-key", &api_key)
.header("anthropic-version", "2023-06-01");
} else {
req_builder = req_builder.header("Authorization", format!("Bearer {}", api_key));
}
let mut result = Vec::new();
result.push(serde_json::json!({ "type": "assistant" }));
match req_builder.send().await {
Ok(response) => {
let status = response.status();
let body = response.text().await.unwrap_or_default();
if !status.is_success() {
log_for_debugging(&format!("Hooks: API error {}: {}", status, body));
result.push(serde_json::json!({ "type": "done" }));
return result;
}
let parsed: serde_json::Value = match serde_json::from_str(&body) {
Ok(v) => v,
Err(e) => {
log_for_debugging(&format!("Hooks: Failed to parse API response: {}", e));
result.push(serde_json::json!({ "type": "done" }));
return result;
}
};
let text = extract_text(&parsed);
if text.is_empty() {
log_for_debugging("Hooks: Empty response from model");
result.push(serde_json::json!({ "type": "done" }));
return result;
}
log_for_debugging(&format!("Hooks: Model response: {}", text));
result.push(serde_json::json!({
"type": "attachment",
"attachment": {
"type": "structured_output",
"data": serde_json::from_str::<serde_json::Value>(&text).unwrap_or_else(|_| {
serde_json::json!({"ok": false, "reason": "Failed to parse model response"})
})
}
}));
}
Err(e) => {
log_for_debugging(&format!("Hooks: Request failed: {}", e));
}
}
result.push(serde_json::json!({ "type": "done" }));
result
}
fn extract_text(response: &serde_json::Value) -> String {
if let Some(content) = response.get("choices").and_then(|c| c.as_array())
.and_then(|c| c.first())
.and_then(|c| c.get("message"))
.and_then(|m| m.get("content"))
.and_then(|c| c.as_str()) {
return content.to_string();
}
if let Some(blocks) = response.get("content").and_then(|c| c.as_array()) {
let mut texts = Vec::new();
for block in blocks {
if let Some(text) = block.get("text").and_then(|t| t.as_str()) {
texts.push(text.to_string());
}
}
if !texts.is_empty() {
return texts.join("\n");
}
}
String::new()
}
fn noop_set_app_state(_updater: &dyn Fn(&mut serde_json::Value)) {
}
fn register_structured_output_enforcement_impl(session_id: &str) {
crate::utils::hooks::hook_helpers::register_structured_output_enforcement(
&noop_set_app_state,
session_id,
);
}
fn clear_session_hooks_impl(session_id: &str) {
clear_session_hooks(&noop_set_app_state, session_id);
}
fn log_event(event_name: &str, _metadata: &serde_json::Value) {
log::debug!("Analytics event: {}", event_name);
}
fn log_for_debugging(msg: &str) {
log::debug!("{}", msg);
}