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:
3 files changed
+174
-6
No files matched your search
@@ -1266,6 +1266,37 @@ mod tests {
|
||||
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]
|
||||
fn session_cost_accumulates_and_survives_yaml_round_trip() {
|
||||
let mut session = Session::default();
|
||||
|
||||
+21
-1
@@ -388,7 +388,12 @@ async fn run(
|
||||
let app = Arc::clone(&ctx.app.config);
|
||||
ctx.use_agent(app.as_ref(), agent, session, abort_signal.clone())
|
||||
.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 {
|
||||
let app: Arc<AppConfig> = Arc::clone(&ctx.app.config);
|
||||
if let Some(prompt) = &cli.temp_role {
|
||||
@@ -525,6 +530,9 @@ async fn run(
|
||||
return Ok(());
|
||||
}
|
||||
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 input = create_input(&ctx, text, &cli.file, 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;
|
||||
}
|
||||
|
||||
// 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 {
|
||||
false => {
|
||||
let result = async {
|
||||
@@ -692,6 +706,12 @@ async fn shell_execute(
|
||||
if code == 0 && app.save_shell_history {
|
||||
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;
|
||||
process::exit(code);
|
||||
}
|
||||
|
||||
+122
-5
@@ -1644,10 +1644,7 @@ async fn ask_inner(
|
||||
};
|
||||
eprintln!("\n📢 {}", color.italic().paint("Compressing the session."),);
|
||||
|
||||
match ctx.compress_session().await {
|
||||
Ok(()) => hooks::fire(HookEvent::SessionCompressed, ctx, &[], None),
|
||||
Err(err) => warn!("Failed to compress the session: {err}"),
|
||||
}
|
||||
auto_compress_session(ctx).await;
|
||||
if let Some(session) = ctx.session.as_mut() {
|
||||
session.set_compressing(false);
|
||||
}
|
||||
@@ -1698,6 +1695,18 @@ fn print_turn_divider() {
|
||||
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 {
|
||||
let config = ctx.auto_continue_config();
|
||||
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)]
|
||||
mod tests {
|
||||
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 anyhow::anyhow;
|
||||
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]
|
||||
fn test_process_command_line() {
|
||||
assert_eq!(parse_command(" ."), Some((".", None)));
|
||||
|
||||
Reference in new issue
Block a user