consciousness/src/subconscious/hook.rs

207 lines
7.1 KiB
Rust
Raw Normal View History

// hook.rs — Claude Code session hook: context injection + agent orchestration
//
// Called on each UserPromptSubmit via the poc-hook binary. Handles
// context loading, chunking, seen-set management, and delegates
// agent orchestration to AgentCycleState.
use std::collections::HashSet;
use std::fs;
use std::io::Write;
use std::path::Path;
use std::process::Command;
use std::time::Instant;
pub use crate::session::HookSession;
pub use super::subconscious::*;
const CHUNK_SIZE: usize = 9000;
/// Run the hook logic on parsed JSON input. Returns output to inject.
pub fn run_hook(input: &str) -> String {
let Some(session) = HookSession::from_json(input) else { return String::new() };
hook(&session)
}
fn chunk_context(ctx: &str, max_bytes: usize) -> Vec<String> {
let mut sections: Vec<String> = Vec::new();
let mut current = String::new();
for line in ctx.lines() {
if line.starts_with("--- ") && line.ends_with(" ---") && !current.is_empty() {
sections.push(std::mem::take(&mut current));
}
if !current.is_empty() {
current.push('\n');
}
current.push_str(line);
}
if !current.is_empty() {
sections.push(current);
}
let mut chunks: Vec<String> = Vec::new();
let mut chunk = String::new();
for section in sections {
if !chunk.is_empty() && chunk.len() + section.len() + 1 > max_bytes {
chunks.push(std::mem::take(&mut chunk));
}
if !chunk.is_empty() {
chunk.push('\n');
}
chunk.push_str(&section);
}
if !chunk.is_empty() {
chunks.push(chunk);
}
chunks
}
fn save_pending_chunks(dir: &Path, session_id: &str, chunks: &[String]) {
let chunks_dir = dir.join(format!("chunks-{}", session_id));
let _ = fs::remove_dir_all(&chunks_dir);
if chunks.is_empty() { return; }
fs::create_dir_all(&chunks_dir).ok();
for (i, chunk) in chunks.iter().enumerate() {
let path = chunks_dir.join(format!("{:04}", i));
fs::write(path, chunk).ok();
}
}
fn pop_pending_chunk(dir: &Path, session_id: &str) -> Option<String> {
let chunks_dir = dir.join(format!("chunks-{}", session_id));
if !chunks_dir.exists() { return None; }
let mut entries: Vec<_> = fs::read_dir(&chunks_dir).ok()?
.flatten()
.filter(|e| e.file_type().map(|t| t.is_file()).unwrap_or(false))
.collect();
entries.sort_by_key(|e| e.file_name());
let first = entries.first()?;
let content = fs::read_to_string(first.path()).ok()?;
fs::remove_file(first.path()).ok();
if fs::read_dir(&chunks_dir).ok().map(|mut d| d.next().is_none()).unwrap_or(true) {
fs::remove_dir(&chunks_dir).ok();
}
Some(content)
}
fn generate_cookie() -> String {
uuid::Uuid::new_v4().as_simple().to_string()[..12].to_string()
}
fn parse_seen_line(line: &str) -> &str {
line.split_once('\t').map(|(_, key)| key).unwrap_or(line)
}
pub fn load_seen(dir: &Path, session_id: &str) -> HashSet<String> {
let path = dir.join(format!("seen-{}", session_id));
if path.exists() {
fs::read_to_string(&path)
.unwrap_or_default()
.lines()
.filter(|s| !s.is_empty())
.map(|s| parse_seen_line(s).to_string())
.collect()
} else {
HashSet::new()
}
}
fn mark_seen(dir: &Path, session_id: &str, key: &str, seen: &mut HashSet<String>) {
if !seen.insert(key.to_string()) { return; }
let path = dir.join(format!("seen-{}", session_id));
if let Ok(mut f) = fs::OpenOptions::new().create(true).append(true).open(path) {
let ts = chrono::Local::now().format("%Y-%m-%dT%H:%M:%S");
writeln!(f, "{}\t{}", ts, key).ok();
}
}
/// Standalone entry point for the Claude Code hook path.
/// Loads saved state, runs cycles, saves state back.
pub fn run_agent_cycles(session: &HookSession) -> AgentCycleOutput {
let mut state = AgentCycleState::new(&session.session_id);
state.restore(&SavedAgentState::load(&session.session_id));
state.trigger(session);
state.save(&session.session_id);
state.last_output
}
fn hook(session: &HookSession) -> String {
let start_time = Instant::now();
let mut out = String::new();
let is_compaction = crate::transcript::detect_new_compaction(
&session.state_dir, &session.session_id, &session.transcript_path,
);
let cookie_path = session.path("cookie");
let is_first = !cookie_path.exists();
let log_dir = dirs::home_dir().unwrap_or_default().join(".consciousness/logs");
fs::create_dir_all(&log_dir).ok();
let log_path = log_dir.join(format!("hook-{}", session.session_id));
let Ok(mut log_f) = fs::OpenOptions::new().create(true).append(true).open(log_path) else { return Default::default(); };
let ts = chrono::Local::now().format("%Y-%m-%dT%H:%M:%S");
let _ = writeln!(log_f, "\n=== {} ({}) {} bytes ===", ts, session.hook_event, out.len());
let _ = writeln!(log_f, "is_first {is_first} is_compaction {is_compaction}");
if is_first || is_compaction {
if is_compaction {
fs::rename(&session.path("seen"), &session.path("seen-prev")).ok();
} else {
fs::remove_file(&session.path("seen")).ok();
fs::remove_file(&session.path("seen-prev")).ok();
}
fs::remove_file(&session.path("returned")).ok();
if is_first {
fs::write(&cookie_path, generate_cookie()).ok();
}
if let Ok(output) = Command::new("poc-memory").args(["admin", "load-context"]).output() {
if output.status.success() {
let ctx = String::from_utf8_lossy(&output.stdout).to_string();
if !ctx.trim().is_empty() {
let mut ctx_seen = session.seen();
for line in ctx.lines() {
if line.starts_with("--- ") && line.ends_with(" ---") {
let inner = &line[4..line.len() - 4];
if let Some(paren) = inner.rfind(" (") {
let key = inner[..paren].trim();
mark_seen(&session.state_dir, &session.session_id, key, &mut ctx_seen);
}
}
}
let chunks = chunk_context(&ctx, CHUNK_SIZE);
if let Some(first) = chunks.first() {
out.push_str(first);
}
save_pending_chunks(&session.state_dir, &session.session_id, &chunks[1..]);
}
}
}
}
if let Some(chunk) = pop_pending_chunk(&session.state_dir, &session.session_id) {
out.push_str(&chunk);
} else {
let cfg = crate::config::get();
if cfg.surface_hooks.iter().any(|h| h == &session.hook_event) {
let cycle_output = run_agent_cycles(&session);
out.push_str(&format_agent_output(&cycle_output));
}
}
let _ = write!(log_f, "{}", out);
let duration = (Instant::now() - start_time).as_secs_f64();
let _ = writeln!(log_f, "\nran in {duration:.2}s");
out
}