Use shared check sandboxing (#7547)

This commit is contained in:
pakrym-oai 2025-12-04 08:34:09 -08:00 committed by GitHub
parent c4e18f1b63
commit f1b7cdc3bd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 47 additions and 77 deletions

View file

@ -154,10 +154,6 @@ impl UnifiedExecSession {
}
pub(super) async fn check_for_sandbox_denial(&self) -> Result<(), UnifiedExecError> {
if self.sandbox_type() == SandboxType::None || !self.has_exited() {
return Ok(());
}
let _ =
tokio::time::timeout(Duration::from_millis(20), self.output_notify.notified()).await;
@ -167,28 +163,40 @@ impl UnifiedExecSession {
aggregated.extend_from_slice(&chunk);
}
let aggregated_text = String::from_utf8_lossy(&aggregated).to_string();
let exit_code = self.exit_code().unwrap_or(-1);
self.check_for_sandbox_denial_with_text(&aggregated_text)
.await?;
Ok(())
}
pub(super) async fn check_for_sandbox_denial_with_text(
&self,
text: &str,
) -> Result<(), UnifiedExecError> {
let sandbox_type = self.sandbox_type();
if sandbox_type == SandboxType::None || !self.has_exited() {
return Ok(());
}
let exit_code = self.exit_code().unwrap_or(-1);
let exec_output = ExecToolCallOutput {
exit_code,
stdout: StreamOutput::new(aggregated_text.clone()),
aggregated_output: StreamOutput::new(aggregated_text.clone()),
stderr: StreamOutput::new(text.to_string()),
aggregated_output: StreamOutput::new(text.to_string()),
..Default::default()
};
if is_likely_sandbox_denied(self.sandbox_type(), &exec_output) {
if is_likely_sandbox_denied(sandbox_type, &exec_output) {
let snippet = formatted_truncate_text(
&aggregated_text,
text,
TruncationPolicy::Tokens(UNIFIED_EXEC_OUTPUT_MAX_TOKENS),
);
let message = if snippet.is_empty() {
format!("Session creation failed with exit code {exit_code}")
format!("Session exited with code {exit_code}")
} else {
snippet
};
return Err(UnifiedExecError::sandbox_denied(message, exec_output));
}
Ok(())
}

View file

@ -154,10 +154,24 @@ impl UnifiedExecSessionManager {
let output = formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens));
let has_exited = session.has_exited();
let exit_code = session.exit_code();
let sandbox_type = session.sandbox_type();
let chunk_id = generate_chunk_id();
let process_id = if has_exited {
None
let process_id = request.process_id.clone();
if has_exited {
let exit = exit_code.unwrap_or(-1);
Self::emit_exec_end_from_context(
context,
&request.command,
cwd,
output.clone(),
exit,
wall_time,
// We always emit the process ID in order to keep consistency between the Begin
// event and the End event.
Some(process_id),
)
.await;
session.check_for_sandbox_denial_with_text(&text).await?;
} else {
// Only store session if not exited.
self.store_session(
@ -166,48 +180,29 @@ impl UnifiedExecSessionManager {
&request.command,
cwd.clone(),
start,
request.process_id.clone(),
process_id,
)
.await;
Some(request.process_id.clone())
};
let original_token_count = approx_token_count(&text);
Self::emit_waiting_status(&context.session, &context.turn, &request.command).await;
};
let original_token_count = approx_token_count(&text);
let response = UnifiedExecResponse {
event_call_id: context.call_id.clone(),
chunk_id,
wall_time,
output,
process_id: process_id.clone(),
process_id: if has_exited {
None
} else {
Some(request.process_id.clone())
},
exit_code,
original_token_count: Some(original_token_count),
session_command: Some(request.command.clone()),
};
if !has_exited {
Self::emit_waiting_status(&context.session, &context.turn, &request.command).await;
}
// If the command completed during this call, emit an ExecCommandEnd via the emitter.
if has_exited {
let exit = response.exit_code.unwrap_or(-1);
Self::emit_exec_end_from_context(
context,
&request.command,
cwd,
response.output.clone(),
exit,
response.wall_time,
// We always emit the process ID in order to keep consistency between the Begin
// event and the End event.
Some(request.process_id),
)
.await;
// Exit code should always be Some
sandboxing::check_sandboxing(sandbox_type, &text, exit_code.unwrap_or_default())?;
}
Ok(response)
}
@ -714,39 +709,6 @@ impl UnifiedExecSessionManager {
}
}
mod sandboxing {
use super::*;
use crate::exec::SandboxType;
use crate::exec::is_likely_sandbox_denied;
use crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_TOKENS;
pub(crate) fn check_sandboxing(
sandbox_type: SandboxType,
text: &str,
exit_code: i32,
) -> Result<(), UnifiedExecError> {
let exec_output = ExecToolCallOutput {
exit_code,
stderr: StreamOutput::new(text.to_string()),
aggregated_output: StreamOutput::new(text.to_string()),
..Default::default()
};
if is_likely_sandbox_denied(sandbox_type, &exec_output) {
let snippet = formatted_truncate_text(
text,
TruncationPolicy::Tokens(UNIFIED_EXEC_OUTPUT_MAX_TOKENS),
);
let message = if snippet.is_empty() {
format!("Session exited with code {exit_code}")
} else {
snippet
};
return Err(UnifiedExecError::sandbox_denied(message, exec_output));
}
Ok(())
}
}
enum SessionStatus {
Alive {
exit_code: Option<i32>,