windmill/backend/windmill-worker/src/ai_executor.rs

use crate::ai::tools::{execute_tool_calls, ToolAbortHandles, ToolExecutionContext};
use crate::ai::utils::{
    add_message_to_conversation, any_tool_needs_previous_result, cleanup_mcp_clients,
    filter_schema_by_input_transforms, find_unique_tool_name, get_flow_context,
    get_flow_job_runnable_and_raw_flow, get_step_name_from_flow, load_mcp_tools,
    parse_raw_script_schema, update_flow_status_module_with_actions,
    update_flow_status_module_with_actions_success,
};
use crate::memory_oss::{read_from_memory, write_to_memory};
use crate::worker_flow::{get_previous_job_result, get_transform_context};
use async_recursion::async_recursion;
use regex::Regex;
use serde_json::value::RawValue;
use std::{collections::HashMap, sync::Arc};
use uuid::Uuid;
#[cfg(feature = "bedrock")]
use windmill_ai::ai_bedrock::check_env_credentials;
#[cfg(feature = "mcp")]
use windmill_mcp::McpClient;

#[cfg(not(feature = "mcp"))]
use crate::ai::tools::McpClientStub as McpClient;
use windmill_ai::{
    ai_providers::AIProvider,
    image_handler::upload_image_to_s3,
    providers::create_query_builder,
    proxy::{
        common_outbound_headers, needs_unavailable_oauth_exchange, retain_effective_credentials,
    },
    query_builder::{BuildRequestArgs, ParsedResponse},
    types::*,
    utils::{pinned_ai_client_for, should_use_structured_output_tool},
};
use windmill_common::{
    cache,
    client::AuthedClient,
    db::DB,
    error::{self, Error},
    flow_conversations::MessageType,
    flow_status::AgentAction,
    flows::{AgentTool, FlowModule, FlowModuleValue, InputTransform, ToolValue},
    get_latest_hash_for_path,
    jobs::JobKind,
    scripts::get_full_hub_script_by_path,
    utils::{StripPath, HTTP_CLIENT},
    worker::{to_raw_value, Connection},
};
use windmill_queue::{cancel_single_job, CanceledBy, MiniPulledJob};

use crate::{
    ai::stream_event_processor::StreamEventProcessor,
    common::{
        build_args_map, resolve_job_timeout, transform_json_value, OccupancyMetrics, StreamNotifier,
    },
    handle_child::{run_future_with_polling_update_job_poller_graceful, GracefulPollOutcome},
};

lazy_static::lazy_static! {
    static ref TOOL_NAME_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_]+$").unwrap();

    static ref AI_AGENT_TOOL_SCHEMA: Box<RawValue> = to_raw_value(&serde_json::json!({
        "type": "object",
        "properties": {
            "user_message": { "type": "string" },
        },
        "required": ["user_message"],
        "additionalProperties": false,
    }));
}

const DEFAULT_MAX_AGENT_ITERATIONS: usize = 10;
const HARD_MAX_AGENT_ITERATIONS: usize = 1000;

fn strip_system_messages(messages: &[OpenAIMessage]) -> Vec<OpenAIMessage> {
    messages
        .iter()
        .filter(|message| message.role != "system")
        .cloned()
        .collect()
}

fn strip_leading_tool_messages(messages: Vec<OpenAIMessage>) -> Vec<OpenAIMessage> {
    match messages.iter().position(|message| message.role != "tool") {
        Some(first_non_tool_index) => messages.into_iter().skip(first_non_tool_index).collect(),
        None => Vec::new(),
    }
}

fn prepare_auto_memory_messages_for_request(
    loaded_messages: &[OpenAIMessage],
    context_length: usize,
) -> Vec<OpenAIMessage> {
    let start_idx = loaded_messages.len().saturating_sub(context_length);
    strip_leading_tool_messages(loaded_messages[start_idx..].to_vec())
}

fn prepare_auto_memory_messages_for_persistence(
    all_messages: &[OpenAIMessage],
    context_length: usize,
) -> Vec<OpenAIMessage> {
    let non_system_messages = strip_system_messages(all_messages);
    let start_idx = non_system_messages.len().saturating_sub(context_length);
    non_system_messages[start_idx..].to_vec()
}

fn find_module_by_id(
    modules: &Vec<FlowModule>,
    target_id: &str,
) -> Result<Option<FlowModule>, Error> {
    let mut found: Option<FlowModule> = None;
    FlowModule::traverse_modules(modules, &mut |module| {
        if found.is_none() && module.id == target_id {
            found = Some(module.clone());
        }
        Ok(())
    })
    .map_err(|e| Error::internal_err(format!("Failed to traverse flow modules: {e}")))?;
    Ok(found)
}

async fn find_ai_agent_tool_module_in_parent_agent(
    modules: &Vec<FlowModule>,
    parent_agent_step_id: &str,
    tool_module_id: &str,
    client: &AuthedClient,
) -> Result<Option<FlowModule>, Error> {
    let Some(parent_agent_module) = find_module_by_id(modules, parent_agent_step_id)? else {
        return Ok(None);
    };

    let FlowModuleValue::AIAgent { tools, agent, .. } = parent_agent_module.get_value()? else {
        return Ok(None);
    };

    // A linked parent carries no tools on the module (they live in the resource, resolved only in
    // the main execution branch). Resolve them from the resource here too, so a nested agent tool
    // of a saved+linked agent can still be located when it runs as its own job.
    let tools = if let Some(agent_ref) = agent.as_deref() {
        let agent_path = agent_ref
            .trim_start_matches("$res:")
            .trim_start_matches("res://");
        // Definitions only: resolving their defaults here would hit the same inaccessible resources.
        let resource_value = client
            .get_resource_value::<serde_json::Value>(agent_path)
            .await
            .map_err(|e| {
                Error::internal_err(format!(
                    "failed to load ai_agent resource {agent_path}: {e}"
                ))
            })?;
        match resource_value {
            serde_json::Value::Object(mut map) => match map.remove("tools") {
                Some(t) => serde_json::from_value::<Vec<AgentTool>>(t).map_err(|e| {
                    Error::internal_err(format!(
                        "invalid tools in ai_agent resource {agent_path}: {e}"
                    ))
                })?,
                None => Vec::new(),
            },
            _ => Vec::new(),
        }
    } else {
        tools
    };

    for tool in tools {
        if tool.id == tool_module_id {
            return Ok(Option::<FlowModule>::from(&tool));
        }
    }

    Ok(None)
}

/// Resolve the `description` sent to the model for an AI agent tool, in priority order:
/// an explicit per-tool description, then one auto-derived from the underlying runnable,
/// then the tool name as the historical last-resort fallback. Blank/whitespace-only values
/// at each level are skipped so a lower-priority source can still apply.
fn resolve_tool_description(
    user_description: Option<String>,
    derived_description: Option<String>,
    tool_name: &str,
) -> String {
    fn non_empty(value: Option<String>) -> Option<String> {
        value
            .map(|s| s.trim().to_string())
            .filter(|s| !s.is_empty())
    }

    non_empty(user_description)
        .or_else(|| non_empty(derived_description))
        .unwrap_or_else(|| tool_name.to_string())
}

/// Fetch a workspace script's stored description by hash, used to auto-derive an AI agent
/// tool's description when the user did not provide an explicit one. Returns `None` when the
/// script has no description or on any lookup error, so the caller falls back to the tool name.
async fn fetch_script_description(db: &DB, w_id: &str, hash: i64) -> Option<String> {
    sqlx::query_scalar!(
        "SELECT description FROM script WHERE hash = $1 AND workspace_id = $2",
        hash,
        w_id,
    )
    .fetch_optional(db)
    .await
    .ok()
    .flatten()
    .map(|d| d.trim().to_string())
    .filter(|d| !d.is_empty())
}

/// Overlay a linked step's host-local tool wiring onto the agent resource's tools. For each tool
/// id present in `tool_inputs`, merge its per-input transforms into that tool's `input_transforms`
/// (step wins). Only `FlowModule` tools carry input transforms; MCP/websearch tools are skipped.
fn overlay_tool_inputs(
    tools: &mut [AgentTool],
    tool_inputs: &HashMap<String, HashMap<String, InputTransform>>,
) {
    if tool_inputs.is_empty() {
        return;
    }
    for tool in tools.iter_mut() {
        let Some(overrides) = tool_inputs.get(&tool.id) else {
            continue;
        };
        let ToolValue::FlowModule(fmv) = &mut tool.value else {
            continue;
        };
        let input_transforms = match fmv {
            FlowModuleValue::Script { input_transforms, .. }
            | FlowModuleValue::RawScript { input_transforms, .. }
            | FlowModuleValue::FlowScript { input_transforms, .. }
            | FlowModuleValue::AIAgent { input_transforms, .. } => input_transforms,
            _ => continue,
        };
        for (key, transform) in overrides {
            input_transforms.insert(key.clone(), transform.clone());
        }
    }
}

pub async fn handle_ai_agent_job(
    // connection
    conn: &Connection,
    db: &DB,

    // agent job
    job: &MiniPulledJob,

    // job execution context
    client: &AuthedClient,
    canceled_by: &mut Option<CanceledBy>,
    mem_peak: &mut i32,
    occupancy_metrics: &mut OccupancyMetrics,
    worker_dir: &str,
    base_internal_url: &str,
    worker_name: &str,
    hostname: &str,
    killpill_rx: &mut tokio::sync::broadcast::Receiver<()>,
    has_stream: &mut bool,
) -> Result<Box<RawValue>, Error> {
    // build_args_map returns None if no $res:/$var: transforms needed, in which case use original args
    let local_args = match build_args_map(job, client, conn).await? {
        Some(transformed) => transformed,
        None => job.args.as_ref().map(|a| a.0.clone()).unwrap_or_default(),
    };

    // Handle dry_run mode - check credentials without making API calls.
    // The credentials check is always invoked inline (provider present, no agent link and no
    // parent flow), so it resolves before any flow/agent-resource context is fetched.
    let is_credentials_check = local_args
        .get("credentials_check")
        .map(|v| v.get().trim() == "true")
        .unwrap_or(false);
    if is_credentials_check {
        let args = serde_json::from_str::<AIAgentArgs>(&serde_json::to_string(&local_args)?)?;
        return handle_credentials_check(&args.provider).await;
    }

    // flow_step_id is set by the flow executor for top-level AI agents.
    // For nested AI agent tools, it's not set (to avoid triggering flow step
    // machinery on a parent that has no v2_job_status row), so we extract the
    // tool module ID from the runnable_path which has the form ".../tools/{id}".
    let flow_step_id = job
        .flow_step_id
        .as_deref()
        .or_else(|| job.runnable_path().rsplit_once("/tools/").map(|(_, id)| id))
        .ok_or_else(|| Error::internal_err("AI agent job has no flow step id".to_string()))?
        .to_string();
    let flow_step_id = &flow_step_id;

    let Some(immediate_parent_job) = &job.parent_job else {
        return Err(Error::internal_err(
            "AI agent job has no parent job".to_string(),
        ));
    };

    let mut flow_job_id = *immediate_parent_job;
    let mut flow_job = get_flow_job_runnable_and_raw_flow(db, &flow_job_id).await?;
    let direct_parent_job_kind = flow_job.kind;
    let direct_parent_job_flow_step_id = flow_job.flow_step_id.clone();

    // If the direct parent is an AI agent (nested tool case), go one level up to the flow.
    if flow_job.kind == JobKind::AIAgent {
        let Some(parent_job_id) = flow_job.parent_job else {
            return Err(Error::internal_err(
                "AI agent parent has no parent job".to_string(),
            ));
        };
        flow_job_id = parent_job_id;
        flow_job = get_flow_job_runnable_and_raw_flow(db, &flow_job_id).await?;

        if !matches!(
            flow_job.kind,
            JobKind::Flow | JobKind::FlowNode | JobKind::FlowPreview
        ) {
            return Err(Error::internal_err(
                "AI agent nesting beyond 2 levels is not supported. \
                 Only flow → agent → nested agent tool is allowed."
                    .to_string(),
            ));
        }
    }

    let flow_data = match flow_job.kind {
        JobKind::Flow | JobKind::FlowNode => {
            cache::job::fetch_flow(db, &flow_job.kind, flow_job.runnable_id).await?
        }
        JobKind::FlowPreview => {
            cache::job::fetch_preview_flow(db, &flow_job_id, flow_job.raw_flow).await?
        }
        _ => {
            return Err(Error::internal_err(
                "expected parent flow, flow preview or flow node for ai agent job".to_string(),
            ));
        }
    };

    let value = flow_data.value();

    let module = if direct_parent_job_kind == JobKind::AIAgent {
        let parent_agent_step_id = direct_parent_job_flow_step_id.as_deref().ok_or_else(|| {
            Error::internal_err("Parent AI agent job has no flow_step_id".to_string())
        })?;
        find_ai_agent_tool_module_in_parent_agent(
            &value.modules,
            parent_agent_step_id,
            flow_step_id,
            client,
        )
        .await?
    } else {
        find_module_by_id(&value.modules, flow_step_id)?
    };

    let Some(module) = module else {
        return Err(Error::internal_err(
            "AI agent module not found in flow".to_string(),
        ));
    };

    let summary = module.summary.clone();

    let FlowModuleValue::AIAgent {
        tools: module_tools,
        omit_output_from_conversation,
        agent,
        tool_inputs,
        ..
    } = module.get_value()?
    else {
        return Err(Error::internal_err(
            "AI agent module is not an AI agent".to_string(),
        ));
    };

    // A linked step takes its brain and tools from the resource and keeps only the flow-local
    // inputs (user_message/user_attachments) of its own; both stay rigid, so the one thing it may
    // bind to this flow is the tools' inputs, overlaid from `tool_inputs` below.
    let (args, tools): (AIAgentArgs, Vec<AgentTool>) = if let Some(agent_ref) = agent.as_deref() {
        let agent_path = agent_ref
            .trim_start_matches("$res:")
            .trim_start_matches("res://");
        // Read raw and interpolate only the brain below. Interpolating the whole resource would also
        // resolve each tool's default `$res:`/`$var:`, which a host flow may be overriding and which
        // may be unreadable to whoever runs this flow — an unused tool could then fail the agent.
        let resource_value = client
            .get_resource_value::<serde_json::Value>(agent_path)
            .await
            .map_err(|e| {
                Error::internal_err(format!(
                    "failed to load ai_agent resource {agent_path}: {e}"
                ))
            })?;
        let mut config = match resource_value {
            serde_json::Value::Object(map) => map,
            _ => {
                return Err(Error::internal_err(format!(
                    "ai_agent resource {agent_path} must be a JSON object"
                )))
            }
        };
        let mut tools = match config.remove("tools") {
            Some(t) => serde_json::from_value::<Vec<AgentTool>>(t).map_err(|e| {
                Error::internal_err(format!(
                    "invalid tools in ai_agent resource {agent_path}: {e}"
                ))
            })?,
            None => Vec::new(),
        };
        overlay_tool_inputs(&mut tools, &tool_inputs);
        let brain = transform_json_value(
            "ai_agent",
            client,
            &job.workspace_id,
            serde_json::Value::Object(config),
            job,
            conn,
            0,
        )
        .await?;
        let mut brain = match brain {
            serde_json::Value::Object(map) => map,
            _ => {
                return Err(Error::internal_err(format!(
                    "ai_agent resource {agent_path} must be a JSON object"
                )))
            }
        };
        // Only after interpolating the resource: these are caller-controlled and already resolved by
        // build_args_map, so passing them through it again would expand contextual values —
        // `$WM_TOKEN` in a user message would reach the model provider.
        for key in ["user_message", "user_attachments"] {
            if let Some(v) = local_args.get(key) {
                brain.insert(
                    key.to_string(),
                    serde_json::from_str(v.get()).unwrap_or(serde_json::Value::Null),
                );
            }
        }
        let args = serde_json::from_value::<AIAgentArgs>(serde_json::Value::Object(brain))
            .map_err(|e| {
                Error::internal_err(format!(
                    "invalid ai_agent resource config {agent_path}: {e}"
                ))
            })?;
        (args, tools)
    } else {
        let args = serde_json::from_str::<AIAgentArgs>(&serde_json::to_string(&local_args)?)?;
        // "Edit" on a linked step clears `agent` but keeps the host's `tool_inputs` until Save or
        // Cancel folds them back, so overlay them here too: a flow persisted mid-edit must still
        // bind its tools to this flow's context rather than the agent author's.
        let mut tools = module_tools;
        overlay_tool_inputs(&mut tools, &tool_inputs);
        (args, tools)
    };

    // Nesting is capped at flow → agent → nested agent. When this job is itself a nested tool,
    // a linked resource's tool set may still contain AIAgent tools (the editor can't constrain a
    // shared resource); don't advertise them — invoking one would only fail the depth check as a
    // third-level agent.
    let tools = if direct_parent_job_kind == JobKind::AIAgent {
        tools
            .into_iter()
            .filter(|t| {
                !matches!(
                    &t.value,
                    ToolValue::FlowModule(FlowModuleValue::AIAgent { .. })
                )
            })
            .collect()
    } else {
        tools
    };

    // Separate Windmill tools from MCP tools, websearch, and extract MCP resource configs
    let mut windmill_modules: Vec<FlowModule> = Vec::new();
    // Explicit per-tool descriptions keyed by tool id. When set, these override the
    // description auto-derived from the underlying runnable when building tool definitions.
    let mut tool_descriptions: HashMap<String, String> = HashMap::new();
    #[allow(unused_mut)]
    let mut mcp_configs: Vec<crate::ai::utils::McpResourceConfig> = Vec::new();
    let mut has_websearch = false;

    for tool in tools {
        match &tool.value {
            #[allow(unused_variables)]
            ToolValue::Mcp(mcp_config) => {
                #[cfg(feature = "mcp")]
                {
                    // This is an MCP tool - extract config
                    tracing::debug!(
                        "MCP server module: path={}, include={:?}, exclude={:?}",
                        mcp_config.resource_path,
                        mcp_config.include_tools,
                        mcp_config.exclude_tools
                    );
                    mcp_configs.push(crate::ai::utils::McpResourceConfig {
                        resource_path: mcp_config.resource_path.clone(),
                        include_tools: Some(mcp_config.include_tools.clone()),
                        exclude_tools: Some(mcp_config.exclude_tools.clone()),
                    });
                }

                #[cfg(not(feature = "mcp"))]
                {
                    tracing::warn!("MCP tool detected but MCP feature is not enabled");
                }
            }
            ToolValue::FlowModule(_) => {
                // Regular Windmill flow module (script, flow, etc.) - convert to FlowModule
                tracing::debug!("Windmill module: {:?}", tool.id);
                if let Some(description) = tool
                    .description
                    .as_ref()
                    .map(|d| d.trim())
                    .filter(|d| !d.is_empty())
                {
                    tool_descriptions.insert(tool.id.clone(), description.to_string());
                }
                if let Some(flow_module) = Option::<FlowModule>::from(&tool) {
                    windmill_modules.push(flow_module);
                }
            }
            ToolValue::Websearch(_) => {
                // WebSearch tool - mark as enabled
                tracing::debug!("WebSearch tool enabled");
                has_websearch = true;
            }
        }
    }

    // Process Windmill flow modules into Tool definitions
    let tools = futures::future::try_join_all(windmill_modules.into_iter().map(|mut t| {
        let conn = conn;
        let db = db;
        let job = job;
        let user_description = tool_descriptions.get(&t.id).cloned();
        async move {
            let Some(summary) = t.summary.as_ref().filter(|s| TOOL_NAME_REGEX.is_match(s)) else {
                return Err(Error::internal_err(format!(
                    "Invalid tool name: {:?}",
                    t.summary
                )));
            };

            // Extract schema, input_transforms, and an auto-derived description from the module value
            let module_value = t.get_value()?;
            let (schema, input_transforms, derived_description) = match &module_value {
                FlowModuleValue::Script {
                    hash,
                    path,
                    tag_override,
                    input_transforms,
                    is_trigger,
                    pass_flow_input_directly,
                } => {
                    let derived_description: Option<String>;
                    let schema = match hash {
                        Some(hash) => {
                            let (_, metadata) = cache::script::fetch(conn, hash.clone()).await?;
                            derived_description =
                                fetch_script_description(db, &job.workspace_id, hash.0).await;
                            Ok::<_, Error>(
                                metadata
                                    .schema
                                    .clone()
                                    .map(|s| RawValue::from_string(s).ok())
                                    .flatten(),
                            )
                        }
                        None => {
                            if path.starts_with("hub/") {
                                let hub_script = get_full_hub_script_by_path(
                                    StripPath(path.to_string()),
                                    &HTTP_CLIENT,
                                    None,
                                )
                                .await?;
                                // Hub scripts carry their free-text description in `summary`.
                                derived_description = hub_script
                                    .summary
                                    .as_ref()
                                    .map(|s| s.trim().to_string())
                                    .filter(|s| !s.is_empty());
                                Ok(Some(hub_script.schema))
                            } else {
                                let hash = get_latest_hash_for_path(
                                    db,
                                    &job.workspace_id,
                                    path.as_str(),
                                    true,
                                )
                                .await?
                                .0;
                                // update module definition to use a fixed hash so all tool calls match the same schema
                                t.value = to_raw_value(&FlowModuleValue::Script {
                                    hash: Some(hash),
                                    path: path.clone(),
                                    tag_override: tag_override.clone(),
                                    input_transforms: input_transforms.clone(),
                                    is_trigger: *is_trigger,
                                    pass_flow_input_directly: *pass_flow_input_directly,
                                });
                                derived_description =
                                    fetch_script_description(db, &job.workspace_id, hash.0).await;
                                let (_, metadata) = cache::script::fetch(conn, hash).await?;
                                Ok(metadata
                                    .schema
                                    .clone()
                                    .map(|s| RawValue::from_string(s).ok())
                                    .flatten())
                            }
                        }
                    }?;
                    (schema, input_transforms, derived_description)
                }
                FlowModuleValue::RawScript { content, language, input_transforms, .. } => {
                    let schema = Some(parse_raw_script_schema(&content, &language)?);
                    (schema, input_transforms, None)
                }
                FlowModuleValue::AIAgent { input_transforms, .. } => {
                    // By convention for AIAgent tools, only user_message is expected to be AI-filled.
                    (
                        Some(
                            RawValue::from_string(AI_AGENT_TOOL_SCHEMA.get().to_string())
                                .expect("AI_AGENT_TOOL_SCHEMA should always be valid JSON"),
                        ),
                        input_transforms,
                        None,
                    )
                }
                _ => {
                    return Err(Error::internal_err(format!(
                        "Unsupported tool: {}",
                        summary
                    )));
                }
            };

            // Filter schema based on user given input transforms
            let schema = if let Some(s) = schema {
                Some(filter_schema_by_input_transforms(s, input_transforms)?)
            } else {
                None
            };

            let description =
                resolve_tool_description(user_description, derived_description, summary);

            Ok(Tool {
                def: ToolDef {
                    r#type: "function".to_string(),
                    function: ToolDefFunction {
                        name: summary.clone(),
                        description: Some(description),
                        parameters: schema.unwrap_or_else(|| {
                            to_raw_value(&serde_json::json!({
                                "type": "object",
                                "properties": {},
                                "required": [],
                            }))
                        }),
                    },
                },
                module: Some(t),
                mcp_source: None,
            })
        }
    }))
    .await?;

    // Load MCP tools if configured
    let mut tools = tools;

    let mcp_clients = if !mcp_configs.is_empty() {
        let (clients, mcp_tools) =
            load_mcp_tools(db, &job.workspace_id, mcp_configs, client).await?;
        tools.extend(mcp_tools);
        clients
    } else {
        HashMap::new()
    };

    let mut inner_occupancy_metrics = occupancy_metrics.clone();

    let stream_notifier = StreamNotifier::new(conn, job);

    if let Some(stream_notifier) = stream_notifier {
        stream_notifier.update_flow_status_with_stream_job();
    }

    let flow_status_job = if direct_parent_job_kind == JobKind::AIAgent {
        None
    } else {
        Some(flow_job_id)
    };

    // Create cancellation signal for graceful shutdown
    let (cancel_tx, cancel_rx) = tokio::sync::watch::channel(false);
    let tool_abort_handles: ToolAbortHandles = Arc::new(std::sync::Mutex::new(Vec::new()));

    /// Grace period for in-flight tool calls to complete after cancellation.
    const CANCEL_GRACE_PERIOD: std::time::Duration = std::time::Duration::from_secs(30);

    let outcome = {
        let agent_fut = run_agent(
            db,
            conn,
            job,
            flow_status_job.as_ref(),
            Some(flow_step_id.as_str()),
            &args,
            &tools,
            &mcp_clients,
            summary.as_deref(),
            client,
            &mut inner_occupancy_metrics,
            worker_dir,
            base_internal_url,
            worker_name,
            hostname,
            killpill_rx,
            has_stream,
            has_websearch,
            omit_output_from_conversation,
            cancel_rx,
            tool_abort_handles.clone(),
        );

        let mut occupancy_opt = Some(occupancy_metrics);

        run_future_with_polling_update_job_poller_graceful(
            job.id,
            job.timeout,
            conn,
            mem_peak,
            canceled_by,
            agent_fut,
            worker_name,
            &job.workspace_id,
            &mut occupancy_opt,
            Box::pin(futures::stream::once(async { 0 })),
            cancel_tx,
            CANCEL_GRACE_PERIOD,
        )
        .await?
    };
    // agent_fut and update_job are now dropped — borrows on mcp_clients and canceled_by released

    // Cleanup MCP clients
    cleanup_mcp_clients(mcp_clients).await;

    let format_cancel_info = |cb: &Option<CanceledBy>| {
        cb.as_ref()
            .map_or(("unknown".to_string(), "unknown".to_string()), |x| {
                (
                    x.username.clone().unwrap_or_default(),
                    x.reason.clone().unwrap_or_default(),
                )
            })
    };

    match outcome {
        GracefulPollOutcome::Ok(result) => Ok(result),
        GracefulPollOutcome::Timeout(ms) => {
            tracing::error!("AI agent timeout after {}s", ms / 1000);
            Err(Error::ExecutionErr(format!(
                "AI agent timeout after (>{}s)",
                ms / 1000
            )))
        }
        GracefulPollOutcome::Cancelled { canceled_by: cb } => {
            let (by, reason) = format_cancel_info(&cb);
            Err(Error::ExecutionErr(format!(
                "Job cancelled by {by} (reason: {reason})"
            )))
        }
        GracefulPollOutcome::CancelledTimeout { canceled_by: cb } => {
            let (by, reason) = format_cancel_info(&cb);
            // Abort any still-running spawned tool tasks
            // unwrap safe: lock is only held briefly for push/drain, no panic possible inside
            for handle in tool_abort_handles.lock().unwrap().drain(..) {
                handle.abort();
            }
            // Hard timeout: clean up orphaned jobs still stuck in v2_job_queue
            cleanup_orphaned_tool_jobs(db, &job.id, &job.workspace_id, cb).await;
            Err(Error::ExecutionErr(format!(
                "Job cancelled by {by} (reason: {reason}, timed out waiting for tool calls)"
            )))
        }
        GracefulPollOutcome::AlreadyCompleted => {
            Err(Error::AlreadyCompleted("Job already completed".to_string()))
        }
    }
}

#[async_recursion]
pub async fn run_agent(
    // connection
    db: &DB,
    conn: &Connection,

    // agent job and flow data
    job: &MiniPulledJob,
    parent_job: Option<&Uuid>,
    flow_step_id_override: Option<&str>,
    args: &AIAgentArgs,
    tools: &[Tool],
    mcp_clients: &HashMap<String, Arc<McpClient>>,
    summary: Option<&str>,

    // job execution context
    client: &AuthedClient,
    occupancy_metrics: &mut OccupancyMetrics,
    worker_dir: &str,
    base_internal_url: &str,
    worker_name: &str,
    hostname: &str,
    killpill_rx: &mut tokio::sync::broadcast::Receiver<()>,
    has_stream: &mut bool,
    has_websearch: bool,
    omit_output_from_conversation: bool,

    // cancellation signal from parent
    cancel_rx: tokio::sync::watch::Receiver<bool>,

    // abort handles for spawned tool tasks
    tool_abort_handles: ToolAbortHandles,
) -> error::Result<Box<RawValue>> {
    let output_type = args.output_type.as_ref().unwrap_or(&OutputType::Text);
    let credentials = args.provider.to_provider_credentials(db).await?;
    let base_url = &credentials.base_url;
    let api_key = credentials.api_key.as_deref().unwrap_or("");

    // Create the query builder for the provider
    let query_builder = create_query_builder(&credentials, args.provider.get_model());

    // Initialize messages
    let mut messages =
        if let Some(system_prompt) = args.system_prompt.clone().filter(|s| !s.is_empty()) {
            vec![OpenAIMessage {
                role: "system".to_string(),
                content: Some(OpenAIContent::Text(system_prompt)),
                ..Default::default()
            }]
        } else {
            vec![]
        };

    // Effective flow_step_id: override for nested agents, otherwise from job
    let effective_flow_step_id: Option<&str> =
        flow_step_id_override.or(job.flow_step_id.as_deref());

    // Fetch flow context for input transforms context, chat and memory
    let mut flow_context = get_flow_context(db, job).await;

    // Determine if we're using manual messages (which bypasses memory)
    let use_manual_messages = matches!(args.memory, Some(Memory::Manual { .. }));

    // Check if user_message is provided and non-empty
    let has_user_message = args
        .user_message
        .as_ref()
        .map(|m| !m.is_empty())
        .unwrap_or(false);

    // Validate: at least one of memory with manual messages or user_message must be provided
    if !use_manual_messages && !has_user_message {
        return Err(Error::internal_err(
            "Either 'memory' with manual messages or 'user_message' must be provided".to_string(),
        ));
    }

    let is_text_output = output_type == &OutputType::Text;

    // Flow-level memory_id (from chat mode) takes precedence over step-level memory_id
    let memory_id = flow_context
        .flow_status
        .as_ref()
        .and_then(|fs| fs.memory_id)
        .or_else(|| {
            // Extract memory_id from Memory::Auto if present
            match &args.memory {
                Some(Memory::Auto { memory_id, .. }) => *memory_id,
                _ => None,
            }
        });

    // Load messages based on history mode
    if matches!(output_type, OutputType::Text) {
        match &args.memory {
            Some(Memory::Manual { messages: manual_messages }) => {
                // Use explicitly provided messages (bypass memory)
                if !manual_messages.is_empty() {
                    messages.extend(manual_messages.clone());
                }
            }
            Some(Memory::Auto { context_length, .. }) => {
                // Auto mode: load from memory
                if let Some(step_id) = effective_flow_step_id {
                    if let Some(memory_id) = memory_id {
                        // Read messages from memory
                        match read_from_memory(db, &job.workspace_id, memory_id, step_id).await {
                            Ok(Some(loaded_messages)) => {
                                let messages_to_load = prepare_auto_memory_messages_for_request(
                                    &loaded_messages,
                                    *context_length,
                                );
                                messages.extend(messages_to_load);
                            }
                            Ok(None) => {}
                            Err(e) => {
                                tracing::error!(
                                    "Failed to read memory for step {}: {}",
                                    step_id,
                                    e
                                );
                            }
                        }
                    }
                }
            }
            _ => {}
        }
    }

    // Extract previous step result only if any tool needs it
    let previous_result = {
        if any_tool_needs_previous_result(&tools) {
            if let Some(ref flow_status) = flow_context.flow_status {
                get_previous_job_result(db, &job.workspace_id, flow_status)
                    .await
                    .ok()
                    .flatten()
            } else {
                None
            }
        } else {
            None
        }
    };

    // Build IdContext for results.stepId syntax
    let id_context = {
        if let Some(ref flow_status) = flow_context.flow_status {
            // Get the step ID from the AI agent's flow step
            let previous_id = effective_flow_step_id
                .map(str::to_string)
                .unwrap_or_else(|| "unknown".to_string());

            Some(get_transform_context(job, &previous_id, flow_status))
        } else {
            None
        }
    };

    // Add user message and attachments as a single user message
    // (Bedrock requires a text block alongside document blocks in the same message)
    {
        let has_message = args
            .user_message
            .as_ref()
            .map(|m| !m.is_empty())
            .unwrap_or(false);
        let has_attachments = args
            .user_attachments
            .as_ref()
            .map(|a| !a.is_empty())
            .unwrap_or(false);

        if has_message && has_attachments {
            let mut parts = vec![ContentPart::Text { text: args.user_message.clone().unwrap() }];
            for attachment in args.user_attachments.as_ref().unwrap() {
                if !attachment.s3.is_empty() {
                    parts.push(ContentPart::S3Object { s3_object: attachment.clone() });
                }
            }
            messages.push(OpenAIMessage {
                role: "user".to_string(),
                content: Some(OpenAIContent::Parts(parts)),
                ..Default::default()
            });
        } else if has_message {
            messages.push(OpenAIMessage {
                role: "user".to_string(),
                content: Some(OpenAIContent::Text(args.user_message.clone().unwrap())),
                ..Default::default()
            });
        } else if has_attachments {
            let mut parts = vec![];
            for attachment in args.user_attachments.as_ref().unwrap() {
                if !attachment.s3.is_empty() {
                    parts.push(ContentPart::S3Object { s3_object: attachment.clone() });
                }
            }
            messages.push(OpenAIMessage {
                role: "user".to_string(),
                content: Some(OpenAIContent::Parts(parts)),
                ..Default::default()
            });
        }
    }

    let mut actions = vec![];
    let mut content = None;
    let mut final_usage: Option<TokenUsage> = None;

    // Check if this provider supports tools with the current output type
    let supports_tools = query_builder.supports_tools_with_output_type(output_type);

    let mut tool_defs: Option<Vec<ToolDef>> = if tools.is_empty() || !supports_tools {
        None
    } else {
        Some(tools.iter().map(|t| t.def.clone()).collect())
    };

    // Handle structured output schema
    let has_output_properties = args
        .output_schema
        .as_ref()
        .and_then(|schema| schema.properties.as_ref())
        .map(|props| !props.is_empty())
        .unwrap_or(false);

    let should_use_structured_output_tool =
        should_use_structured_output_tool(&args.provider.kind, &args.provider.model);
    let mut used_structured_output_tool = false;
    let mut structured_output_tool_name: Option<String> = None;

    // For text output with schema, handle structured output
    if has_output_properties && is_text_output {
        let schema = args.output_schema.as_ref().unwrap();
        if should_use_structured_output_tool {
            // Anthropic uses a tool for structured output
            let unique_tool_name = find_unique_tool_name("structured_output", tool_defs.as_deref());
            structured_output_tool_name = Some(unique_tool_name.clone());

            let output_tool = ToolDef {
                r#type: "function".to_string(),
                function: ToolDefFunction {
                    name: unique_tool_name,
                    description: Some(
                        "This tool MUST be used last to return a structured JSON object as the final output."
                            .to_string(),
                    ),
                    parameters: to_raw_value(&schema),
                },
            };
            if let Some(ref mut existing_tools) = tool_defs {
                existing_tools.push(output_tool);
            } else {
                tool_defs = Some(vec![output_tool]);
            }
        }
        // For non-Anthropic providers, response_format is handled by the query builder
    }

    let user_wants_streaming = args.streaming.unwrap_or(false);
    *has_stream = user_wants_streaming && is_text_output;

    let mut final_events_str = String::new();

    // Always create a StreamEventProcessor for text output (use silent mode if user doesn't want streaming)
    let stream_event_processor = if is_text_output {
        if user_wants_streaming {
            Some(StreamEventProcessor::new(conn, job))
        } else {
            Some(StreamEventProcessor::new_silent())
        }
    } else {
        None
    };

    let chat_enabled = flow_context
        .flow_status
        .as_ref()
        .and_then(|fs| fs.chat_input_enabled)
        .unwrap_or(false);
    let persist_output_to_conversation = chat_enabled && !omit_output_from_conversation;

    let step_name = get_step_name_from_flow(summary.as_deref(), effective_flow_step_id);

    let max_iterations = args
        .max_iterations
        .map(|m| m.clamp(1, HARD_MAX_AGENT_ITERATIONS))
        .unwrap_or(DEFAULT_MAX_AGENT_ITERATIONS);

    // Main agent loop
    for i in 0..max_iterations {
        // Check if parent was canceled — stop iterating but let current tool calls finish
        if *cancel_rx.borrow() {
            return Err(Error::ExecutionErr("Job cancelled".to_string()));
        }

        if used_structured_output_tool {
            break;
        }

        // Handle AWS Bedrock provider specially using the official SDK
        let parsed = if credentials.provider == AIProvider::AWSBedrock {
            #[cfg(feature = "bedrock")]
            {
                let region = credentials
                    .region
                    .as_deref()
                    .unwrap_or(windmill_ai::ai_providers::USE_ENV_REGION);
                // Use Bedrock SDK via dedicated query builder
                windmill_ai::providers::bedrock::BedrockQueryBuilder::default()
                    .execute_request(
                        &messages,
                        tool_defs.as_deref(),
                        args.provider.get_model(),
                        args.temperature,
                        args.provider.get_reasoning_effort(),
                        args.max_completion_tokens,
                        api_key,
                        region,
                        stream_event_processor.as_ref().map(|p| p.boxed_sink()),
                        client,
                        &job.workspace_id,
                        structured_output_tool_name.as_deref(),
                        credentials.aws_access_key_id.as_deref(),
                        credentials.aws_secret_access_key.as_deref(),
                        credentials.aws_session_token.as_deref(),
                    )
                    .await?
            }
            #[cfg(not(feature = "bedrock"))]
            {
                return Err(Error::internal_err(
                    "AWS Bedrock support is not enabled. Build with 'bedrock' feature.".to_string(),
                ));
            }
        } else {
            // For all other providers, use the HTTP client approach
            let build_args = BuildRequestArgs {
                messages: &messages,
                tools: tool_defs.as_deref(),
                model: args.provider.get_model(),
                temperature: args.temperature,
                reasoning_effort: args.provider.get_reasoning_effort(),
                max_tokens: args.max_completion_tokens,
                output_schema: args.output_schema.as_ref(),
                output_type,
                system_prompt: args.system_prompt.as_deref(),
                user_message: args.user_message.as_deref().unwrap_or(""),
                attachments: args.user_attachments.as_deref(),
                has_websearch,
            };

            let request_body = query_builder
                .build_request(&build_args, client, &job.workspace_id)
                .await?;

            let endpoint =
                query_builder.get_endpoint(base_url, args.provider.get_model(), output_type);
            let auth_headers = query_builder.get_auth_headers(api_key, base_url, output_type);
            // A worker cannot run the client credentials exchange, so an OAuth resource
            // has no token here: the request would carry an empty credential and come
            // back 401.
            if needs_unavailable_oauth_exchange(
                &credentials,
                args.provider.resource.token_url.as_deref(),
                &auth_headers,
            ) {
                return Err(Error::ExecutionErr(format!(
                    "The {:?} resource authenticates with OAuth, which AI agent steps do not \
                     support. Set an API key on the resource, or carry the provider's credential \
                     header in its `headers`.",
                    credentials.provider
                )));
            }
            let auth_headers = retain_effective_credentials(&credentials, auth_headers);

            let timeout = resolve_job_timeout(conn, &job.workspace_id, job.id, job.timeout)
                .await
                .0;

            let trailing_headers = common_outbound_headers(&credentials).collect::<Vec<_>>();

            // `endpoint` derives from the user-controlled provider base_url, so pin
            // DNS to the SSRF-validated address: the connect must not rebind to an
            // internal IP between the check and the request (TOCTOU).
            let pinned_ai_client = pinned_ai_client_for(base_url).await?;

            // Helper to build HTTP request with headers
            let build_http_request = |body: String| {
                let mut req = pinned_ai_client
                    .post(&endpoint)
                    .timeout(timeout)
                    .header("Content-Type", "application/json");

                for (header_name, header_value) in &auth_headers {
                    req = req.header(*header_name, header_value.clone());
                }

                for (header_name, header_value) in &trailing_headers {
                    req = req.header(header_name.as_str(), header_value.as_str());
                }

                req.body(body)
            };

            let resp = build_http_request(request_body.clone())
                .send()
                .await
                .map_err(|e| Error::internal_err(format!("Failed to call API: {}", e)))?;

            // Check if request failed and we should retry without stream_options
            let resp = match resp.error_for_status_ref() {
                Ok(_) => resp,
                Err(e) => {
                    let status = resp.status();
                    let text = resp
                        .text()
                        .await
                        .unwrap_or_else(|_| "<failed to read body>".to_string());

                    // Retry without stream_options if provider supports it and error suggests incompatibility
                    // Common error patterns: 400 Bad Request with mentions of stream_options or include_usage
                    let should_retry = query_builder.supports_retry_without_usage()
                        && status.as_u16() == 400
                        && (text.contains("stream_options")
                            || text.contains("include_usage")
                            || text.contains("Additional properties are not allowed"));

                    if should_retry {
                        tracing::info!(
                            "Retrying request without stream_options due to provider incompatibility"
                        );

                        let retry_body = query_builder
                            .build_request_without_usage(&build_args, client, &job.workspace_id)
                            .await?;

                        let retry_resp =
                            build_http_request(retry_body).send().await.map_err(|e| {
                                Error::internal_err(format!("Failed to call API on retry: {}", e))
                            })?;

                        match retry_resp.error_for_status_ref() {
                            Ok(_) => retry_resp,
                            Err(retry_e) => {
                                let retry_text = retry_resp
                                    .text()
                                    .await
                                    .unwrap_or_else(|_| "<failed to read body>".to_string());
                                return Err(Error::internal_err(format!(
                                    "API error on retry: {} - {}",
                                    retry_e, retry_text
                                )));
                            }
                        }
                    } else {
                        return Err(Error::internal_err(format!("API error: {} - {}", e, text)));
                    }
                }
            };

            if let Some(ref stream_event_processor) = stream_event_processor {
                query_builder
                    .parse_streaming_response(resp, stream_event_processor.boxed_sink())
                    .await?
            } else {
                query_builder.parse_image_response(resp).await?
            }
        };

        match parsed {
            ParsedResponse::Text {
                content: response_content,
                tool_calls,
                events_str,
                annotations,
                used_websearch,
                usage,
            } => {
                // Accumulate usage from this iteration
                if let Some(u) = usage {
                    match &mut final_usage {
                        Some(existing) => existing.accumulate(&u),
                        None => final_usage = Some(u),
                    }
                }
                if let Some(events_str) = events_str {
                    final_events_str.push_str(&events_str);
                }

                // Add websearch tool message if websearch was used
                if used_websearch {
                    actions.push(AgentAction::WebSearch {});
                    if let Some(parent_job) = parent_job {
                        update_flow_status_module_with_actions(db, parent_job, &actions).await?;
                        update_flow_status_module_with_actions_success(db, parent_job, true)
                            .await?;
                    }
                    messages.push(OpenAIMessage {
                        role: "tool".to_string(),
                        content: Some(OpenAIContent::Text(
                            "Used websearch tool successfully".to_string(),
                        )),
                        agent_action: Some(AgentAction::WebSearch {}),
                        ..Default::default()
                    });
                    if persist_output_to_conversation {
                        if let Some(memory_id) = memory_id {
                            let agent_job_id = job.id;
                            let db_clone = db.clone();
                            let message_content = "Used websearch tool successfully".to_string();
                            let step_name = step_name.clone();
                            tokio::spawn(async move {
                                if let Err(e) = add_message_to_conversation(
                                    &db_clone,
                                    &memory_id,
                                    Some(agent_job_id),
                                    &message_content,
                                    MessageType::Tool,
                                    &step_name,
                                    true,
                                )
                                .await
                                {
                                    tracing::warn!(
                                        "Failed to add websearch tool message to conversation {}: {}",
                                        memory_id,
                                        e
                                    );
                                }
                            });
                        }
                    }
                }

                if let Some(ref response_content) = response_content {
                    actions.push(AgentAction::Message {});
                    messages.push(OpenAIMessage {
                        role: "assistant".to_string(),
                        content: Some(OpenAIContent::Text(response_content.clone())),
                        agent_action: Some(AgentAction::Message {}),
                        annotations: if annotations.is_empty() {
                            None
                        } else {
                            Some(annotations.clone())
                        },
                        ..Default::default()
                    });

                    if let Some(parent_job) = parent_job {
                        update_flow_status_module_with_actions(db, parent_job, &actions).await?;
                        update_flow_status_module_with_actions_success(db, parent_job, true)
                            .await?;
                    }

                    content = Some(OpenAIContent::Text(response_content.clone()));

                    // Add assistant message to conversation if chat_input_enabled
                    if persist_output_to_conversation && !response_content.is_empty() {
                        if let Some(memory_id) = memory_id {
                            let agent_job_id = job.id;
                            let db_clone = db.clone();
                            let message_content = response_content.clone();
                            let step_name = step_name.clone();

                            // Spawn task because we do not need to wait for the result
                            tokio::spawn(async move {
                                if let Err(e) = add_message_to_conversation(
                                    &db_clone,
                                    &memory_id,
                                    Some(agent_job_id),
                                    &message_content,
                                    MessageType::Assistant,
                                    &step_name,
                                    true,
                                )
                                .await
                                {
                                    tracing::warn!(
                                        "Failed to add assistant message to conversation {}: {}",
                                        memory_id,
                                        e
                                    );
                                }
                            });
                        }
                    }
                }

                if tool_calls.is_empty() {
                    break;
                } else if i == max_iterations - 1 {
                    #[derive(serde::Serialize)]
                    struct MaxIterError<'a> {
                        message: String,
                        name: &'static str,
                        #[serde(skip_serializing_if = "Option::is_none")]
                        step_id: Option<&'a str>,
                        result: MaxIterPartialResult<'a>,
                    }
                    #[derive(serde::Serialize)]
                    struct MaxIterPartialResult<'a> {
                        messages: &'a [OpenAIMessage],
                    }
                    return Err(Error::ExecutionRawError(
                        serde_json::value::to_raw_value(&MaxIterError {
                            message: format!(
                                "AI agent reached max iterations ({}), you can either increase max_iterations or enable the \"continue on error\" option from the advanced options of the step.",
                                max_iterations
                            ),
                            name: "ExecutionErr",
                            step_id: effective_flow_step_id,
                            result: MaxIterPartialResult { messages: &messages },
                        })?,
                    ));
                }

                messages.push(OpenAIMessage {
                    role: "assistant".to_string(),
                    tool_calls: Some(tool_calls.clone()),
                    ..Default::default()
                });

                // Handle tool calls using extracted tools module
                let tool_execution_ctx = ToolExecutionContext {
                    db,
                    conn,
                    job,
                    parent_job,
                    summary: &summary,
                    flow_step_id_override,
                    client,
                    worker_dir,
                    base_internal_url,
                    worker_name,
                    hostname,
                    occupancy_metrics,
                    killpill_rx,
                    stream_event_processor: stream_event_processor.as_ref(),
                    flow_context: &mut flow_context,
                    omit_output_from_conversation,
                    previous_result: &previous_result,
                    id_context: &id_context,
                    tool_abort_handles: tool_abort_handles.clone(),
                };

                let (tool_messages, tool_content, tool_used_structured_output) =
                    execute_tool_calls(
                        tool_execution_ctx,
                        &tool_calls,
                        &tools,
                        mcp_clients,
                        &mut actions,
                        &mut final_events_str,
                        &structured_output_tool_name,
                    )
                    .await?;

                messages.extend(tool_messages);
                if let Some(tc) = tool_content {
                    content = Some(tc);
                }
                used_structured_output_tool = tool_used_structured_output;

                // Check cancellation after tool calls complete to avoid a wasted LLM call
                if *cancel_rx.borrow() {
                    return Err(Error::ExecutionErr("Job cancelled".to_string()));
                }
            }
            ParsedResponse::Image { base64_data } => {
                // For image output, upload to S3 and track in conversation
                let s3_object =
                    upload_image_to_s3(&base64_data, &job.workspace_id, &job.id, client).await?;

                let content = to_raw_value(&s3_object);

                // Add assistant message to conversation if chat_input_enabled
                if persist_output_to_conversation {
                    if let Some(memory_id) = memory_id {
                        let agent_job_id = job.id;
                        let db_clone = db.clone();

                        // Create extended version with type discriminator for conversation storage
                        // This avoids conflicts with outputs that are of the same format as S3 objects
                        let s3_with_type = S3ObjectWithType {
                            s3_object: s3_object.clone(),
                            r#type: "windmill_s3_object".to_string(),
                        };

                        let message_content = serde_json::to_string(&s3_with_type)
                            .unwrap_or_else(|_| content.get().to_string());

                        // Spawn task because we do not need to wait for the result
                        tokio::spawn(async move {
                            if let Err(e) = add_message_to_conversation(
                                &db_clone,
                                &memory_id,
                                Some(agent_job_id),
                                &message_content,
                                MessageType::Assistant,
                                &step_name,
                                true,
                            )
                            .await
                            {
                                tracing::warn!(
                                    "Failed to add assistant message to conversation {}: {}",
                                    memory_id,
                                    e
                                );
                            }
                        });
                    }
                }

                // Return early since image generation is complete
                return Ok(content);
            }
        }
    }

    // Return the final result
    let final_messages: Vec<Message> = messages
        .iter()
        .map(|m| Message { message: m, agent_action: m.agent_action.as_ref() })
        .collect();

    // Parse content as JSON for structured output, fallback to string if it fails
    let output_value = match content {
        Some(content_str) => match has_output_properties {
            true => match content_str {
                OpenAIContent::Text(text) => {
                    serde_json::from_str::<Box<RawValue>>(&text).map_err(|_e| {
                        Error::internal_err(format!("Failed to parse structured output: {}", text))
                    })
                }
                OpenAIContent::Parts(_parts) => Err(Error::internal_err(
                    "Failed to parse structured output".to_string(),
                )),
            },
            false => Ok(match content_str {
                OpenAIContent::Text(text) => to_raw_value(&text),
                OpenAIContent::Parts(parts) => to_raw_value(&parts),
            }),
        }?,
        None => to_raw_value(&""),
    };

    // Wait for stream event processor to finish persisting events (if any)
    if let Some(handle) = {
        if let Some(stream_event_processor) = stream_event_processor {
            stream_event_processor.to_handle()
        } else {
            None
        }
    } {
        if let Err(e) = handle.await {
            return Err(Error::internal_err(format!(
                "Error waiting for stream event processor: {}",
                e
            )));
        }
    }

    // Persist complete conversation to memory at the end (only if in auto mode with context length)
    // Skip memory persistence if using manual messages (bypass memory entirely)
    // final_messages contains the complete history (old messages + new ones)
    if matches!(output_type, OutputType::Text) && !use_manual_messages {
        if let Some(Memory::Auto { context_length, .. }) = &args.memory {
            if let Some(step_id) = effective_flow_step_id {
                // Extract OpenAIMessages from final_messages
                let all_messages: Vec<OpenAIMessage> =
                    final_messages.iter().map(|m| m.message.clone()).collect();

                if !all_messages.is_empty() {
                    let messages_to_persist = prepare_auto_memory_messages_for_persistence(
                        &all_messages,
                        *context_length,
                    );

                    if let Some(memory_id) = memory_id {
                        if let Err(e) = write_to_memory(
                            db,
                            &job.workspace_id,
                            memory_id,
                            step_id,
                            &messages_to_persist,
                        )
                        .await
                        {
                            tracing::error!(
                                "Failed to persist {} messages to memory for step {}: {}",
                                messages_to_persist.len(),
                                step_id,
                                e
                            );
                        }
                    }
                }
            }
        }
    }

    Ok(to_raw_value(&AIAgentResult {
        output: output_value,
        messages: final_messages,
        wm_stream: if !final_events_str.is_empty() {
            Some(final_events_str)
        } else {
            None
        },
        usage: if final_usage.as_ref().map(|u| u.is_empty()).unwrap_or(true) {
            None
        } else {
            final_usage
        },
    }))
}

#[cfg(test)]
mod tests {
    use super::*;

    fn text_message(role: &str, content: &str) -> OpenAIMessage {
        OpenAIMessage {
            role: role.to_string(),
            content: Some(OpenAIContent::Text(content.to_string())),
            ..Default::default()
        }
    }

    #[test]
    fn overlay_tool_inputs_binds_matching_flowmodule_tool_only() {
        fn js(expr: &str) -> InputTransform {
            InputTransform::Javascript { expr: expr.to_string() }
        }
        fn script_tool(id: &str, key: &str, expr: &str) -> AgentTool {
            let mut its = HashMap::new();
            its.insert(key.to_string(), js(expr));
            AgentTool {
                id: id.to_string(),
                summary: None,
                description: None,
                value: ToolValue::FlowModule(FlowModuleValue::Script {
                    input_transforms: its,
                    path: "u/test/tool".to_string(),
                    hash: None,
                    tag_override: None,
                    is_trigger: None,
                    pass_flow_input_directly: None,
                }),
            }
        }
        fn script_its(tool: &AgentTool) -> &HashMap<String, InputTransform> {
            let ToolValue::FlowModule(FlowModuleValue::Script { input_transforms, .. }) =
                &tool.value
            else {
                panic!("expected script tool")
            };
            input_transforms
        }

        // "a" gets rebound, "b" is left alone, the MCP tool is skipped even though it has an override.
        let mut tools = vec![
            script_tool("a", "x", "authoring_flow_expr"),
            script_tool("b", "y", "keep_me"),
            AgentTool {
                id: "m".to_string(),
                summary: None,
                description: None,
                value: ToolValue::Mcp(windmill_common::flows::McpToolValue {
                    resource_path: "u/test/mcp".to_string(),
                    include_tools: vec![],
                    exclude_tools: vec![],
                }),
            },
        ];

        let mut tool_inputs: HashMap<String, HashMap<String, InputTransform>> = HashMap::new();
        tool_inputs.insert(
            "a".to_string(),
            HashMap::from([
                ("x".to_string(), js("flow_input.tenant")),
                ("z".to_string(), js("results.step1")),
            ]),
        );
        tool_inputs.insert(
            "m".to_string(),
            HashMap::from([("q".to_string(), js("ignored"))]),
        );

        overlay_tool_inputs(&mut tools, &tool_inputs);

        // "a": existing key replaced, new key added.
        let a = script_its(&tools[0]);
        assert!(
            matches!(a.get("x"), Some(InputTransform::Javascript { expr }) if expr == "flow_input.tenant")
        );
        assert!(
            matches!(a.get("z"), Some(InputTransform::Javascript { expr }) if expr == "results.step1")
        );
        // "b": no override for it, untouched.
        let b = script_its(&tools[1]);
        assert!(
            matches!(b.get("y"), Some(InputTransform::Javascript { expr }) if expr == "keep_me")
        );
        // MCP tool: not a FlowModule, left as-is.
        assert!(matches!(&tools[2].value, ToolValue::Mcp(_)));
    }

    #[test]
    fn tool_description_prefers_explicit_over_derived_and_name() {
        assert_eq!(
            resolve_tool_description(
                Some("  Use to look up a user by id  ".to_string()),
                Some("derived from script".to_string()),
                "get_user"
            ),
            "Use to look up a user by id"
        );
    }

    #[test]
    fn tool_description_falls_back_to_derived_when_no_explicit() {
        assert_eq!(
            resolve_tool_description(None, Some("Sync resources".to_string()), "sync_tool"),
            "Sync resources"
        );
        // A blank explicit description must not shadow a usable derived one.
        assert_eq!(
            resolve_tool_description(
                Some("   ".to_string()),
                Some("Sync resources".to_string()),
                "sync_tool"
            ),
            "Sync resources"
        );
    }

    #[test]
    fn tool_description_falls_back_to_name_when_nothing_usable() {
        assert_eq!(resolve_tool_description(None, None, "my_tool"), "my_tool");
        assert_eq!(
            resolve_tool_description(Some("  ".to_string()), Some("".to_string()), "my_tool"),
            "my_tool"
        );
    }

    #[test]
    fn auto_memory_request_preserves_messages_within_context_window() {
        let loaded_messages = vec![
            text_message("system", "instructions-a"),
            text_message("user", "first-user"),
            text_message("assistant", "first-assistant"),
            text_message("system", "instructions-b"),
            text_message("user", "second-user"),
            text_message("assistant", "second-assistant"),
        ];

        let prepared = prepare_auto_memory_messages_for_request(&loaded_messages, 3);
        let roles: Vec<&str> = prepared
            .iter()
            .map(|message| message.role.as_str())
            .collect();
        let contents: Vec<&str> = prepared
            .iter()
            .map(|message| match message.content.as_ref() {
                Some(OpenAIContent::Text(text)) => text.as_str(),
                _ => "",
            })
            .collect();

        assert_eq!(roles, vec!["system", "user", "assistant"]);
        assert_eq!(
            contents,
            vec!["instructions-b", "second-user", "second-assistant"]
        );
    }

    #[test]
    fn auto_memory_request_drops_leading_tool_messages() {
        let loaded_messages = vec![
            text_message("tool", "stale-tool-result"),
            text_message("user", "hello"),
            text_message("assistant", "hi"),
        ];

        let prepared = prepare_auto_memory_messages_for_request(&loaded_messages, 10);
        let roles: Vec<&str> = prepared
            .iter()
            .map(|message| message.role.as_str())
            .collect();

        assert_eq!(roles, vec!["user", "assistant"]);
    }

    #[test]
    fn auto_memory_persistence_excludes_system_messages() {
        let all_messages = vec![
            text_message("system", "instructions"),
            text_message("user", "hello"),
            text_message("assistant", "hi"),
            text_message("system", "duplicate-instructions"),
            text_message("user", "follow-up"),
        ];

        let persisted = prepare_auto_memory_messages_for_persistence(&all_messages, 10);
        let roles: Vec<&str> = persisted
            .iter()
            .map(|message| message.role.as_str())
            .collect();

        assert_eq!(roles, vec!["user", "assistant", "user"]);
    }
}

/// Handle credentials check mode - check credentials without making API calls
async fn handle_credentials_check(provider: &ProviderWithResource) -> Result<Box<RawValue>, Error> {
    let result = match &provider.kind {
        #[cfg(feature = "bedrock")]
        AIProvider::AWSBedrock => {
            let check = check_env_credentials().await;
            serde_json::json!({
                "credentials_check": true,
                "provider": "aws_bedrock",
                "credentials": {
                    "available": check.available,
                    "access_key_id_prefix": check.access_key_id_prefix,
                    "region": check.region,
                    "error": check.error
                }
            })
        }
        #[cfg(not(feature = "bedrock"))]
        AIProvider::AWSBedrock => {
            serde_json::json!({
                "credentials_check": true,
                "provider": "aws_bedrock",
                "error": "AWS Bedrock support is not enabled. Build with 'bedrock' feature."
            })
        }
        other => {
            serde_json::json!({
                "credentials_check": true,
                "provider": format!("{:?}", other),
                "message": "Credentials check not implemented for this provider"
            })
        }
    };

    serde_json::value::to_raw_value(&result).map_err(|e| Error::internal_err(e.to_string()))
}

/// Hard-timeout fallback: force-cancel any descendant jobs still in v2_job_queue
/// so they don't stay as zombies.
async fn cleanup_orphaned_tool_jobs(
    db: &DB,
    parent_job_id: &Uuid,
    w_id: &str,
    canceled_by: Option<CanceledBy>,
) {
    let username = canceled_by
        .as_ref()
        .and_then(|cb| cb.username.clone())
        .unwrap_or_else(|| "unknown".to_string());
    let reason = canceled_by
        .as_ref()
        .and_then(|cb| cb.reason.clone())
        .unwrap_or_else(|| {
            format!(
                "parent AI agent {} was cancelled and tool call did not complete in time",
                parent_job_id
            )
        });

    // Find direct child jobs still in v2_job_queue (agent tool jobs are always direct children)
    let orphaned_ids: Vec<Uuid> = match sqlx::query_scalar!(
        r#"SELECT j.id FROM v2_job j
            JOIN v2_job_queue q ON q.id = j.id
            WHERE j.parent_job = $1 AND j.workspace_id = $2"#,
        parent_job_id,
        w_id,
    )
    .fetch_all(db)
    .await
    {
        Ok(ids) => ids,
        Err(e) => {
            tracing::error!(
                "Failed to find orphaned tool jobs for {}: {}",
                parent_job_id,
                e
            );
            return;
        }
    };

    if orphaned_ids.is_empty() {
        return;
    }

    tracing::warn!(
        "Cleaning up {} orphaned tool jobs for cancelled AI agent {}",
        orphaned_ids.len(),
        parent_job_id,
    );

    for job_id in &orphaned_ids {
        let queued_job = match windmill_queue::get_queued_job_v2(db, job_id).await {
            Ok(Some(j)) => j,
            Ok(None) => continue,
            Err(e) => {
                tracing::error!("Failed to fetch orphaned tool job {}: {}", job_id, e);
                continue;
            }
        };

        let tx = match db.begin().await {
            Ok(tx) => tx,
            Err(e) => {
                tracing::error!(
                    "Failed to begin transaction for orphaned job {}: {}",
                    job_id,
                    e
                );
                continue;
            }
        };

        match cancel_single_job(
            &username,
            Some(reason.clone()),
            queued_job,
            w_id,
            tx,
            db,
            true,
        )
        .await
        {
            Ok((tx, _)) => {
                if let Err(e) = tx.commit().await {
                    tracing::error!(
                        "Failed to commit cancel for orphaned tool job {}: {}",
                        job_id,
                        e
                    );
                }
            }
            Err(e) => {
                // warn not error: job may have completed between fetch and cancel (expected race)
                tracing::warn!("Failed to force-cancel orphaned tool job {}: {}", job_id, e);
            }
        }
    }
}