fix(TASK-035): close agent-bracket leaks and prove compression/role-hook seams

- Open the top-level agent bracket only at committed dispatch sites so
  inspection-flag early returns (--list-sessions, --info, --macro,
  --rebuild-rag, --init-*, --build-tools) never leak agent.started
- Close the agent bracket ahead of process::exit in shell_execute's
  execute branch, where the dispatch wrapper's close is unreachable
- Extract the ask-loop auto-compression seam and cover the
  session.compressed success path at both call sites with sink tests
- Prove Session.role_hooks never serializes and that injected
  role_hooks YAML is ignored on deserialization
This commit is contained in:
Dark-Alex-17 committed 2026-09-16 13:44:35 -06:00
1 parent 8710737964
commit 0c6b7c8594
3 files changed
+174 -6

No files matched your search

+31
View File
@@ -1266,6 +1266,37 @@ mod tests {
assert!(!yaml.contains("token_usage")); assert!(!yaml.contains("token_usage"));
} }
#[test]
fn session_role_hooks_is_not_serialized() {
let content = "---\nhooks:\n turn.completed:\n - name: role-probe\n command: role-cmd\n---\nPrompt";
let mut session = Session::default();
session.set_role(Role::new("hooked", content));
assert!(session.role_hooks().is_some());
let yaml = serde_yaml::to_string(&session).unwrap();
assert!(
!yaml.contains("role_hooks"),
"the role-hooks snapshot must never reach session YAML"
);
assert!(
!yaml.contains("role-cmd"),
"no hook command text may leak into session YAML under any key"
);
}
#[test]
fn session_deserialization_ignores_injected_role_hooks() {
// A session file is attacker-writable data: a `role_hooks:` key
// crafted into it must never deserialize into executable hook
// config — hooks are only ever re-read from the role file itself.
let yaml = "model: ''\nmessages: []\nrole_hooks:\n session.started:\n - name: injected\n command: touch /tmp/pwned\n";
let session: Session = serde_yaml::from_str(yaml).unwrap();
assert!(
session.role_hooks().is_none(),
"injected role_hooks in session YAML must be ignored"
);
}
#[test] #[test]
fn session_cost_accumulates_and_survives_yaml_round_trip() { fn session_cost_accumulates_and_survives_yaml_round_trip() {
let mut session = Session::default(); let mut session = Session::default();
+21 -1
View File
@@ -388,7 +388,12 @@ async fn run(
let app = Arc::clone(&ctx.app.config); let app = Arc::clone(&ctx.app.config);
ctx.use_agent(app.as_ref(), agent, session, abort_signal.clone()) ctx.use_agent(app.as_ref(), agent, session, abort_signal.clone())
.await?; .await?;
ctx.top_level_agent_started(); // The top-level agent bracket does NOT open here: the inspection
// flags below (--list-sessions, --info, --macro, --rebuild-rag,
// --init-*) return early without running the agent, and an open
// bracket on those paths would leak `agent.started` with no
// terminal event. It opens at the dispatch sites further down,
// where an agent run is actually committed.
} else { } else {
let app: Arc<AppConfig> = Arc::clone(&ctx.app.config); let app: Arc<AppConfig> = Arc::clone(&ctx.app.config);
if let Some(prompt) = &cli.temp_role { if let Some(prompt) = &cli.temp_role {
@@ -525,6 +530,9 @@ async fn run(
return Ok(()); return Ok(());
} }
if cli.execute && !is_repl { if cli.execute && !is_repl {
// Dispatch committed: open the top-level agent bracket only now,
// past every inspection-flag early return above.
ctx.top_level_agent_started();
let result = async { let result = async {
let input = create_input(&ctx, text, &cli.file, abort_signal.clone()).await?; let input = create_input(&ctx, text, &cli.file, abort_signal.clone()).await?;
shell_execute(&mut ctx, &SHELL, input, abort_signal.clone()).await shell_execute(&mut ctx, &SHELL, input, abort_signal.clone()).await
@@ -550,6 +558,12 @@ async fn run(
return acp::run_acp_server(ctx, abort_signal).await; return acp::run_acp_server(ctx, abort_signal).await;
} }
// Dispatch committed (headless directive or REPL): open the top-level
// agent bracket only now, past every inspection-flag early return
// above. The ACP server stays unbracketed deliberately. The headless
// arm closes the bracket below; the REPL closes it at teardown (or on
// an interactive agent switch).
ctx.top_level_agent_started();
match is_repl { match is_repl {
false => { false => {
let result = async { let result = async {
@@ -692,6 +706,12 @@ async fn shell_execute(
if code == 0 && app.save_shell_history { if code == 0 && app.save_shell_history {
let _ = append_to_shell_history(&shell.name, &eval_str, code); let _ = append_to_shell_history(&shell.name, &eval_str, code);
} }
// `process::exit` never returns, so the dispatch
// wrapper's bracket close in `run()` is unreachable
// from here: close the top-level agent bracket now.
// The agent run itself succeeded — the exit code
// belongs to the user's shell command, not to us.
ctx.top_level_agent_finished(None);
hooks::drain_pending(EXIT_HOOK_DRAIN_TIMEOUT).await; hooks::drain_pending(EXIT_HOOK_DRAIN_TIMEOUT).await;
process::exit(code); process::exit(code);
} }
+122 -5
View File
@@ -1644,10 +1644,7 @@ async fn ask_inner(
}; };
eprintln!("\n📢 {}", color.italic().paint("Compressing the session."),); eprintln!("\n📢 {}", color.italic().paint("Compressing the session."),);
match ctx.compress_session().await { auto_compress_session(ctx).await;
Ok(()) => hooks::fire(HookEvent::SessionCompressed, ctx, &[], None),
Err(err) => warn!("Failed to compress the session: {err}"),
}
if let Some(session) = ctx.session.as_mut() { if let Some(session) = ctx.session.as_mut() {
session.set_compressing(false); session.set_compressing(false);
} }
@@ -1698,6 +1695,18 @@ fn print_turn_divider() {
println!(); println!();
} }
/// Auto-compression seam for the ask loop: `session.compressed` fires only
/// when compression actually succeeded. A failed auto-compression is
/// survivable mid-session, so the error is logged and swallowed rather
/// than aborting the turn — unlike `.compress session`, which surfaces
/// the failure to the user.
async fn auto_compress_session(ctx: &mut RequestContext) {
match ctx.compress_session().await {
Ok(()) => hooks::fire(HookEvent::SessionCompressed, ctx, &[], None),
Err(err) => warn!("Failed to compress the session: {err}"),
}
}
fn should_continue(ctx: &RequestContext) -> bool { fn should_continue(ctx: &RequestContext) -> bool {
let config = ctx.auto_continue_config(); let config = ctx.auto_continue_config();
ctx.app.config.function_calling_support ctx.app.config.function_calling_support
@@ -2005,7 +2014,8 @@ pub fn split_args_text(line: &str, is_win: bool) -> (Vec<String>, &str) {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::config::{AppState, Session, WorkingMode}; use crate::client::{ClientConfig, Model};
use crate::config::{AppState, RoleLike, Session, WorkingMode};
use crate::hooks::{HookDef, HooksMap, test_sink}; use crate::hooks::{HookDef, HooksMap, test_sink};
use anyhow::anyhow; use anyhow::anyhow;
use serial_test::serial; use serial_test::serial;
@@ -2180,6 +2190,113 @@ mod tests {
); );
} }
/// A ctx whose compression can actually SUCCEED offline: `dry_run`
/// makes the summarization call echo instead of hitting the network,
/// the default openai client config makes the session model
/// resolvable, and the session already holds a user exchange so
/// `compress_session` has something to summarize.
fn compressible_ctx(marker: &str) -> RequestContext {
let mut hooks_map = HooksMap::default();
hooks_map.insert(
"session.compressed".to_string(),
vec![HookDef {
name: format!("{marker}-session.compressed"),
command: "true".to_string(),
}],
);
let mut app = AppState::test_default();
app.config = Arc::new(AppConfig {
hooks: hooks_map,
dry_run: true,
clients: vec![ClientConfig::default()],
..Default::default()
});
let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Repl);
let mut session = Session::default();
session.set_model(Model::new("openai", "test-compress-model"));
ctx.session = Some(session);
let input = Input::from_str(&ctx, "hello there", None).unwrap();
ctx.session
.as_mut()
.unwrap()
.add_message(&input, "hi")
.unwrap();
ctx
}
fn session_compressed_count(marker: &str) -> usize {
test_sink::snapshot()
.iter()
.filter(|capture| capture.hook_name == format!("{marker}-session.compressed"))
.count()
}
#[tokio::test]
#[serial]
async fn compress_command_success_fires_session_compressed_once() {
let _sink = test_sink::install();
let marker = "compress-ok-w2p";
let mut ctx = compressible_ctx(marker);
let abort_signal = create_abort_signal();
// Boxed: `run_repl_command`'s state machine is far larger than a test
// thread's stack.
let result = Box::pin(run_repl_command(
&mut ctx,
abort_signal,
".compress session",
))
.await;
assert!(
result.is_ok(),
"dry-run compression must succeed: {result:?}"
);
assert_eq!(
session_compressed_count(marker),
1,
"a successful `.compress session` must fire session.compressed exactly once"
);
}
#[tokio::test]
#[serial]
async fn auto_compress_success_fires_session_compressed_once() {
// Covers the ask-loop call site through its extracted seam;
// driving `needs_compression` itself requires a full LLM turn and
// is exercised end-to-end by integration coverage, not here.
let _sink = test_sink::install();
let marker = "auto-compress-ok-f8r";
let mut ctx = compressible_ctx(marker);
auto_compress_session(&mut ctx).await;
assert_eq!(
session_compressed_count(marker),
1,
"a successful auto-compression must fire session.compressed exactly once"
);
}
#[tokio::test]
#[serial]
async fn auto_compress_failure_fires_no_session_compressed_and_swallows_error() {
let _sink = test_sink::install();
let marker = "auto-compress-fail-d3k";
let mut ctx = compressible_ctx(marker);
// An empty session makes compress_session fail; the seam must
// swallow the error (no panic, no propagation) and fire nothing.
ctx.session = Some(Session::default());
auto_compress_session(&mut ctx).await;
assert_eq!(
session_compressed_count(marker),
0,
"a failed auto-compression must fire no session.compressed hook"
);
}
#[test] #[test]
fn test_process_command_line() { fn test_process_command_line() {
assert_eq!(parse_command(" ."), Some((".", None))); assert_eq!(parse_command(" ."), Some((".", None)));