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"));
|
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
@@ -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
@@ -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)));
|
||||||
|
|||||||
Reference in new issue
Block a user