Use shared check sandboxing (#7547)
This commit is contained in:
parent
c4e18f1b63
commit
f1b7cdc3bd
2 changed files with 47 additions and 77 deletions
|
|
@ -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(())
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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>,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue