use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex as StdMutex, OnceLock, Weak};
use bamboo_agent_core::{Message, PendingQuestion, Session};
use bamboo_domain::session::runtime_state::{AgentRuntimeState, PlanModeState, PlanModeStatus};
use bamboo_domain::{
latest_response_occurrence, ResponseOccurrence, SessionPermissionMode,
CONSUMED_CLARIFICATION_IDS_KEY, CONSUMED_RESPONSE_OCCURRENCES_KEY,
};
use bamboo_tools::permission::{PermissionDecisionKind, PermissionDecisionReceipt, PermissionType};
use chrono::Utc;
use dashmap::mapref::entry::Entry;
use dashmap::DashMap;
use tokio::sync::{Mutex, OwnedMutexGuard};
use super::errors::RespondError;
use super::execute::mark_startup_handoff;
use super::provider_model::{derive_model_ref, persist_legacy_model_provider, persist_model_ref};
use super::repository::SessionAccess;
use super::types::RespondInput;
const CLARIFICATION_RESUME_PENDING_KEY: &str = "clarification_resume_pending";
const CONCLUSION_WITH_OPTIONS_RESUME_PENDING_KEY: &str = "conclusion_with_options_resume_pending";
type AppliedPendingResponse = (
String,
Option<PlanModeTransition>,
Vec<(PermissionType, String)>,
);
struct PendingResponseGate {
lock: Arc<Mutex<()>>,
waiters: AtomicUsize,
}
impl PendingResponseGate {
fn new() -> Self {
Self {
lock: Arc::new(Mutex::new(())),
waiters: AtomicUsize::new(0),
}
}
}
struct PendingResponseWaiter(Arc<PendingResponseGate>);
impl Drop for PendingResponseWaiter {
fn drop(&mut self) {
self.0.waiters.fetch_sub(1, Ordering::SeqCst);
}
}
fn pending_response_locks() -> &'static DashMap<String, Weak<PendingResponseGate>> {
static LOCKS: OnceLock<DashMap<String, Weak<PendingResponseGate>>> = OnceLock::new();
LOCKS.get_or_init(DashMap::new)
}
fn pending_response_lock(session_id: &str) -> Arc<PendingResponseGate> {
let locks = pending_response_locks();
locks.retain(|_, lock| lock.strong_count() > 0);
match locks.entry(session_id.to_string()) {
Entry::Occupied(mut entry) => {
if let Some(lock) = entry.get().upgrade() {
lock
} else {
let lock = Arc::new(PendingResponseGate::new());
entry.insert(Arc::downgrade(&lock));
lock
}
}
Entry::Vacant(entry) => {
let lock = Arc::new(PendingResponseGate::new());
entry.insert(Arc::downgrade(&lock));
lock
}
}
}
pub struct PendingResponseGuard {
session_id: String,
_gate: Arc<PendingResponseGate>,
_guard: OwnedMutexGuard<()>,
}
impl PendingResponseGuard {
fn ensure_session(&self, session_id: &str) -> Result<(), RespondError> {
if self.session_id == session_id {
Ok(())
} else {
Err(RespondError::InvalidResponse(
"response transaction guard belongs to a different session".to_string(),
))
}
}
}
pub async fn acquire_pending_response_guard(session_id: &str) -> PendingResponseGuard {
let gate = pending_response_lock(session_id);
gate.waiters.fetch_add(1, Ordering::SeqCst);
let waiter = PendingResponseWaiter(gate.clone());
let guard = gate.lock.clone().lock_owned().await;
drop(waiter);
PendingResponseGuard {
session_id: session_id.to_string(),
_gate: gate,
_guard: guard,
}
}
#[doc(hidden)]
pub fn pending_response_waiter_count(session_id: &str) -> usize {
pending_response_locks()
.get(session_id)
.and_then(|entry| entry.upgrade())
.map(|gate| gate.waiters.load(Ordering::SeqCst))
.unwrap_or(0)
}
pub async fn inspect_pending_response_guarded(
repo: &dyn SessionAccess,
session_id: &str,
guard: &PendingResponseGuard,
) -> Result<Option<Session>, RespondError> {
guard.ensure_session(session_id)?;
Ok(repo.inspect_for_response(session_id).await?)
}
pub const PERMISSION_REEXECUTE_METADATA_KEY: &str = "permission.reexecute_tool_call_id";
pub const PERMISSION_REEXECUTE_GENERATION_METADATA_KEY: &str =
"permission.reexecute_request_generation";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ResponseSource {
Human,
Gold,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PlanModeTransition {
Entered {
reason: Option<String>,
pre_permission_mode: String,
entered_at: chrono::DateTime<chrono::Utc>,
status: PlanModeStatus,
plan_file_path: Option<String>,
},
Exited {
approved: bool,
restored_mode: String,
plan: Option<String>,
},
}
pub async fn submit_pending_response(
repo: &dyn SessionAccess,
input: RespondInput,
) -> Result<
(
Session,
String,
Option<PlanModeTransition>,
Vec<(PermissionType, String)>,
),
RespondError,
> {
submit_pending_response_checked(repo, input, None).await
}
pub async fn submit_pending_response_checked(
repo: &dyn SessionAccess,
input: RespondInput,
expected_tool_call_id: Option<String>,
) -> Result<
(
Session,
String,
Option<PlanModeTransition>,
Vec<(PermissionType, String)>,
),
RespondError,
> {
submit_pending_response_with_source_checked(
repo,
input,
expected_tool_call_id,
ResponseSource::Human,
)
.await
}
pub async fn submit_pending_response_checked_guarded(
repo: &dyn SessionAccess,
input: RespondInput,
expected_tool_call_id: Option<String>,
guard: &PendingResponseGuard,
) -> Result<
(
Session,
String,
Option<PlanModeTransition>,
Vec<(PermissionType, String)>,
),
RespondError,
> {
submit_pending_response_with_source_checked_guarded(
repo,
input,
expected_tool_call_id,
ResponseSource::Human,
guard,
)
.await
}
pub async fn submit_pending_permission_response_checked_guarded(
repo: &dyn SessionAccess,
input: RespondInput,
expected_tool_call_id: Option<String>,
permission_receipt: PermissionDecisionReceipt,
guard: &PendingResponseGuard,
) -> Result<
(
Session,
String,
Option<PlanModeTransition>,
Vec<(PermissionType, String)>,
),
RespondError,
> {
submit_pending_response_with_source_checked_guarded_inner(
repo,
input,
expected_tool_call_id,
ResponseSource::Human,
Some(permission_receipt),
guard,
)
.await
}
pub async fn submit_pending_response_with_source(
repo: &dyn SessionAccess,
input: RespondInput,
response_source: ResponseSource,
) -> Result<
(
Session,
String,
Option<PlanModeTransition>,
Vec<(PermissionType, String)>,
),
RespondError,
> {
submit_pending_response_with_source_checked(repo, input, None, response_source).await
}
pub async fn submit_pending_response_with_source_checked(
repo: &dyn SessionAccess,
input: RespondInput,
expected_tool_call_id: Option<String>,
response_source: ResponseSource,
) -> Result<
(
Session,
String,
Option<PlanModeTransition>,
Vec<(PermissionType, String)>,
),
RespondError,
> {
let guard = acquire_pending_response_guard(&input.session_id).await;
submit_pending_response_with_source_checked_guarded(
repo,
input,
expected_tool_call_id,
response_source,
&guard,
)
.await
}
pub async fn submit_pending_response_with_source_checked_guarded(
repo: &dyn SessionAccess,
input: RespondInput,
expected_tool_call_id: Option<String>,
response_source: ResponseSource,
guard: &PendingResponseGuard,
) -> Result<
(
Session,
String,
Option<PlanModeTransition>,
Vec<(PermissionType, String)>,
),
RespondError,
> {
submit_pending_response_with_source_checked_guarded_inner(
repo,
input,
expected_tool_call_id,
response_source,
None,
guard,
)
.await
}
async fn submit_pending_response_with_source_checked_guarded_inner(
repo: &dyn SessionAccess,
input: RespondInput,
expected_tool_call_id: Option<String>,
response_source: ResponseSource,
permission_receipt: Option<PermissionDecisionReceipt>,
guard: &PendingResponseGuard,
) -> Result<
(
Session,
String,
Option<PlanModeTransition>,
Vec<(PermissionType, String)>,
),
RespondError,
> {
guard.ensure_session(&input.session_id)?;
let applied = Arc::new(StdMutex::new(None));
let mutation_outcome = applied.clone();
let mutation_input = input.clone();
let mutation_expected_tool_call_id = expected_tool_call_id.clone();
let mutation_permission_receipt = permission_receipt.clone();
let session = repo
.mutate_for_response(
&input.session_id,
Box::new(move |session| {
let outcome = apply_pending_response(
session,
&mutation_input,
mutation_expected_tool_call_id.as_deref(),
response_source,
mutation_permission_receipt.as_ref(),
)?;
*mutation_outcome
.lock()
.expect("response outcome lock poisoned") = Some(outcome);
Ok(())
}),
)
.await?
.ok_or_else(|| RespondError::NotFound(input.session_id.clone()))?;
let (user_response, plan_mode_transition, permission_grants) = applied
.lock()
.expect("response outcome lock poisoned")
.take()
.expect("successful response mutation records its outcome");
tracing::info!(
"[{}] Response processed successfully, agent loop can resume",
input.session_id
);
Ok((
session,
user_response,
plan_mode_transition,
permission_grants,
))
}
fn apply_pending_response(
session: &mut Session,
input: &RespondInput,
expected_tool_call_id: Option<&str>,
response_source: ResponseSource,
permission_receipt: Option<&PermissionDecisionReceipt>,
) -> Result<AppliedPendingResponse, RespondError> {
let pending = session
.pending_question
.take()
.ok_or(RespondError::NoPendingQuestion)?;
if let Some(expected) = expected_tool_call_id {
if pending.tool_call_id != expected {
let actual = pending.tool_call_id.clone();
session.pending_question = Some(pending);
return Err(RespondError::PendingQuestionMismatch {
expected: expected.to_string(),
actual,
});
}
}
if let Some(receipt) = permission_receipt {
if receipt.session_id != input.session_id
|| receipt.decision.request_id != pending.tool_call_id
{
session.pending_question = Some(pending);
return Err(RespondError::InvalidResponse(
"permission receipt identity does not match the pending question".to_string(),
));
}
let current_generation = session
.messages
.iter()
.rev()
.find(|message| message.tool_call_id.as_deref() == Some(pending.tool_call_id.as_str()))
.and_then(permission_request_generation);
if current_generation.as_deref() != Some(receipt.decision.request_generation.as_str()) {
session.pending_question = Some(pending);
return Err(RespondError::InvalidResponse(
"permission receipt generation does not match the pending operation".to_string(),
));
}
}
if permission_receipt.is_none() {
if let Err(error_message) = validate_pending_response(&pending, &input.user_response) {
session.pending_question = Some(pending);
return Err(RespondError::InvalidResponse(error_message));
}
}
let tool_call_id = pending.tool_call_id.clone();
tracing::debug!(
"[{}] Looking for tool result message with tool_call_id: {}",
input.session_id,
tool_call_id
);
let reviewed_plan = extract_exit_plan_from_tool_result_message(session, &tool_call_id);
let typed_permission_approved = permission_receipt.is_some_and(|receipt| {
matches!(
receipt.decision.decision,
PermissionDecisionKind::AllowOnce
| PermissionDecisionKind::AllowSession
| PermissionDecisionKind::AllowWorkspace
| PermissionDecisionKind::AllowGlobal
)
});
let permission_approved = permission_receipt
.map(|_| typed_permission_approved)
.unwrap_or_else(|| is_permission_approval(&input.user_response));
let permission_grants = if permission_approved {
extract_permission_grants_from_tool_result_message(session, &tool_call_id)
} else {
Vec::new()
};
let should_reexecute = permission_receipt
.map(|_| typed_permission_approved)
.unwrap_or(!permission_grants.is_empty());
if should_reexecute {
session.metadata.insert(
PERMISSION_REEXECUTE_METADATA_KEY.to_string(),
tool_call_id.clone(),
);
if let Some(receipt) = permission_receipt {
session.metadata.insert(
PERMISSION_REEXECUTE_GENERATION_METADATA_KEY.to_string(),
receipt.decision.request_generation.clone(),
);
} else {
session
.metadata
.remove(PERMISSION_REEXECUTE_GENERATION_METADATA_KEY);
}
} else if permission_receipt.is_some() {
session.metadata.remove(PERMISSION_REEXECUTE_METADATA_KEY);
session
.metadata
.remove(PERMISSION_REEXECUTE_GENERATION_METADATA_KEY);
}
let found = update_or_append_tool_result_message(
session,
&tool_call_id,
&input.user_response,
response_source,
);
if let Some(receipt) = permission_receipt {
debug_assert!(persist_permission_decision_receipt(
session,
&tool_call_id,
receipt
));
}
if found {
tracing::info!(
"[{}] Updated existing tool result message",
input.session_id
);
} else {
tracing::warn!(
"[{}] Tool result message not found for tool_call_id: {}, added fallback message",
input.session_id,
tool_call_id
);
}
let plan_mode_transition =
apply_plan_mode_transition(session, &pending, &input.user_response, reviewed_plan);
session.clear_pending_question();
record_consumed_clarification(session, &tool_call_id);
session.metadata.remove("runtime.suspend_reason");
session.metadata.insert(
CLARIFICATION_RESUME_PENDING_KEY.to_string(),
"true".to_string(),
);
session.metadata.insert(
CONCLUSION_WITH_OPTIONS_RESUME_PENDING_KEY.to_string(),
"true".to_string(),
);
mark_startup_handoff(session);
let request_model_ref = derive_model_ref(
input.model_ref.as_ref(),
input.provider.as_deref(),
input.model.as_deref(),
);
if let Some(model_ref) = request_model_ref.as_ref() {
persist_model_ref(session, model_ref);
} else {
persist_legacy_model_provider(session, input.model.as_deref(), input.provider.as_deref());
}
if let Some(reasoning_effort) = input.reasoning_effort {
session.reasoning_effort = Some(reasoning_effort);
}
Ok((
input.user_response.clone(),
plan_mode_transition,
permission_grants,
))
}
fn record_consumed_clarification(session: &mut Session, tool_call_id: &str) {
let mut legacy_consumed = session
.metadata
.get(CONSUMED_CLARIFICATION_IDS_KEY)
.and_then(|value| serde_json::from_str::<Vec<String>>(value).ok())
.unwrap_or_default();
legacy_consumed.retain(|existing| existing != tool_call_id);
legacy_consumed.push(tool_call_id.to_string());
if legacy_consumed.len() > 64 {
legacy_consumed.drain(..legacy_consumed.len() - 64);
}
if let Ok(serialized) = serde_json::to_string(&legacy_consumed) {
session
.metadata
.insert(CONSUMED_CLARIFICATION_IDS_KEY.to_string(), serialized);
}
let Some(occurrence) = latest_response_occurrence(session, tool_call_id) else {
return;
};
let mut consumed = session
.metadata
.get(CONSUMED_RESPONSE_OCCURRENCES_KEY)
.and_then(|value| serde_json::from_str::<Vec<ResponseOccurrence>>(value).ok())
.unwrap_or_default();
consumed.retain(|existing| existing != &occurrence);
consumed.push(occurrence);
if consumed.len() > 64 {
consumed.drain(..consumed.len() - 64);
}
if let Ok(serialized) = serde_json::to_string(&consumed) {
session
.metadata
.insert(CONSUMED_RESPONSE_OCCURRENCES_KEY.to_string(), serialized);
}
}
fn apply_plan_mode_transition(
session: &mut Session,
pending: &PendingQuestion,
user_response: &str,
reviewed_plan: Option<String>,
) -> Option<PlanModeTransition> {
match pending.tool_name.as_str() {
"EnterPlanMode" if user_response.to_lowercase().contains("enter plan mode") => {
let pre_mode = session
.agent_runtime_state
.as_ref()
.map(|state| state.effective_permission_mode().as_str().to_string())
.unwrap_or_else(|| SessionPermissionMode::Default.as_str().to_string());
let entered_at = Utc::now();
let status = PlanModeStatus::Exploring;
let runtime_state = session
.agent_runtime_state
.get_or_insert_with(|| AgentRuntimeState::new(uuid::Uuid::new_v4().to_string()));
runtime_state.plan_mode = Some(PlanModeState {
entered_at,
pre_permission_mode: pre_mode.clone(),
plan_file_path: None,
status,
});
tracing::info!(
session_id = %session.id,
"Entered plan mode"
);
Some(PlanModeTransition::Entered {
reason: Some(pending.question.clone()),
pre_permission_mode: pre_mode,
entered_at,
status,
plan_file_path: None,
})
}
"ExitPlanMode" if is_exit_plan_mode_approved(user_response) => {
let restored_mode = session
.agent_runtime_state
.as_ref()
.map(|state| state.effective_permission_mode().as_str().to_string())
.unwrap_or_else(|| "default".to_string());
if let Some(ref mut runtime_state) = session.agent_runtime_state {
runtime_state.plan_mode = None;
}
tracing::info!(
session_id = %session.id,
"Exited plan mode"
);
Some(PlanModeTransition::Exited {
approved: true,
restored_mode,
plan: reviewed_plan,
})
}
_ => None,
}
}
fn is_exit_plan_mode_approved(user_response: &str) -> bool {
let lower = user_response.to_lowercase();
lower.contains("approve") && !lower.contains("stay in plan mode")
}
pub fn validate_pending_response(
pending: &PendingQuestion,
user_response: &str,
) -> Result<(), String> {
if pending.allow_custom {
return Ok(());
}
let valid = pending.options.iter().any(|option| option == user_response);
if valid {
Ok(())
} else {
let options_str = pending.options.join(", ");
Err(format!("Response must be one of: {options_str}"))
}
}
pub fn update_or_append_tool_result_message(
session: &mut Session,
tool_call_id: &str,
user_response: &str,
response_source: ResponseSource,
) -> bool {
for message in session.messages.iter_mut().rev() {
if message.tool_call_id.as_deref() == Some(tool_call_id) {
if let Ok(payload) = serde_json::from_str::<serde_json::Value>(&message.content) {
if payload.get("status").and_then(serde_json::Value::as_str)
== Some("awaiting_permission_approval")
{
if let Some(request) = payload.get("permission_request") {
insert_message_metadata(message, "permission_request", request.clone());
}
}
}
message.content = selected_message_content(user_response, response_source);
message.tool_success = Some(true);
return true;
}
}
session.add_message(bamboo_agent_core::Message::tool_result_with_status(
tool_call_id,
selected_message_content(user_response, response_source),
true,
));
false
}
fn insert_message_metadata(message: &mut Message, key: &str, value: serde_json::Value) {
let metadata = message
.metadata
.get_or_insert_with(|| serde_json::Value::Object(Default::default()));
if !metadata.is_object() {
let previous = std::mem::replace(metadata, serde_json::Value::Object(Default::default()));
metadata
.as_object_mut()
.expect("replacement metadata is an object")
.insert("previous_metadata".to_string(), previous);
}
metadata
.as_object_mut()
.expect("message metadata is an object")
.insert(key.to_string(), value);
}
fn permission_request_generation(message: &Message) -> Option<String> {
message
.metadata
.as_ref()
.and_then(|metadata| metadata.get("permission_request"))
.and_then(|request| request.get("request_generation"))
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned)
.or_else(|| {
serde_json::from_str::<serde_json::Value>(&message.content)
.ok()?
.get("permission_request")?
.get("request_generation")?
.as_str()
.map(ToOwned::to_owned)
})
}
fn persist_permission_decision_receipt(
session: &mut Session,
tool_call_id: &str,
receipt: &PermissionDecisionReceipt,
) -> bool {
let Some(message) = session
.messages
.iter_mut()
.rev()
.find(|message| message.tool_call_id.as_deref() == Some(tool_call_id))
else {
return false;
};
if permission_request_generation(message).as_deref()
!= Some(receipt.decision.request_generation.as_str())
{
return false;
}
insert_message_metadata(
message,
"permission_decision_receipt",
serde_json::to_value(receipt).expect("permission receipt is serializable"),
);
true
}
fn selected_message_content(user_response: &str, response_source: ResponseSource) -> String {
match response_source {
ResponseSource::Human => format!("Selected response: {}", user_response),
ResponseSource::Gold => format!("Auto-selected response (gold): {}", user_response),
}
}
fn extract_exit_plan_from_tool_result_message(
session: &Session,
tool_call_id: &str,
) -> Option<String> {
let message = session
.messages
.iter()
.rev()
.find(|message| message.tool_call_id.as_deref() == Some(tool_call_id))?;
let payload = serde_json::from_str::<serde_json::Value>(&message.content).ok()?;
payload
.get("plan")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn is_permission_approval(user_response: &str) -> bool {
user_response.trim().eq_ignore_ascii_case("approve")
}
fn extract_permission_grants_from_tool_result_message(
session: &Session,
tool_call_id: &str,
) -> Vec<(PermissionType, String)> {
let message = match session
.messages
.iter()
.rev()
.find(|message| message.tool_call_id.as_deref() == Some(tool_call_id))
{
Some(message) => message,
None => return Vec::new(),
};
let payload = match serde_json::from_str::<serde_json::Value>(&message.content) {
Ok(payload) => payload,
Err(_) => return Vec::new(),
};
if payload.get("status").and_then(|value| value.as_str())
!= Some("awaiting_permission_approval")
{
return Vec::new();
}
let parse_one = |value: &serde_json::Value| -> Option<(PermissionType, String)> {
let type_value = value
.get("permission_type")
.or_else(|| value.get("type"))?
.clone();
let perm_type: PermissionType = serde_json::from_value(type_value).ok()?;
let resource = value.get("resource")?.as_str()?.trim().to_string();
if resource.is_empty() {
return None;
}
Some((perm_type, resource))
};
if let Some(array) = payload
.get("permissions")
.and_then(|value| value.as_array())
{
array.iter().filter_map(parse_one).collect()
} else {
parse_one(&payload).into_iter().collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Mutex;
use tokio::sync::Notify;
struct TestSessionAccess {
session: Mutex<Session>,
loads: AtomicUsize,
saves: AtomicUsize,
block_first_save: AtomicBool,
save_started: Notify,
release_save: Notify,
}
impl TestSessionAccess {
fn with_pending() -> Self {
Self::with_pending_and_blocked_save(false)
}
fn with_pending_and_blocked_save(block_first_save: bool) -> Self {
let mut session = Session::new("sess-1", "test-model");
session.pending_question = Some(make_pending("ConclusionWithOptions"));
Self {
session: Mutex::new(session),
loads: AtomicUsize::new(0),
saves: AtomicUsize::new(0),
block_first_save: AtomicBool::new(block_first_save),
save_started: Notify::new(),
release_save: Notify::new(),
}
}
}
#[async_trait::async_trait]
impl SessionAccess for TestSessionAccess {
async fn load_session(
&self,
_id: &str,
) -> Result<Option<Session>, super::super::errors::SessionLoadError> {
Ok(Some(self.session.lock().unwrap().clone()))
}
async fn load_or_create(
&self,
_id: &str,
_model: &str,
) -> Result<Session, super::super::errors::SessionLoadError> {
Ok(self.session.lock().unwrap().clone())
}
async fn load_merged(
&self,
_id: &str,
) -> Result<Option<Session>, super::super::errors::SessionLoadError> {
self.loads.fetch_add(1, Ordering::SeqCst);
Ok(Some(self.session.lock().unwrap().clone()))
}
async fn save_session(
&self,
session: &mut Session,
) -> Result<(), super::super::errors::SessionSaveError> {
if self.block_first_save.swap(false, Ordering::SeqCst) {
self.save_started.notify_one();
self.release_save.notified().await;
}
*self.session.lock().unwrap() = session.clone();
self.saves.fetch_add(1, Ordering::SeqCst);
Ok(())
}
async fn save_and_cache(
&self,
session: &mut Session,
) -> Result<(), super::super::errors::SessionSaveError> {
self.save_session(session).await
}
}
fn respond_input() -> RespondInput {
RespondInput {
session_id: "sess-1".to_string(),
user_response: "A".to_string(),
model: None,
model_ref: None,
provider: None,
reasoning_effort: None,
}
}
fn make_pending(tool_name: &str) -> PendingQuestion {
PendingQuestion {
tool_call_id: "call-1".to_string(),
tool_name: tool_name.to_string(),
question: "Question?".to_string(),
options: vec!["A".to_string(), "B".to_string()],
allow_custom: false,
source: bamboo_agent_core::PendingQuestionSource::PauseTool,
}
}
#[tokio::test]
async fn expected_tool_call_id_accepts_the_current_question() {
let repo = TestSessionAccess::with_pending();
submit_pending_response_checked(&repo, respond_input(), Some("call-1".to_string()))
.await
.expect("matching identity should submit");
assert_eq!(repo.saves.load(Ordering::SeqCst), 1);
assert!(repo.session.lock().unwrap().pending_question.is_none());
}
#[tokio::test]
async fn stale_tool_call_id_is_rejected_without_consuming_the_question() {
let repo = TestSessionAccess::with_pending();
let error =
submit_pending_response_checked(&repo, respond_input(), Some("stale-call".to_string()))
.await
.expect_err("stale identity must fail");
assert!(matches!(
error,
RespondError::PendingQuestionMismatch {
ref expected,
ref actual,
} if expected == "stale-call" && actual == "call-1"
));
assert_eq!(repo.saves.load(Ordering::SeqCst), 0);
assert_eq!(
repo.session
.lock()
.unwrap()
.pending_question
.as_ref()
.map(|question| question.tool_call_id.as_str()),
Some("call-1")
);
}
#[tokio::test]
async fn omitted_tool_call_guard_remains_backwards_compatible() {
let repo = TestSessionAccess::with_pending();
submit_pending_response(&repo, respond_input())
.await
.expect("legacy client should remain accepted");
assert_eq!(repo.saves.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn concurrent_responses_consume_a_pending_question_once() {
let repo = Arc::new(TestSessionAccess::with_pending_and_blocked_save(true));
let first_repo = repo.clone();
let first = tokio::spawn(async move {
submit_pending_response_checked(
first_repo.as_ref(),
respond_input(),
Some("call-1".to_string()),
)
.await
});
repo.save_started.notified().await;
let second_entered = Arc::new(Notify::new());
let second_repo = repo.clone();
let second_entered_task = second_entered.clone();
let second = tokio::spawn(async move {
second_entered_task.notify_one();
submit_pending_response_checked(
second_repo.as_ref(),
respond_input(),
Some("call-1".to_string()),
)
.await
});
second_entered.notified().await;
tokio::task::yield_now().await;
assert_eq!(
repo.loads.load(Ordering::SeqCst),
1,
"the second responder must wait before loading the pending question"
);
repo.release_save.notify_one();
first.await.unwrap().expect("first response should win");
let error = second
.await
.unwrap()
.expect_err("second response must observe the consumed question");
assert!(matches!(error, RespondError::NoPendingQuestion));
assert_eq!(repo.loads.load(Ordering::SeqCst), 2);
assert_eq!(repo.saves.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn cancelled_response_waiter_does_not_leak_diagnostics_or_the_gate() {
let session_id = "cancelled-response-waiter";
let owner = acquire_pending_response_guard(session_id).await;
let waiter = tokio::spawn(async move { acquire_pending_response_guard(session_id).await });
tokio::time::timeout(std::time::Duration::from_secs(1), async {
while pending_response_waiter_count(session_id) != 1 {
tokio::task::yield_now().await;
}
})
.await
.expect("waiter should reach the gate");
waiter.abort();
let _ = waiter.await;
assert_eq!(pending_response_waiter_count(session_id), 0);
drop(owner);
tokio::time::timeout(
std::time::Duration::from_secs(1),
acquire_pending_response_guard(session_id),
)
.await
.expect("cancelled waiter must not retain the gate");
}
#[test]
fn enter_plan_mode_activates_plan_mode_state() {
let mut session = Session::new("sess-1", "test-model");
let pending = make_pending("EnterPlanMode");
apply_plan_mode_transition(&mut session, &pending, "Enter plan mode", None);
assert!(session.agent_runtime_state.is_some());
let state = session.agent_runtime_state.unwrap();
assert!(state.plan_mode.is_some());
let plan = state.plan_mode.unwrap();
assert_eq!(plan.status, PlanModeStatus::Exploring);
assert_eq!(plan.pre_permission_mode, "default");
}
#[test]
fn enter_plan_mode_does_nothing_when_not_approved() {
let mut session = Session::new("sess-1", "test-model");
let pending = make_pending("EnterPlanMode");
apply_plan_mode_transition(&mut session, &pending, "Stay in normal mode", None);
assert!(session.agent_runtime_state.is_none());
}
#[test]
fn plan_mode_preserves_and_restores_typed_auto_request() {
let mut session = Session::new("sess-auto-plan", "test-model");
session
.agent_runtime_state
.get_or_insert_with(|| AgentRuntimeState::new("run-1"))
.set_permission_mode(SessionPermissionMode::Auto);
let enter = make_pending("EnterPlanMode");
let transition =
apply_plan_mode_transition(&mut session, &enter, "Enter plan mode", None).unwrap();
assert!(matches!(
transition,
PlanModeTransition::Entered {
ref pre_permission_mode,
..
} if pre_permission_mode == "auto"
));
assert_eq!(
session
.agent_runtime_state
.as_ref()
.unwrap()
.plan_mode
.as_ref()
.unwrap()
.pre_permission_mode,
"auto"
);
let exit = make_pending("ExitPlanMode");
let transition = apply_plan_mode_transition(
&mut session,
&exit,
"Approve (Auto mode)",
Some("Reviewed plan".to_string()),
)
.unwrap();
assert!(matches!(
transition,
PlanModeTransition::Exited {
ref restored_mode,
..
} if restored_mode == "auto"
));
assert_eq!(
session
.agent_runtime_state
.as_ref()
.unwrap()
.effective_permission_mode(),
SessionPermissionMode::Auto
);
}
#[test]
fn exit_plan_mode_does_not_restore_over_a_newer_typed_mode() {
let mut session = Session::new("sess-plan-patch", "test-model");
let state = session
.agent_runtime_state
.get_or_insert_with(|| AgentRuntimeState::new("run-1"));
state.set_permission_mode(SessionPermissionMode::Auto);
state.plan_mode = Some(PlanModeState {
entered_at: Utc::now(),
pre_permission_mode: "auto".to_string(),
plan_file_path: None,
status: PlanModeStatus::AwaitingApproval,
});
state.set_permission_mode(SessionPermissionMode::Bypass);
let transition = apply_plan_mode_transition(
&mut session,
&make_pending("ExitPlanMode"),
"Approve (Default mode)",
None,
)
.unwrap();
assert!(matches!(
transition,
PlanModeTransition::Exited {
ref restored_mode,
..
} if restored_mode == "bypass"
));
let state = session.agent_runtime_state.unwrap();
assert!(state.plan_mode.is_none());
assert_eq!(
state.effective_permission_mode(),
SessionPermissionMode::Bypass
);
}
#[test]
fn exit_plan_mode_clears_plan_mode_state() {
let mut session = Session::new("sess-1", "test-model");
session.agent_runtime_state = Some(AgentRuntimeState::new("run-1"));
session.agent_runtime_state.as_mut().unwrap().plan_mode = Some(PlanModeState {
entered_at: Utc::now(),
pre_permission_mode: "default".to_string(),
plan_file_path: None,
status: PlanModeStatus::AwaitingApproval,
});
let pending = make_pending("ExitPlanMode");
apply_plan_mode_transition(
&mut session,
&pending,
"Approve (Default mode)",
Some("Reviewed plan".to_string()),
);
assert!(session.agent_runtime_state.unwrap().plan_mode.is_none());
}
#[test]
fn exit_plan_mode_keeps_plan_mode_when_not_approved() {
let mut session = Session::new("sess-1", "test-model");
session.agent_runtime_state = Some(AgentRuntimeState::new("run-1"));
session.agent_runtime_state.as_mut().unwrap().plan_mode = Some(PlanModeState {
entered_at: Utc::now(),
pre_permission_mode: "default".to_string(),
plan_file_path: None,
status: PlanModeStatus::AwaitingApproval,
});
let pending = make_pending("ExitPlanMode");
apply_plan_mode_transition(&mut session, &pending, "Stay in plan mode", None);
assert!(session.agent_runtime_state.unwrap().plan_mode.is_some());
}
#[test]
fn exit_plan_mode_ignores_other_tools() {
let mut session = Session::new("sess-1", "test-model");
let pending = make_pending("ConclusionWithOptions");
apply_plan_mode_transition(&mut session, &pending, "Approve", None);
assert!(session.agent_runtime_state.is_none());
}
#[test]
fn is_exit_plan_mode_approved_detects_approval() {
assert!(is_exit_plan_mode_approved("Approve (Default mode)"));
assert!(is_exit_plan_mode_approved("Approve (Accept edits mode)"));
assert!(!is_exit_plan_mode_approved("Stay in plan mode"));
assert!(!is_exit_plan_mode_approved("Edit plan first"));
}
#[test]
fn extract_exit_plan_from_tool_result_message_reads_plan_payload() {
let mut session = Session::new("sess-1", "test-model");
let mut tool_message = bamboo_agent_core::Message::tool_result(
"call-1",
serde_json::json!({
"plan": "# Plan\n\n1. Step"
})
.to_string(),
);
tool_message.tool_success = Some(true);
session.add_message(tool_message);
let plan = extract_exit_plan_from_tool_result_message(&session, "call-1");
assert_eq!(plan.as_deref(), Some("# Plan\n\n1. Step"));
}
#[test]
fn selected_permission_preserves_typed_request_and_receipt_in_non_visible_metadata() {
let mut session = Session::new("sess-1", "test-model");
session.add_message(bamboo_agent_core::Message::tool_result(
"permission-1",
serde_json::json!({
"status": "awaiting_permission_approval",
"permission_request": {
"request_id": "permission-1",
"request_generation": "generation-1",
"session_id": "sess-1",
"allowed_decisions": ["allow_once", "deny_once"]
}
})
.to_string(),
));
session
.messages
.last_mut()
.expect("permission tool result")
.metadata = Some(serde_json::json!("legacy-metadata"));
assert!(update_or_append_tool_result_message(
&mut session,
"permission-1",
"Approve",
ResponseSource::Human,
));
let receipt = PermissionDecisionReceipt {
session_id: "sess-1".to_string(),
decision: bamboo_tools::permission::PermissionDecision {
request_id: "permission-1".to_string(),
request_generation: "generation-1".to_string(),
decision: bamboo_tools::permission::PermissionDecisionKind::AllowOnce,
matcher_id: None,
expected_policy_revision: Some(4),
confirm_global: false,
},
decided_at: Utc::now(),
};
assert!(persist_permission_decision_receipt(
&mut session,
"permission-1",
&receipt
));
let message = session
.messages
.iter()
.find(|message| message.tool_call_id.as_deref() == Some("permission-1"))
.expect("permission tool result");
assert_eq!(message.content, "Selected response: Approve");
assert_eq!(
message
.metadata
.as_ref()
.and_then(|metadata| metadata.get("permission_request"))
.and_then(|request| request.get("request_id"))
.and_then(serde_json::Value::as_str),
Some("permission-1")
);
assert_eq!(
message
.metadata
.as_ref()
.and_then(|metadata| metadata.get("previous_metadata")),
Some(&serde_json::json!("legacy-metadata"))
);
assert_eq!(
message
.metadata
.as_ref()
.and_then(|metadata| metadata.get("permission_decision_receipt"))
.cloned()
.and_then(|receipt| {
serde_json::from_value::<PermissionDecisionReceipt>(receipt).ok()
}),
Some(receipt)
);
}
#[test]
fn reused_tool_call_id_requires_current_permission_generation() {
let mut session = Session::new("sess-1", "test-model");
for generation in ["generation-old", "generation-current"] {
session.add_message(bamboo_agent_core::Message::tool_result(
"permission-reused",
serde_json::json!({
"status": "awaiting_permission_approval",
"permission_type": "execute_command",
"resource": format!("resource-{generation}"),
"permission_request": {
"request_id": "permission-reused",
"request_generation": generation,
"session_id": "sess-1",
"allowed_decisions": ["allow_once", "deny_once"]
}
})
.to_string(),
));
}
session.set_pending_question_with_source(
"permission-reused".to_string(),
"Bash".to_string(),
"Allow the current operation?".to_string(),
vec!["Approve".to_string(), "Deny".to_string()],
false,
bamboo_agent_core::PendingQuestionSource::PauseTool,
);
let input = RespondInput {
session_id: "sess-1".to_string(),
user_response: "Approve".to_string(),
model: None,
model_ref: None,
provider: None,
reasoning_effort: None,
};
let receipt = |generation: &str| PermissionDecisionReceipt {
session_id: "sess-1".to_string(),
decision: bamboo_tools::permission::PermissionDecision {
request_id: "permission-reused".to_string(),
request_generation: generation.to_string(),
decision: bamboo_tools::permission::PermissionDecisionKind::AllowOnce,
matcher_id: None,
expected_policy_revision: None,
confirm_global: false,
},
decided_at: Utc::now(),
};
let stale = receipt("generation-old");
assert!(matches!(
apply_pending_response(
&mut session,
&input,
Some("permission-reused"),
ResponseSource::Human,
Some(&stale),
),
Err(RespondError::InvalidResponse(message))
if message.contains("generation")
));
assert!(session.pending_question.is_some());
let current = receipt("generation-current");
apply_pending_response(
&mut session,
&input,
Some("permission-reused"),
ResponseSource::Human,
Some(¤t),
)
.expect("current generation resolves the parked operation");
assert!(session.pending_question.is_none());
assert_eq!(
session
.metadata
.get(PERMISSION_REEXECUTE_GENERATION_METADATA_KEY)
.map(String::as_str),
Some("generation-current")
);
assert!(session.messages[0].content.contains("generation-old"));
assert_eq!(session.messages[1].content, "Selected response: Approve");
assert_eq!(
session.messages[1]
.metadata
.as_ref()
.and_then(|metadata| metadata.get("permission_decision_receipt"))
.and_then(|receipt| receipt.get("decision"))
.and_then(|decision| decision.get("request_generation"))
.and_then(serde_json::Value::as_str),
Some("generation-current")
);
}
#[test]
fn typed_permission_receipt_controls_replay_independent_of_display_options() {
let session = || {
let mut session = Session::new("sess-typed", "test-model");
session.add_message(bamboo_agent_core::Message::tool_result(
"permission-localized",
serde_json::json!({
"status": "awaiting_permission_approval",
"permission_type": "execute_command",
"resource": "cargo test --workspace",
"permission_request": {
"request_id": "permission-localized",
"request_generation": "generation-localized",
"session_id": "sess-typed",
"allowed_decisions": ["allow_once", "deny_once"]
}
})
.to_string(),
));
session.set_pending_question_with_source(
"permission-localized".to_string(),
"Bash".to_string(),
"允许执行吗?".to_string(),
vec!["允许".to_string(), "拒绝".to_string()],
false,
bamboo_agent_core::PendingQuestionSource::PauseTool,
);
session
};
let receipt = |decision| PermissionDecisionReceipt {
session_id: "sess-typed".to_string(),
decision: bamboo_tools::permission::PermissionDecision {
request_id: "permission-localized".to_string(),
request_generation: "generation-localized".to_string(),
decision,
matcher_id: None,
expected_policy_revision: None,
confirm_global: false,
},
decided_at: Utc::now(),
};
let mut allowed = session();
let allow_input = RespondInput {
session_id: "sess-typed".to_string(),
user_response: "已由结构化决定允许".to_string(),
model: None,
model_ref: None,
provider: None,
reasoning_effort: None,
};
apply_pending_response(
&mut allowed,
&allow_input,
Some("permission-localized"),
ResponseSource::Human,
Some(&receipt(PermissionDecisionKind::AllowOnce)),
)
.expect("typed allow must not depend on localized display options");
assert_eq!(
allowed
.metadata
.get(PERMISSION_REEXECUTE_METADATA_KEY)
.map(String::as_str),
Some("permission-localized")
);
assert_eq!(
allowed
.metadata
.get(PERMISSION_REEXECUTE_GENERATION_METADATA_KEY)
.map(String::as_str),
Some("generation-localized")
);
let mut denied = session();
denied.metadata.insert(
PERMISSION_REEXECUTE_METADATA_KEY.to_string(),
"stale-call".to_string(),
);
denied.metadata.insert(
PERMISSION_REEXECUTE_GENERATION_METADATA_KEY.to_string(),
"stale-generation".to_string(),
);
let deny_input = RespondInput {
session_id: "sess-typed".to_string(),
user_response: "Approve".to_string(),
model: None,
model_ref: None,
provider: None,
reasoning_effort: None,
};
apply_pending_response(
&mut denied,
&deny_input,
Some("permission-localized"),
ResponseSource::Human,
Some(&receipt(PermissionDecisionKind::DenyOnce)),
)
.expect("typed deny must not depend on display text");
assert!(!denied
.metadata
.contains_key(PERMISSION_REEXECUTE_METADATA_KEY));
assert!(!denied
.metadata
.contains_key(PERMISSION_REEXECUTE_GENERATION_METADATA_KEY));
let mut legacy = session();
assert!(matches!(
apply_pending_response(
&mut legacy,
&allow_input,
Some("permission-localized"),
ResponseSource::Human,
None,
),
Err(RespondError::InvalidResponse(_))
));
assert!(legacy.pending_question.is_some());
}
}