bamboo_engine/session_app/

respond.rs

//! Respond use case: submit a user response to a pending question.

use bamboo_agent_core::{PendingQuestion, Session};
use bamboo_domain::session::runtime_state::{AgentRuntimeState, PlanModeState, PlanModeStatus};
use bamboo_tools::permission::PermissionType;
use chrono::Utc;

use super::errors::RespondError;
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";

/// Session-metadata key marking a tool call that was approved through a permission
/// prompt and must be RE-EXECUTED on resume. The gated tool never actually ran
/// (the permission gate intercepted it before execution), so on approval the
/// server resume adapter re-runs it and writes the real output back — instead of
/// leaving the model to infer/fabricate it. Value = the tool_call_id.
pub const PERMISSION_REEXECUTE_METADATA_KEY: &str = "permission.reexecute_tool_call_id";

#[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>,
    },
}

/// Submit a pending response: load session, validate, update messages,
/// apply plan mode transitions, persist, and return the updated session.
///
/// The caller (handler) is responsible for auto-resume triggering.
pub async fn submit_pending_response(
    repo: &dyn SessionAccess,
    input: RespondInput,
) -> Result<
    (
        Session,
        String,
        Option<PlanModeTransition>,
        Vec<(PermissionType, String)>,
    ),
    RespondError,
> {
    submit_pending_response_with_source(repo, input, ResponseSource::Human).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,
> {
    // ---- Load session (merged for respond to pick up in-memory pending question) ----
    let mut session = repo
        .load_merged(&input.session_id)
        .await?
        .ok_or_else(|| RespondError::NotFound(input.session_id.clone()))?;

    // ---- Take pending question ----
    let pending = session
        .pending_question
        .take()
        .ok_or(RespondError::NoPendingQuestion)?;

    // ---- Validate response ----
    if let Err(error_message) = validate_pending_response(&pending, &input.user_response) {
        // Put the pending question back when validation fails.
        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);

    // Permission grants implied by approving a permission prompt. Read from the
    // (still-unmodified) synthesized tool-result payload, BEFORE it is overwritten
    // by the user's selection below.
    let permission_grants = if is_permission_approval(&input.user_response) {
        extract_permission_grants_from_tool_result_message(&session, &tool_call_id)
    } else {
        Vec::new()
    };
    if !permission_grants.is_empty() {
        // Approved a permission prompt: mark the gated tool call for re-execution
        // on resume so the operation actually runs (real output) rather than the
        // model inferring it. Consumed by the server resume adapter.
        session.metadata.insert(
            PERMISSION_REEXECUTE_METADATA_KEY.to_string(),
            tool_call_id.clone(),
        );
    }

    // ---- Update or append tool result message ----
    let found = update_or_append_tool_result_message(
        &mut session,
        &tool_call_id,
        &input.user_response,
        response_source,
    );
    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
        );
    }

    // ---- Plan mode state transitions ----
    let plan_mode_transition =
        apply_plan_mode_transition(&mut session, &pending, &input.user_response, reviewed_plan);

    // ---- Clear pending question and set resume marker ----
    session.clear_pending_question();
    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(),
    );

    // ---- Merge model/reasoning from request ----
    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(&mut session, model_ref);
    } else {
        persist_legacy_model_provider(
            &mut session,
            input.model.as_deref(),
            input.provider.as_deref(),
        );
    }
    if let Some(reasoning_effort) = input.reasoning_effort {
        session.reasoning_effort = Some(reasoning_effort);
    }

    // ---- Save ----
    repo.save_and_cache(&mut session).await?;

    tracing::info!(
        "[{}] Response processed successfully, agent loop can resume",
        input.session_id
    );

    Ok((
        session,
        input.user_response,
        plan_mode_transition,
        permission_grants,
    ))
}

/// Apply plan mode state transitions based on the pending question tool and user response.
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()
                .and_then(|s| s.plan_mode.as_ref())
                .map(|p| p.pre_permission_mode.clone())
                .unwrap_or_else(|| "default".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()
                .and_then(|state| state.plan_mode.as_ref())
                .map(|plan| plan.pre_permission_mode.clone())
                .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,
    }
}

/// Check if the user response approves exiting plan mode.
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")
}

// ---- Internal helpers ----

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 &mut session.messages {
        if message.tool_call_id.as_deref() == Some(tool_call_id) {
            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 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()
        .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)
}

/// Detect whether the user response approves a pending permission request.
///
/// Permission prompts (synthesized by the permission gate, and the
/// `request_permissions` tool) offer exactly `["Approve", "Deny"]`.
fn is_permission_approval(user_response: &str) -> bool {
    user_response.trim().eq_ignore_ascii_case("approve")
}

/// Extract the permission grants implied by an approved permission prompt.
///
/// Reads the pending tool-result message (still the synthesized
/// `awaiting_permission_approval` payload, before it is overwritten by the
/// user's selection) and returns the `(PermissionType, resource)` pairs the
/// caller should grant for the session. Handles both the single-gated-tool shape
/// (top-level `permission_type` + `resource`) and the `request_permissions` shape
/// (a `permissions` array).
fn extract_permission_grants_from_tool_result_message(
    session: &Session,
    tool_call_id: &str,
) -> Vec<(PermissionType, String)> {
    let message = match session
        .messages
        .iter()
        .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::*;

    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,
        }
    }

    #[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 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"));
    }
}