Rebasing, going about this another way

app/src/ai/agent/api.rs

@@ -124,6 +124,9 @@ pub struct RequestParams {
     pub research_agent_enabled: bool,
     pub orchestration_enabled: bool,
     pub supported_tools_override: Option<Vec<warp_multi_agent_api::ToolType>>,
+    /// The root task ID for the conversation — needed for direct Bedrock streaming
+    /// since optimistic tasks don't appear in the proto task_context.
+    pub root_task_id: Option<String>,
     /// The conversation ID of the parent agent that spawned this child agent, if any.
     pub parent_agent_id: Option<String>,
     /// The display name for this agent (e.g. "Agent 1"), assigned by the orchestrator.
@@ -235,7 +238,7 @@ impl RequestParams {
         let user_workspaces = UserWorkspaces::as_ref(app);
         let api_keys = ApiKeyManager::as_ref(app).api_keys_for_request(
             user_workspaces.is_byo_api_key_enabled(),
-            user_workspaces.is_aws_bedrock_credentials_enabled(app),
+            user_workspaces.is_bedrock_enabled(app),
         );
         let allow_use_of_warp_credits_with_byok =
             *AISettings::as_ref(app).can_use_warp_credits_with_byok;
@@ -307,6 +310,11 @@ impl RequestParams {
             research_agent_enabled,
             orchestration_enabled,
             supported_tools_override: request_input.supported_tools_override.clone(),
+            root_task_id: request_input
+                .input_messages
+                .keys()
+                .next()
+                .map(|id| id.to_string()),
             parent_agent_id: None,
             agent_name: None,
         }

app/src/ai/agent/api/impl.rs

@@ -5,12 +5,14 @@ use futures_util::StreamExt;
 use warp_core::features::FeatureFlag;
 use warp_multi_agent_api as api;
 
+use crate::ai::bedrock::client::{BedrockClient, BedrockClientConfig};
 use crate::server::server_api::ServerApi;
 
 use super::{convert_to::convert_input, ConvertToAPITypeError, RequestParams, ResponseStream};
 
 pub async fn generate_multi_agent_output(
     server_api: Arc<ServerApi>,
+    bedrock_config: Option<BedrockClientConfig>,
     mut params: RequestParams,
     cancellation_rx: futures::channel::oneshot::Receiver<()>,
 ) -> Result<ResponseStream, ConvertToAPITypeError> {
@@ -129,6 +131,101 @@ pub async fn generate_multi_agent_output(
         mcp_context: params.mcp_context.map(Into::into),
     };
 
+    if let Some(config) = bedrock_config {
+        let model_id_for_fallback_check = request
+            .settings
+            .as_ref()
+            .and_then(|s| s.model_config.as_ref())
+            .map(|mc| mc.base.clone())
+            .unwrap_or_default();
+        let is_arn = model_id_for_fallback_check.starts_with("arn:");
+        let fallback_to_warp = config.fallback_to_warp && !is_arn;
+        if is_arn && config.fallback_to_warp {
+            log::info!("[bedrock] Fallback disabled for ARN-based model (not available on Warp server)");
+        }
+        match BedrockClient::from_config(config).await {
+            Ok(bedrock) => {
+                let task_id = params.root_task_id.clone().unwrap_or_else(|| {
+                    request
+                        .task_context
+                        .as_ref()
+                        .and_then(|tc| tc.tasks.first())
+                        .map(|t| t.id.clone())
+                        .unwrap_or_else(|| uuid::Uuid::new_v4().to_string())
+                });
+
+                log::info!("[bedrock] Starting stream with task_id={task_id}");
+
+                let model_id = request
+                    .settings
+                    .as_ref()
+                    .and_then(|s| s.model_config.as_ref())
+                    .map(|mc| mc.base.clone())
+                    .unwrap_or_default();
+
+                log::info!("[bedrock] Model: {model_id}");
+
+                let messages =
+                    crate::ai::bedrock::convert_request::extract_messages_from_request(&request);
+                let system_prompt =
+                    crate::ai::bedrock::convert_request::extract_system_prompt(&request);
+                let tools = crate::ai::bedrock::convert_request::extract_tools(&request);
+
+                log::info!(
+                    "[bedrock] Sending {} messages, system_prompt={}, tools={}",
+                    messages.len(),
+                    system_prompt.is_some(),
+                    tools.len()
+                );
+
+                match bedrock
+                    .converse_stream(
+                        &model_id,
+                        &task_id,
+                        messages,
+                        system_prompt,
+                        tools,
+                        8192,
+                        None,
+                        true,
+                    )
+                    .await
+                {
+                    Ok(stream) => {
+                        let output_stream = stream.take_until(cancellation_rx);
+                        return Ok(Box::pin(output_stream));
+                    }
+                    Err(e) => {
+                        if fallback_to_warp {
+                            log::warn!("Bedrock stream failed, falling back to server: {e}");
+                        } else {
+                            let err = Arc::new(crate::server::server_api::AIApiError::Stream {
+                                stream_type: "bedrock_converse",
+                                source: anyhow::anyhow!("{e}"),
+                            });
+                            let (tx, rx) = async_channel::unbounded();
+                            let _ = tx.send(Err(err)).await;
+                            return Ok(Box::pin(rx));
+                        }
+                    }
+                }
+            }
+            Err(e) => {
+                if fallback_to_warp {
+                    log::warn!("Bedrock client creation failed, falling back to server: {e}");
+                } else {
+                    let err = Arc::new(crate::server::server_api::AIApiError::Stream {
+                        stream_type: "bedrock_converse",
+                        source: anyhow::anyhow!("{e}"),
+                    });
+                    let (tx, rx) = async_channel::unbounded();
+                    let _ = tx.send(Err(err)).await;
+                    return Ok(Box::pin(rx));
+                }
+            }
+        }
+    }
+
     let response_stream = server_api.generate_multi_agent_output(&request).await;
     match response_stream {
         Ok(stream) => {

app/src/ai/agent/api/impl_tests.rs

@@ -37,6 +37,7 @@ fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool
         research_agent_enabled: false,
         orchestration_enabled: false,
         supported_tools_override: None,
+        root_task_id: None,
         parent_agent_id: None,
         agent_name: None,
     }