a2a-agents 0.6.0

use a2a_rs::Artifact;
use async_trait::async_trait;
use chrono::Utc;
use serde::Deserialize;
use serde_json::{Map, Value, json};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tracing::{debug, error, info, instrument, warn};
use uuid::Uuid;

use a2a_rs::application::{HasPushNotifier, HasStreaming, HasTaskLifecycle, TaskStatusBroadcast};
use a2a_rs::domain::{A2AError, ContextId, Message, Part, Role, Task, TaskId, TaskState, part};
use a2a_rs::port::message_handler::AsyncMessageHandler;

use super::types::*;
use a2a_agents_common::llm::{ChatMessage, LlmProvider, LlmRequest, ToolCallAccumulator};

// NOTE: Task storage is handled by ConnectRpcAdapter + SQLx/InMemory storage
// This handler is stateless and only processes messages

/// Build the typed metadata that marks an artifact as a (possibly partial)
/// tool-call delta, so a UI can distinguish it from a text chunk and reassemble
/// the call's arguments by `call_id`.
fn tool_call_metadata(call_id: &str, tool_name: &str, partial: bool) -> Map<String, Value> {
    let mut m = Map::new();
    m.insert("type".to_string(), json!("tool-call"));
    m.insert("call_id".to_string(), json!(call_id));
    m.insert("tool_name".to_string(), json!(tool_name));
    m.insert("partial".to_string(), json!(partial));
    m
}

/// Metrics for tracking handler performance
#[derive(Debug, Default, Clone)]
pub struct HandlerMetrics {
    pub total_requests: u64,
    pub successful_requests: u64,
    pub validation_errors: u64,
    pub processing_errors: u64,
    pub forms_generated: u64,
    pub approvals_processed: u64,
    pub auto_approvals: u64,
}

impl HandlerMetrics {
    fn increment_requests(&mut self) {
        self.total_requests += 1;
    }

    fn increment_success(&mut self) {
        self.successful_requests += 1;
    }

    fn increment_validation_errors(&mut self) {
        self.validation_errors += 1;
    }

    fn increment_processing_errors(&mut self) {
        self.processing_errors += 1;
    }

    fn increment_forms(&mut self) {
        self.forms_generated += 1;
    }

    fn increment_approvals(&mut self) {
        self.approvals_processed += 1;
    }

    fn increment_auto_approvals(&mut self) {
        self.auto_approvals += 1;
    }

    fn log_metrics(&self) {
        info!(
            total_requests = self.total_requests,
            successful_requests = self.successful_requests,
            validation_errors = self.validation_errors,
            processing_errors = self.processing_errors,
            forms_generated = self.forms_generated,
            approvals_processed = self.approvals_processed,
            auto_approvals = self.auto_approvals,
            success_rate = if self.total_requests > 0 {
                (self.successful_requests as f64 / self.total_requests as f64) * 100.0
            } else {
                0.0
            },
            "Handler metrics"
        );
    }
}

/// Reimbursement message handler with proper JSON parsing and validation
/// Reimbursement handler that manages task history through a task manager
#[derive(Clone)]
pub struct ReimbursementHandler {
    /// Per-task lifecycle port (create/get/update_status)
    task_lifecycle: Arc<dyn a2a_rs::port::AsyncTaskLifecycle>,
    /// Streaming port for broadcasting artifact updates
    streaming: Arc<dyn a2a_rs::port::AsyncStreamingHandler>,
    /// Push-notifier port for out-of-band webhook delivery on each transition
    push_notifier: Arc<dyn a2a_rs::port::AsyncPushNotifier>,
    validation_rules: ValidationRules,
    #[allow(dead_code)]
    file_metadata_store: Arc<Mutex<HashMap<String, Map<String, Value>>>>,
    metrics: Arc<Mutex<HandlerMetrics>>,
    llm_provider: Option<Arc<dyn LlmProvider>>,
}

// The handler holds both the lifecycle and streaming ports, so it hosts the
// `TaskStatusBroadcast` mixin (see `.claude/rules/hexagonal_architecture.md`
// §9): every status transition it drives goes through `update_and_broadcast`,
// announcing the change to streaming subscribers. Storage is persistence-only
// and does not self-broadcast on mutation, so without these accessors the
// agent's updates — including the background worker's final state — would
// silently stop reaching subscribers and push notifications.
impl HasTaskLifecycle for ReimbursementHandler {
    fn lifecycle(&self) -> &dyn a2a_rs::port::AsyncTaskLifecycle {
        self.task_lifecycle.as_ref()
    }
}

impl HasStreaming for ReimbursementHandler {
    fn streaming(&self) -> &dyn a2a_rs::port::AsyncStreamingHandler {
        self.streaming.as_ref()
    }
}

impl HasPushNotifier for ReimbursementHandler {
    fn push_notifier(&self) -> &dyn a2a_rs::port::AsyncPushNotifier {
        self.push_notifier.as_ref()
    }
}

#[allow(dead_code)]
impl ReimbursementHandler {
    pub fn new(
        task_lifecycle: impl a2a_rs::port::AsyncTaskLifecycle + 'static,
        streaming: impl a2a_rs::port::AsyncStreamingHandler + 'static,
        push_notifier: impl a2a_rs::port::AsyncPushNotifier + 'static,
    ) -> Self {
        // Initialize an AI client from the environment (OpenRouter → Gemini →
        // OpenAI). `None` disables conversational features.
        let llm_provider = a2a_agents_common::llm::provider_from_env();

        Self::with_llm(task_lifecycle, streaming, push_notifier, llm_provider)
    }

    pub fn with_llm(
        task_lifecycle: impl a2a_rs::port::AsyncTaskLifecycle + 'static,
        streaming: impl a2a_rs::port::AsyncStreamingHandler + 'static,
        push_notifier: impl a2a_rs::port::AsyncPushNotifier + 'static,
        llm_provider: Option<Arc<dyn LlmProvider>>,
    ) -> Self {
        Self {
            task_lifecycle: Arc::new(task_lifecycle),
            streaming: Arc::new(streaming),
            push_notifier: Arc::new(push_notifier),
            validation_rules: ValidationRules::default(),
            file_metadata_store: Arc::new(Mutex::new(HashMap::new())),
            metrics: Arc::new(Mutex::new(HandlerMetrics::default())),
            llm_provider,
        }
    }

    pub fn with_validation_rules(mut self, rules: ValidationRules) -> Self {
        self.validation_rules = rules;
        self
    }

    /// Merge metadata from parts into a combined metadata map
    fn merge_metadata(&self, target: &mut Map<String, Value>, source: &Map<String, Value>) {
        for (key, value) in source {
            // Handle special metadata keys that might influence processing
            match key.as_str() {
                "expense_type" | "category_hint" | "date_format" | "currency" => {
                    target.insert(key.clone(), value.clone());
                }
                "auto_approve" if value.is_boolean() => {
                    target.insert(key.clone(), value.clone());
                }
                "priority" if value.is_string() => {
                    target.insert(key.clone(), value.clone());
                }
                _ => {
                    // Store other metadata for reference
                    if !key.starts_with("_") {
                        // Skip internal metadata
                        target.insert(key.clone(), value.clone());
                    }
                }
            }
        }
    }

    /// Store file metadata for later retrieval
    fn store_file_metadata(&self, file_id: &str, metadata: Map<String, Value>) {
        if let Ok(mut store) = self.file_metadata_store.lock() {
            store.insert(file_id.to_string(), metadata);
        }
    }

    /// Retrieve file metadata
    fn get_file_metadata(&self, file_id: &str) -> Option<Map<String, Value>> {
        self.file_metadata_store
            .lock()
            .ok()
            .and_then(|store| store.get(file_id).cloned())
    }

    /// Get current metrics
    pub fn get_metrics(&self) -> HandlerMetrics {
        self.metrics.lock().map(|m| m.clone()).unwrap_or_default()
    }

    /// Log current metrics
    pub fn log_metrics(&self) {
        if let Ok(metrics) = self.metrics.lock() {
            metrics.log_metrics();
        }
    }

    /// Update metrics based on processing results
    fn update_metrics(&self, response: &ReimbursementResponse, auto_approved: bool) {
        if let Ok(mut metrics) = self.metrics.lock() {
            metrics.increment_requests();

            match response {
                ReimbursementResponse::Form { .. } => {
                    metrics.increment_forms();
                    metrics.increment_success();
                }
                ReimbursementResponse::Result { status, .. } => {
                    metrics.increment_approvals();
                    if auto_approved {
                        metrics.increment_auto_approvals();
                    }
                    match status {
                        super::types::ProcessingStatus::Approved
                        | super::types::ProcessingStatus::Pending => {
                            metrics.increment_success();
                        }
                        _ => {}
                    }
                }
                ReimbursementResponse::Error { code, .. } => {
                    if code == "VALIDATION_ERROR" {
                        metrics.increment_validation_errors();
                    } else {
                        metrics.increment_processing_errors();
                    }
                }
            }
        }
    }

    /// System prompt for the AI reimbursement assistant
    fn get_system_prompt(&self) -> String {
        r#"You are an AI-powered reimbursement assistant. You MUST respond with valid JSON only.

Company Policy:
- Expenses under $100: auto-approved (if all info provided + receipt if amount > $25)
- Expenses $100+: requires manager approval (mark as "pending")
- Categories: travel, meals, supplies, equipment, training, other
- Required info: date, amount, purpose, category
- Receipt rules: ONLY required when amount is GREATER than $25 (>$25)
  * $25 or less: NO receipt needed
  * Over $25 (e.g. $25.01, $26, $30): Receipt REQUIRED

CRITICAL: You MUST respond with ONLY valid JSON in this exact format:

{
  "action": "approved",
  "message": "Your conversational message here",
  "extracted_data": {
    "date": "2025-10-24",
    "amount": 25.00,
    "currency": "USD",
    "purpose": "Office supplies",
    "category": "supplies",
    "has_receipt": true
  },
  "needs_clarification": [],
  "approval_reason": "Expense under $100 with receipt"
}

Action values (choose ONE):
- "approved": All info present + amount<$100 + (receipt uploaded if amount>$25)
- "pending": All info present + amount>=$100
- "question": Missing required info OR (amount>$25 AND no receipt uploaded)
- "error": Cannot process (system error, not user error)

Rules:
1. ALWAYS return valid JSON (no extra text before or after)
2. Use "question" action when info is missing
3. Set extracted_data fields to null if unknown
4. When user uploads file, set has_receipt=true and acknowledge in message
5. Be conversational in the "message" field only
6. IMPORTANT: Check math carefully - $23 is NOT over $25, $26 IS over $25

Example response when asking for info:
{
  "action": "question",
  "message": "What was the amount of the expense?",
  "extracted_data": {
    "date": null,
    "amount": null,
    "currency": "USD",
    "purpose": null,
    "category": null,
    "has_receipt": false
  },
  "needs_clarification": ["amount"],
  "approval_reason": null
}"#
        .to_string()
    }

    /// Extract text content from message parts
    fn extract_text_from_message(&self, message: &Message) -> String {
        let mut text_parts = Vec::new();
        for part in &message.parts {
            if let Some(part::Content::Text(text)) = &part.content {
                text_parts.push(text.clone());
            }
        }
        text_parts.join(" ")
    }

    /// Check if a message has file attachments
    fn has_file_attachments(&self, message: &Message) -> (bool, Vec<String>) {
        let mut has_files = false;
        let mut file_names = Vec::new();

        for part in &message.parts {
            match &part.content {
                Some(part::Content::Raw(_)) | Some(part::Content::Url(_)) => {
                    has_files = true;
                    if !part.filename.is_empty() {
                        file_names.push(part.filename.clone());
                    } else {
                        file_names.push("unnamed file".to_string());
                    }
                }
                _ => {}
            }
        }

        (has_files, file_names)
    }

    /// Process a reimbursement request using AI
    async fn process_with_ai(
        &self,
        task_id: &str,
        user_message: &str,
        task: Option<&Task>,
        current_message: &Message,
    ) -> Result<ReimbursementResponse, String> {
        let llm_provider = self
            .llm_provider
            .as_ref()
            .ok_or_else(|| "LLM provider not available - cannot process request".to_string())?;

        let system_prompt = self.get_system_prompt();

        // Build conversation history
        let mut history = vec![ChatMessage::system(system_prompt)];

        if let Some(task) = task {
            for msg in &task.history {
                match msg.role {
                    buffa::EnumValue::Known(Role::User) => {
                        let text = self.extract_text_from_message(msg);
                        let (has_files, file_names) = self.has_file_attachments(msg);

                        let message_with_context = if has_files {
                            format!(
                                "{}\n[User uploaded file(s): {}]",
                                text,
                                file_names.join(", ")
                            )
                        } else {
                            text
                        };

                        if !message_with_context.is_empty() {
                            history.push(ChatMessage::user(message_with_context));
                        }
                    }
                    buffa::EnumValue::Known(Role::Agent) => {
                        let text = self.extract_text_from_message(msg);
                        if !text.is_empty() {
                            history.push(ChatMessage::assistant(text));
                        }
                    }
                    _ => {}
                }
            }
        }

        // Check if current message has files
        let (has_files, file_names) = self.has_file_attachments(current_message);
        let current_message_text = if has_files {
            format!(
                "{}\n[User uploaded file(s): {}]",
                user_message,
                file_names.join(", ")
            )
        } else {
            user_message.to_string()
        };

        // Add current user message
        history.push(ChatMessage::user(current_message_text));

        let request = LlmRequest::new(history)
            .temperature(0.7)
            .max_tokens(500)
            .force_json(true);

        let mut stream = llm_provider
            .chat_completion_stream(request)
            .await
            .map_err(|e| e.to_string())?;

        let mut ai_response = String::new();
        let artifact_id = uuid::Uuid::new_v4().to_string();
        let mut tool_calls = ToolCallAccumulator::new();

        use futures::StreamExt;

        while let Some(event) = stream.next().await {
            match event {
                Ok(a2a_agents_common::llm::LlmStreamEvent::ContentChunk(chunk)) => {
                    ai_response.push_str(&chunk);

                    // Push partial output to UI as an appending Artifact Update
                    let artifact = Artifact {
                        artifact_id: artifact_id.clone(),
                        name: "AI Generating...".to_string(),
                        description: String::new(),
                        parts: vec![Part::text(chunk)],
                        metadata: buffa::MessageField::none(),
                        extensions: Vec::new(),
                        ..Default::default()
                    };

                    let update_event = a2a_rs::domain::TaskArtifactUpdateEvent {
                        task_id: task_id.to_string(),
                        context_id: current_message.context_id.clone(),
                        kind: "artifact-update".to_string(),
                        artifact,
                        append: Some(true),
                        last_chunk: Some(false),
                        metadata: None,
                    };

                    let _ = self
                        .streaming
                        .broadcast_artifact_update(task_id, update_event)
                        .await;
                }
                Ok(a2a_agents_common::llm::LlmStreamEvent::ToolCallChunk {
                    id,
                    name,
                    arguments,
                }) => {
                    // Stream the tool-call delta as a typed, structured artifact so
                    // a UI can distinguish it from text and reassemble arguments by
                    // call id (see the `tool-call` metadata marker).
                    ai_response.push_str(&arguments);
                    let partial = tool_calls.push(&id, name.as_deref(), &arguments);
                    let tool_name = partial
                        .name
                        .clone()
                        .unwrap_or_else(|| "unknown".to_string());

                    let artifact = Artifact {
                        artifact_id: artifact_id.clone(),
                        name: format!("Tool Call: {tool_name}"),
                        description: String::new(),
                        parts: vec![Part::text(arguments)],
                        metadata: buffa::MessageField::none(),
                        extensions: Vec::new(),
                        ..Default::default()
                    };

                    let update_event = a2a_rs::domain::TaskArtifactUpdateEvent {
                        task_id: task_id.to_string(),
                        context_id: current_message.context_id.clone(),
                        kind: "artifact-update".to_string(),
                        artifact,
                        append: Some(true),
                        last_chunk: Some(false),
                        metadata: Some(tool_call_metadata(&id, &tool_name, true)),
                    };

                    let _ = self
                        .streaming
                        .broadcast_artifact_update(task_id, update_event)
                        .await;
                }
                Ok(a2a_agents_common::llm::LlmStreamEvent::ToolCall(call)) => {
                    // Reconcile the authoritative final call and emit a terminal,
                    // non-partial artifact carrying the complete arguments.
                    let metadata = tool_call_metadata(&call.id, &call.name, false);
                    let artifact = Artifact {
                        artifact_id: artifact_id.clone(),
                        name: format!("Tool Call: {}", call.name),
                        description: String::new(),
                        parts: vec![Part::text(call.arguments.clone())],
                        metadata: buffa::MessageField::none(),
                        extensions: Vec::new(),
                        ..Default::default()
                    };

                    let update_event = a2a_rs::domain::TaskArtifactUpdateEvent {
                        task_id: task_id.to_string(),
                        context_id: current_message.context_id.clone(),
                        kind: "artifact-update".to_string(),
                        artifact,
                        append: Some(false),
                        last_chunk: Some(true),
                        metadata: Some(metadata),
                    };

                    tool_calls.finalize(call);

                    let _ = self
                        .streaming
                        .broadcast_artifact_update(task_id, update_event)
                        .await;
                }
                Ok(a2a_agents_common::llm::LlmStreamEvent::Reasoning(chunk)) => {
                    // Reasoning-model thinking: surface it as a distinct artifact
                    // but keep it OUT of `ai_response` — the request forces JSON,
                    // and reasoning text would corrupt the parsed answer.
                    let artifact = Artifact {
                        artifact_id: artifact_id.clone(),
                        name: "AI Thinking...".to_string(),
                        description: String::new(),
                        parts: vec![Part::text(chunk)],
                        metadata: buffa::MessageField::none(),
                        extensions: Vec::new(),
                        ..Default::default()
                    };

                    let update_event = a2a_rs::domain::TaskArtifactUpdateEvent {
                        task_id: task_id.to_string(),
                        context_id: current_message.context_id.clone(),
                        kind: "artifact-update".to_string(),
                        artifact,
                        append: Some(true),
                        last_chunk: Some(false),
                        metadata: None,
                    };

                    let _ = self
                        .streaming
                        .broadcast_artifact_update(task_id, update_event)
                        .await;
                }
                Err(e) => {
                    tracing::error!("LLM Stream error: {}", e);
                    break;
                }
            }
        }

        if ai_response.is_empty() {
            return Err("AI returned no text content".to_string());
        }

        info!(
            response_length = ai_response.len(),
            "Received AI response via stream"
        );

        // Parse AI response as JSON
        self.parse_ai_response(&ai_response)
    }

    /// Parse the AI's JSON response into a ReimbursementResponse
    fn parse_ai_response(&self, ai_response: &str) -> Result<ReimbursementResponse, String> {
        // Try to extract JSON from the response (AI might wrap it in markdown)
        let json_str = if let Some(start) = ai_response.find('{') {
            if let Some(end) = ai_response.rfind('}') {
                &ai_response[start..=end]
            } else {
                ai_response
            }
        } else {
            ai_response
        };

        #[derive(Debug, Deserialize)]
        struct AiResponse {
            action: String,
            message: String,
            #[serde(default)]
            extracted_data: Option<ExtractedData>,
            #[serde(default)]
            needs_clarification: Vec<String>,
            #[serde(default)]
            #[allow(dead_code)]
            approval_reason: Option<String>,
        }

        #[derive(Debug, Deserialize)]
        struct ExtractedData {
            #[allow(dead_code)]
            date: Option<String>,
            amount: Option<f64>,
            currency: Option<String>,
            #[allow(dead_code)]
            purpose: Option<String>,
            #[allow(dead_code)]
            category: Option<String>,
            #[allow(dead_code)]
            has_receipt: Option<bool>,
        }

        let parsed: AiResponse = serde_json::from_str(json_str).map_err(|e| {
            format!(
                "Failed to parse AI response as JSON: {}. Response was: {}",
                e, ai_response
            )
        })?;

        match parsed.action.as_str() {
            "approved" | "pending" => {
                let extracted = parsed
                    .extracted_data
                    .ok_or("Missing extracted_data for approval")?;

                let amount = Money::Number {
                    amount: extracted.amount.ok_or("Missing amount")?,
                    currency: extracted.currency.unwrap_or_else(|| "USD".to_string()),
                };

                let status = if parsed.action == "approved" {
                    ProcessingStatus::Approved
                } else {
                    ProcessingStatus::Pending
                };

                let details = ProcessingDetails {
                    approved_amount: Some(amount.clone()),
                    approval_date: Some(Utc::now().to_rfc3339()),
                    approver: Some("AI Assistant".to_string()),
                    rejection_reason: None,
                    required_documents: if extracted.amount.unwrap_or(0.0) > 25.0 {
                        Some(vec!["receipt".to_string()])
                    } else {
                        None
                    },
                };

                Ok(ReimbursementResponse::Result {
                    request_id: format!("req_{}", Uuid::new_v4().simple()),
                    status,
                    message: Some(parsed.message),
                    details: Some(details),
                })
            }
            "question" => {
                // AI needs more information - return as conversational response
                Ok(ReimbursementResponse::Error {
                    code: "NEEDS_MORE_INFO".to_string(),
                    message: parsed.message,
                    field: None,
                    suggestions: if !parsed.needs_clarification.is_empty() {
                        Some(parsed.needs_clarification)
                    } else {
                        None
                    },
                })
            }
            "error" => Ok(ReimbursementResponse::Error {
                code: "PROCESSING_ERROR".to_string(),
                message: parsed.message,
                field: None,
                suggestions: None,
            }),
            _ => Err(format!("Unknown action from AI: {}", parsed.action)),
        }
    }

    /// Parse message content into a reimbursement request
    #[instrument(skip(self, message), fields(message_id = %message.message_id, parts_count = message.parts.len()))]
    fn parse_request(&self, message: &Message) -> Result<ReimbursementRequest, A2AError> {
        // Extract all content from message parts
        let mut text_content = Vec::new();
        let mut data_content: Option<Map<String, Value>> = None;
        let mut file_ids = Vec::new();
        let mut metadata_hints: Map<String, Value> = Map::new();

        for (idx, part) in message.parts.iter().enumerate() {
            let mut meta_map = Map::new();
            if let Some(meta) = part.metadata.as_option()
                && let Ok(Value::Object(map)) = serde_json::to_value(meta)
            {
                meta_map = map;
            }

            match &part.content {
                Some(part::Content::Text(text)) => {
                    debug!(
                        part_index = idx,
                        text_length = text.len(),
                        "Processing text part"
                    );
                    text_content.push(text.clone());
                    // Extract any metadata hints for processing
                    if !meta_map.is_empty() {
                        debug!(metadata_keys = ?meta_map.keys().collect::<Vec<_>>(), "Found text metadata");
                        self.merge_metadata(&mut metadata_hints, &meta_map);
                    }
                }
                Some(part::Content::Data(data)) => {
                    let mut data_map = Map::new();
                    if let Ok(Value::Object(map)) = serde_json::to_value(&**data) {
                        data_map = map;
                    }
                    debug!(part_index = idx, data_keys = ?data_map.keys().collect::<Vec<_>>(), "Processing data part");
                    // Merge multiple data parts
                    if let Some(ref mut existing) = data_content {
                        for (k, v) in data_map {
                            existing.insert(k, v);
                        }
                    } else {
                        data_content = Some(data_map);
                    }
                    // Extract metadata
                    if !meta_map.is_empty() {
                        debug!(metadata_keys = ?meta_map.keys().collect::<Vec<_>>(), "Found data metadata");
                        self.merge_metadata(&mut metadata_hints, &meta_map);
                    }
                }
                Some(part::Content::Raw(_)) | Some(part::Content::Url(_)) => {
                    // Store file references with metadata
                    let file_id = if !part.filename.is_empty() {
                        part.filename.clone()
                    } else {
                        format!("file_{}", Uuid::new_v4().simple())
                    };
                    info!(part_index = idx, file_id = %file_id, mime_type = ?part.media_type, "Processing file part");
                    file_ids.push(file_id.clone());

                    // Store file metadata for later processing
                    let mut file_meta = meta_map.clone();
                    file_meta.insert("file_id".to_string(), Value::String(file_id.clone()));
                    if !part.media_type.is_empty() {
                        file_meta.insert(
                            "mime_type".to_string(),
                            Value::String(part.media_type.clone()),
                        );
                    }
                    debug!(file_id = %file_id, metadata_keys = ?file_meta.keys().collect::<Vec<_>>(), "Storing file metadata");
                    self.store_file_metadata(&file_id, file_meta);
                }
                None => {}
            }
        }

        // Apply metadata hints to data if available
        if let Some(mut data) = data_content {
            // Merge metadata hints into data for enhanced processing
            for (key, value) in metadata_hints {
                if !data.contains_key(&key) {
                    data.insert(key, value);
                }
            }
            data_content = Some(data);
        }

        // Try to parse structured data first
        if let Some(data) = data_content {
            // Try parsing as complete request types
            if let Ok(request) =
                serde_json::from_value::<ReimbursementRequest>(Value::Object(data.clone()))
            {
                return Ok(request);
            }

            // Check if it's a status query
            if let Some(request_id) = data.get("request_id").and_then(|v| v.as_str())
                && (data.len() == 1
                    || data.get("action").and_then(|v| v.as_str()) == Some("status"))
            {
                return Ok(ReimbursementRequest::StatusQuery {
                    request_id: request_id.to_string(),
                });
            }

            // Try to build an initial or form submission request
            let request_id = data.get("request_id").and_then(|v| v.as_str());
            let date = data
                .get("date")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string());
            let amount = self.parse_money_from_value(data.get("amount"));
            let purpose = data
                .get("purpose")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string());

            // Check for category from both direct field and metadata hints
            let category = data
                .get("category")
                .and_then(|v| serde_json::from_value::<ExpenseCategory>(v.clone()).ok())
                .or_else(|| {
                    // Try category_hint from metadata
                    let hint = data.get("category_hint")?.as_str()?;
                    serde_json::from_value::<ExpenseCategory>(Value::String(hint.to_string())).ok()
                })
                .or_else(|| {
                    // Try expense_type from metadata
                    let expense_type = data.get("expense_type")?.as_str()?;
                    Some(match expense_type.to_lowercase().as_str() {
                        "travel" => ExpenseCategory::Travel,
                        "meals" | "meal" => ExpenseCategory::Meals,
                        "supplies" | "supply" => ExpenseCategory::Supplies,
                        "equipment" => ExpenseCategory::Equipment,
                        "training" => ExpenseCategory::Training,
                        _ => ExpenseCategory::Other,
                    })
                });

            let notes = data
                .get("notes")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string());

            let receipt_files = if !file_ids.is_empty() {
                Some(file_ids.clone())
            } else {
                data.get("receipt_files")
                    .and_then(|v| serde_json::from_value::<Vec<String>>(v.clone()).ok())
            };

            if let Some(request_id) = request_id
                && let (Some(date), Some(amount), Some(purpose), Some(category)) = (
                    date.clone(),
                    amount.clone(),
                    purpose.clone(),
                    category.clone(),
                )
            {
                // Form submission with request_id
                return Ok(ReimbursementRequest::FormSubmission {
                    request_id: request_id.to_string(),
                    date,
                    amount,
                    purpose,
                    category,
                    receipt_files,
                    notes,
                });
            } else {
                // Initial request without request_id
                return Ok(ReimbursementRequest::Initial {
                    date,
                    amount,
                    purpose,
                    category,
                    receipt_files,
                });
            }
        }

        // Fall back to text parsing
        if !text_content.is_empty() {
            let combined_text = text_content.join(" ");
            return self.parse_text_request(&combined_text, file_ids);
        }

        Err(A2AError::InvalidParams(
            "No valid request data found in message".to_string(),
        ))
    }

    /// Parse money from various JSON value formats with metadata awareness
    fn parse_money_from_value_with_metadata(
        &self,
        value: Option<&Value>,
        metadata: &Map<String, Value>,
    ) -> Option<Money> {
        // Check for currency hint in metadata
        let default_currency = metadata
            .get("currency")
            .and_then(|v| v.as_str())
            .unwrap_or("USD")
            .to_string();

        match value? {
            Value::String(s) => Some(Money::String(s.clone())),
            Value::Number(n) => n.as_f64().map(|amount| Money::Number {
                amount,
                currency: default_currency,
            }),
            Value::Object(obj) => {
                let amount = obj.get("amount")?.as_f64()?;
                let currency = obj
                    .get("currency")
                    .and_then(|v| v.as_str())
                    .unwrap_or(&default_currency)
                    .to_string();
                Some(Money::Number { amount, currency })
            }
            _ => None,
        }
    }

    /// Parse money from various JSON value formats (legacy method for compatibility)
    fn parse_money_from_value(&self, value: Option<&Value>) -> Option<Money> {
        self.parse_money_from_value_with_metadata(value, &Map::new())
    }

    /// Parse text content for reimbursement information
    fn parse_text_request(
        &self,
        text: &str,
        file_ids: Vec<String>,
    ) -> Result<ReimbursementRequest, A2AError> {
        let lower_text = text.to_lowercase();

        // Check if it's a status query
        if lower_text.contains("status")
            && text.contains("req_")
            && let Some(start) = text.find("req_")
        {
            let request_id: String = text[start..]
                .chars()
                .take_while(|c| c.is_alphanumeric() || *c == '_')
                .collect();
            return Ok(ReimbursementRequest::StatusQuery { request_id });
        }

        // Parse as initial request
        let mut date = None;
        let mut amount = None;
        let mut purpose = None;
        let mut category = None;

        // Extract amount
        if let Some(dollar_pos) = text.find('$') {
            let amount_str: String = text[dollar_pos + 1..]
                .chars()
                .take_while(|c| c.is_numeric() || *c == '.')
                .collect();
            if let Ok(parsed_amount) = amount_str.parse::<f64>() {
                amount = Some(Money::Number {
                    amount: parsed_amount,
                    currency: "USD".to_string(),
                });
            }
        }

        // Extract date patterns (MM/DD/YYYY, YYYY-MM-DD, etc.)
        let date_patterns = [
            r"\d{1,2}/\d{1,2}/\d{4}",
            r"\d{4}-\d{2}-\d{2}",
            r"\d{1,2}-\d{1,2}-\d{4}",
        ];
        for pattern in &date_patterns {
            if let Ok(regex) = regex::Regex::new(pattern)
                && let Some(mat) = regex.find(text)
            {
                date = Some(mat.as_str().to_string());
                break;
            }
        }

        // Extract purpose (text after "for" or "Description:")
        if let Some(desc_pos) = text.find("Description:") {
            // Extract text from "Description:" until next line or end
            let desc_start = desc_pos + "Description:".len();
            let purpose_text = text[desc_start..].lines().next().unwrap_or("").trim();
            if !purpose_text.is_empty() {
                purpose = Some(purpose_text.to_string());
            }
        } else if let Some(for_pos) = lower_text.find(" for ") {
            let purpose_text = text[for_pos + 5..].trim();
            if !purpose_text.is_empty() {
                purpose = Some(purpose_text.to_string());
            }
        }

        // Detect category from keywords
        if lower_text.contains("travel")
            || lower_text.contains("flight")
            || lower_text.contains("hotel")
        {
            category = Some(ExpenseCategory::Travel);
        } else if lower_text.contains("meal")
            || lower_text.contains("lunch")
            || lower_text.contains("dinner")
        {
            category = Some(ExpenseCategory::Meals);
        } else if lower_text.contains("office")
            || lower_text.contains("supply")
            || lower_text.contains("supplies")
        {
            category = Some(ExpenseCategory::Supplies);
        } else if lower_text.contains("equipment") || lower_text.contains("hardware") {
            category = Some(ExpenseCategory::Equipment);
        } else if lower_text.contains("training") || lower_text.contains("course") {
            category = Some(ExpenseCategory::Training);
        } else if lower_text.contains("other") {
            category = Some(ExpenseCategory::Other);
        }

        let receipt_files = if !file_ids.is_empty() {
            Some(file_ids)
        } else {
            None
        };

        Ok(ReimbursementRequest::Initial {
            date,
            amount,
            purpose,
            category,
            receipt_files,
        })
    }

    /// Validate a reimbursement request
    #[instrument(skip(self, request), fields(request_type = ?std::mem::discriminant(request)))]
    fn validate_request(&self, request: &ReimbursementRequest) -> Result<(), A2AError> {
        match request {
            ReimbursementRequest::Initial { amount, .. } => {
                // Validate amount if provided
                if let Some(money) = amount
                    && let Err(e) = money.validate()
                {
                    warn!(error = %e, "Amount validation failed");
                    return Err(A2AError::ValidationError {
                        field: "amount".to_string(),
                        message: e,
                    });
                }
                Ok(())
            }
            ReimbursementRequest::FormSubmission {
                date,
                amount,
                purpose,
                category,
                ..
            } => {
                // Validate required fields
                if date.trim().is_empty() {
                    return Err(A2AError::ValidationError {
                        field: "date".to_string(),
                        message: "Date is required".to_string(),
                    });
                }

                if purpose.trim().is_empty() {
                    return Err(A2AError::ValidationError {
                        field: "purpose".to_string(),
                        message: "Purpose is required".to_string(),
                    });
                }

                // Validate amount
                if let Err(e) = amount.validate() {
                    return Err(A2AError::ValidationError {
                        field: "amount".to_string(),
                        message: e,
                    });
                }

                // Validate against rules
                if !self.validation_rules.allowed_categories.contains(category) {
                    return Err(A2AError::ValidationError {
                        field: "category".to_string(),
                        message: format!("Category {:?} is not allowed", category),
                    });
                }

                // TODO: Add more validation (date range, amount limits, etc.)

                Ok(())
            }
            ReimbursementRequest::StatusQuery { request_id } => {
                if !request_id.starts_with("req_") {
                    return Err(A2AError::ValidationError {
                        field: "request_id".to_string(),
                        message: "Invalid request ID format".to_string(),
                    });
                }
                Ok(())
            }
        }
    }

    /// Process a reimbursement request and generate appropriate response
    #[instrument(skip(self, request), fields(request_type = ?std::mem::discriminant(&request)))]
    fn process_request(
        &self,
        request: ReimbursementRequest,
    ) -> Result<ReimbursementResponse, A2AError> {
        match &request {
            ReimbursementRequest::Initial {
                date,
                amount,
                purpose,
                category,
                receipt_files,
            } => {
                // Generate a new request ID
                let request_id = format!("req_{}", Uuid::new_v4().simple());

                // Check if we have all required fields for auto-processing
                let has_all_fields = date.is_some()
                    && amount.is_some()
                    && purpose.is_some()
                    && !purpose.as_ref().unwrap_or(&String::new()).trim().is_empty()
                    && category.is_some();

                // Check for auto-approval based on amount (if we have all fields)
                if has_all_fields {
                    let auto_approve = match amount.as_ref().unwrap() {
                        Money::Number { amount, .. } if *amount < 100.0 => true,
                        Money::Number { .. } => false, // >= $100
                        Money::String(s) => {
                            // Try to parse string amount
                            let amount_str = s.trim_start_matches('$').replace(",", "");
                            if let Ok(amt) = amount_str.parse::<f64>() {
                                amt < 100.0
                            } else {
                                false
                            }
                        }
                    };

                    if auto_approve {
                        // Auto-approve small expenses with all required fields
                        let details = ProcessingDetails {
                            approved_amount: amount.clone(),
                            approval_date: Some(Utc::now().to_rfc3339()),
                            approver: Some("System Auto-Approval".to_string()),
                            rejection_reason: None,
                            required_documents: None,
                        };

                        return Ok(ReimbursementResponse::Result {
                            request_id: request_id.clone(),
                            status: ProcessingStatus::Approved,
                            message: Some(format!(
                                "✅ Your expense of {} has been automatically approved! No further action needed.",
                                amount.as_ref().unwrap().to_formatted_string()
                            )),
                            details: Some(details),
                        });
                    }
                }

                // Otherwise, return a form for completion/review
                let form_data = FormData {
                    request_id: request_id.clone(),
                    date: date.clone().or_else(|| Some(String::new())),
                    amount: amount
                        .as_ref()
                        .map(|m| m.to_formatted_string())
                        .or_else(|| Some(String::new())),
                    purpose: purpose.clone().or_else(|| Some(String::new())),
                    category: category
                        .as_ref()
                        .map(|c| format!("{:?}", c).to_lowercase())
                        .or_else(|| Some("other".to_string())),
                    receipt_files: receipt_files.clone(),
                    notes: None,
                };

                // Create form schema
                let form_schema = self.create_form_schema();

                Ok(ReimbursementResponse::Form {
                    form: form_schema,
                    form_data,
                    instructions: Some(
                        "Please complete all required fields for your reimbursement request. \
                        Receipts are required for expenses over $25."
                            .to_string(),
                    ),
                })
            }
            ReimbursementRequest::FormSubmission {
                request_id,
                date: _,
                amount,
                purpose: _,
                category: _,
                receipt_files,
                notes: _,
            } => {
                // Store the request
                // Process receipt files with metadata
                let mut receipts = vec![];
                if let Some(files) = receipt_files {
                    for file_id in files {
                        if let Some(metadata) = self.get_file_metadata(file_id) {
                            receipts.push(ReceiptMetadata {
                                file_id: file_id.clone(),
                                file_name: metadata
                                    .get("file_name")
                                    .and_then(|v| v.as_str())
                                    .unwrap_or(file_id)
                                    .to_string(),
                                mime_type: metadata
                                    .get("mime_type")
                                    .and_then(|v| v.as_str())
                                    .unwrap_or("application/octet-stream")
                                    .to_string(),
                                size_bytes: metadata
                                    .get("size_bytes")
                                    .and_then(|v| v.as_u64())
                                    .unwrap_or(0)
                                    as usize,
                                upload_timestamp: Some(Utc::now().to_rfc3339()),
                                extracted_data: None, // TODO: OCR integration
                            });
                        }
                    }
                }

                // Check for auto-approval based on amount
                // In a real system, we'd get this from the message metadata or business rules
                // For now, we'll auto-approve small amounts
                let auto_approve =
                    matches!(amount, Money::Number { amount, .. } if *amount < 100.0);

                let status = if auto_approve {
                    ProcessingStatus::Approved
                } else {
                    ProcessingStatus::Pending
                };

                // NOTE: Task storage is handled by ConnectRpcAdapter
                // We just process the message and return a response
                // (receipts metadata is still tracked in file_metadata_store if needed)

                // In a real implementation, this would trigger workflow processing
                let details = ProcessingDetails {
                    approved_amount: Some(amount.clone()),
                    approval_date: Some(Utc::now().to_rfc3339()),
                    approver: Some("System Auto-Approval".to_string()),
                    rejection_reason: None,
                    required_documents: None,
                };

                Ok(ReimbursementResponse::Result {
                    request_id: request_id.clone(),
                    status,
                    message: Some(if auto_approve {
                        "Your reimbursement request has been auto-approved for amounts under $100."
                            .to_string()
                    } else {
                        "Your reimbursement request has been submitted for review.".to_string()
                    }),
                    details: if auto_approve { Some(details) } else { None },
                })
            }
            ReimbursementRequest::StatusQuery { request_id } => {
                // NOTE: Actual task status is managed by ConnectRpcAdapter/SQLx
                // This handler doesn't have access to task storage directly
                // The client should use tasks/get instead for real status
                {
                    Err(A2AError::ValidationError {
                        field: "request_id".to_string(),
                        message: format!("Request {} not found", request_id),
                    })
                }
            }
        }
    }

    /// Create form schema for reimbursement requests
    fn create_form_schema(&self) -> FormSchema {
        let mut properties = Map::new();

        // Date field
        properties.insert(
            "date".to_string(),
            json!({
                "type": "string",
                "format": "date",
                "title": "Expense Date",
                "description": "Date when the expense was incurred"
            }),
        );

        // Amount field
        properties.insert(
            "amount".to_string(),
            json!({
                "type": "string",
                "format": "money",
                "title": "Amount",
                "description": "Total amount to be reimbursed",
                "pattern": r"^\$?\d+(\.\d{2})?$"
            }),
        );

        // Purpose field
        properties.insert(
            "purpose".to_string(),
            json!({
                "type": "string",
                "title": "Business Purpose",
                "description": "Business justification for the expense",
                "minLength": 10
            }),
        );

        // Category field
        properties.insert(
            "category".to_string(),
            json!({
                "type": "string",
                "title": "Expense Category",
                "enum": ["travel", "meals", "supplies", "equipment", "training", "other"],
                "description": "Category of the expense"
            }),
        );

        // Notes field (optional)
        properties.insert(
            "notes".to_string(),
            json!({
                "type": "string",
                "title": "Additional Notes",
                "description": "Any additional information (optional)"
            }),
        );

        // Request ID (hidden/readonly)
        properties.insert(
            "request_id".to_string(),
            json!({
                "type": "string",
                "title": "Request ID",
                "readOnly": true
            }),
        );

        FormSchema {
            schema_type: "object".to_string(),
            properties,
            required: vec![
                "request_id".to_string(),
                "date".to_string(),
                "amount".to_string(),
                "purpose".to_string(),
                "category".to_string(),
            ],
            dependencies: None,
        }
    }

    /// Convert response to message parts
    fn response_to_parts(&self, response: ReimbursementResponse) -> Vec<Part> {
        match response {
            ReimbursementResponse::Form { ref form, .. } => {
                // Serialize as JSON for both text and data parts
                if let Ok(json_str) = serde_json::to_string_pretty(&response) {
                    let mut metadata = Map::new();
                    metadata.insert(
                        "response_type".to_string(),
                        Value::String("form".to_string()),
                    );
                    metadata.insert(
                        "form_id".to_string(),
                        Value::String(
                            form.properties
                                .get("request_id")
                                .and_then(|v| v.get("default"))
                                .and_then(|v| v.as_str())
                                .unwrap_or("unknown")
                                .to_string(),
                        ),
                    );

                    let metadata_struct: buffa_types::google::protobuf::Struct =
                        serde_json::from_value(Value::Object(metadata)).unwrap_or_default();
                    let val: serde_json::Value =
                        serde_json::from_str(&json_str).unwrap_or_default();
                    let proto_val: buffa_types::google::protobuf::Value =
                        serde_json::from_value(val).unwrap_or_default();

                    vec![
                        Part::text_with_metadata(json_str.clone(), metadata_struct.clone()),
                        Part::data_builder(proto_val)
                            .with_metadata(metadata_struct)
                            .build(),
                    ]
                } else {
                    vec![Part::text("Failed to serialize response".to_string())]
                }
            }
            ReimbursementResponse::Result {
                ref request_id,
                ref status,
                ..
            } => {
                if let Ok(json_str) = serde_json::to_string_pretty(&response) {
                    let mut metadata = Map::new();
                    metadata.insert(
                        "response_type".to_string(),
                        Value::String("result".to_string()),
                    );
                    metadata.insert("request_id".to_string(), Value::String(request_id.clone()));
                    metadata.insert("status".to_string(), Value::String(format!("{:?}", status)));

                    let metadata_struct: buffa_types::google::protobuf::Struct =
                        serde_json::from_value(Value::Object(metadata)).unwrap_or_default();
                    let val: serde_json::Value =
                        serde_json::from_str(&json_str).unwrap_or_default();
                    let proto_val: buffa_types::google::protobuf::Value =
                        serde_json::from_value(val).unwrap_or_default();

                    vec![
                        Part::text_with_metadata(json_str.clone(), metadata_struct.clone()),
                        Part::data_builder(proto_val)
                            .with_metadata(metadata_struct)
                            .build(),
                    ]
                } else {
                    vec![Part::text("Failed to serialize response".to_string())]
                }
            }
            ReimbursementResponse::Error {
                ref message,
                ref code,
                ..
            } => {
                let mut metadata = Map::new();
                metadata.insert(
                    "response_type".to_string(),
                    Value::String("error".to_string()),
                );
                metadata.insert("error_code".to_string(), Value::String(code.clone()));

                let metadata_struct: buffa_types::google::protobuf::Struct =
                    serde_json::from_value(Value::Object(metadata)).unwrap_or_default();

                vec![Part::text_with_metadata(message.clone(), metadata_struct)]
            }
        }
    }
}

#[async_trait]
impl AsyncMessageHandler for ReimbursementHandler {
    #[instrument(skip(self, message), fields(
        task_id = %task_id,
        message_id = %message.message_id,
        session_id = ?_session_id,
        parts_count = message.parts.len()
    ))]
    async fn process_message(
        &self,
        task_id: &str,
        message: &Message,
        _session_id: Option<&str>,
    ) -> Result<Task, A2AError> {
        error!(
            "🚨 HANDLER CALLED: Processing reimbursement request for task_id={}",
            task_id
        );
        info!("Processing reimbursement request");

        let id: TaskId = task_id.parse()?;

        // Check if task exists and get its current state
        let existing_task = if self.task_lifecycle.exists(&id).await? {
            Some(self.task_lifecycle.get(&id, Some(50)).await?)
        } else {
            None
        };

        // If task doesn't exist, create it
        if existing_task.is_none() {
            let context_id = if message.context_id.is_empty() {
                Uuid::new_v4().to_string()
            } else {
                message.context_id.clone()
            };
            let context_id: ContextId = context_id.parse()?;
            self.task_lifecycle.create(&id, &context_id).await?;
        }

        // Check if this task already has a completed/approved expense
        // If so, treat this as a follow-up conversation message
        if let Some(ref task) = existing_task
            && task.status.state == TaskState::Completed
        {
            // This is a follow-up to a completed task
            // Add user message to history
            self.update_and_broadcast(&id, TaskState::Working, Some(message.clone()))
                .await?;

            // Send a simple acknowledgment
            let response_message = Message::builder()
                .role(Role::Agent)
                .parts(vec![Part::text("Your expense has already been processed. Is there anything else I can help you with?".to_string())])
                .message_id(Uuid::new_v4().to_string())
                .context_id(message.context_id.clone())
                .build();

            // Add agent response and return
            let final_task = self
                .update_and_broadcast(&id, TaskState::Completed, Some(response_message))
                .await?;

            return Ok(final_task);
        }

        // Add the user's message to history first
        self.update_and_broadcast(&id, TaskState::Working, Some(message.clone()))
            .await?;

        // Send immediate acknowledgment
        let ack_message = Message::builder()
            .role(Role::Agent)
            .parts(vec![Part::text("Processing your request...".to_string())])
            .message_id(Uuid::new_v4().to_string())
            .context_id(message.context_id.clone())
            .build();

        self.update_and_broadcast(&id, TaskState::Working, Some(ack_message))
            .await?;

        // Clone what we need for the background task
        let handler = self.clone();
        let task_id_owned = task_id.to_string();
        let id_owned = id.clone();
        let message_owned = message.clone();
        let context_id = message.context_id.clone();

        // Log before spawning
        info!(task_id = %task_id, "About to spawn background worker for AI processing");

        // Spawn background task to process with AI
        let spawn_result = tokio::task::spawn(async move {
            info!(task_id = %task_id_owned, "🚀 Background worker STARTED - beginning AI processing");

            // Extract the user message text and optional task
            let text_content = handler.extract_text_from_message(&message_owned);

            // Get current task for context
            let current_task = match handler.task_lifecycle.get(&id_owned, Some(50)).await {
                Ok(task) => {
                    info!(task_id = %task_id_owned, history_count = task.history.len(), "Retrieved task for AI processing");
                    Some(task)
                }
                Err(e) => {
                    error!(task_id = %task_id_owned, error = %e, "Failed to get task for AI processing");
                    None
                }
            };

            // Process with AI
            info!(task_id = %task_id_owned, "Calling AI for processing");
            let (response, auto_approved) = match handler
                .process_with_ai(
                    &task_id_owned,
                    &text_content,
                    current_task.as_ref(),
                    &message_owned,
                )
                .await
            {
                Ok(resp) => {
                    info!(task_id = %task_id_owned, response_type = ?std::mem::discriminant(&resp), "AI processed request successfully");
                    let auto_approved = matches!(&resp, ReimbursementResponse::Result { status, .. } if matches!(status, super::types::ProcessingStatus::Approved));
                    (resp, auto_approved)
                }
                Err(e) => {
                    error!(task_id = %task_id_owned, error = %e, "AI processing failed");
                    let resp = ReimbursementResponse::Error {
                        code: "AI_PROCESSING_ERROR".to_string(),
                        message: format!("I'm having trouble processing your request. Please try again or provide more details. Error: {}", e),
                        field: None,
                        suggestions: Some(vec![
                            "Provide the date of the expense".to_string(),
                            "Specify the amount".to_string(),
                            "Describe the business purpose".to_string(),
                            "Select a category (travel, meals, supplies, equipment, training, or other)".to_string(),
                        ]),
                    };
                    (resp, false)
                }
            };

            // Update metrics
            handler.update_metrics(&response, auto_approved);

            // Determine task state based on response type
            let task_state = match &response {
                ReimbursementResponse::Form { .. } => TaskState::InputRequired,
                ReimbursementResponse::Result { status, .. } => match status {
                    ProcessingStatus::RequiresAdditionalInfo => TaskState::InputRequired,
                    _ => TaskState::Completed,
                },
                ReimbursementResponse::Error { code, .. } => {
                    if code == "NEEDS_MORE_INFO" {
                        TaskState::InputRequired
                    } else {
                        TaskState::Failed
                    }
                }
            };

            // Create response message
            let response_parts = handler.response_to_parts(response);
            let response_message = Message::builder()
                .role(Role::Agent)
                .parts(response_parts)
                .message_id(Uuid::new_v4().to_string())
                .context_id(context_id)
                .build();

            // Update task with AI response
            info!(task_id = %task_id_owned, new_state = ?task_state, "Updating task with AI response");
            match handler
                .update_and_broadcast(&id_owned, task_state, Some(response_message))
                .await
            {
                Ok(updated_task) => {
                    info!(
                        task_id = %task_id_owned,
                        task_state = ?task_state,
                        history_count = updated_task.history.len(),
                        "Background worker: Successfully updated task with AI response (push notification should have been sent)"
                    );
                }
                Err(e) => {
                    error!(task_id = %task_id_owned, error = %e, "Background worker: Failed to update task with AI response");
                }
            }

            info!(task_id = %task_id_owned, "✅ Background worker COMPLETED successfully");
        });

        // Log that the task was spawned
        info!(task_id = %task_id, background_task_id = ?spawn_result.id(), "Background worker spawned successfully");

        // Return immediately - client will get updates via polling or push notifications
        info!(task_id = %task_id, "Returning immediate acknowledgment, AI processing in background");

        // Get the updated task with the acknowledgment message
        let final_task = self.task_lifecycle.get(&id, Some(50)).await?;
        Ok(final_task)
    }

    async fn validate_message(&self, message: &Message) -> Result<(), A2AError> {
        if message.parts.is_empty() {
            return Err(A2AError::ValidationError {
                field: "message.parts".to_string(),
                message: "Message must contain at least one part".to_string(),
            });
        }

        // Validate individual parts
        for (idx, part) in message.parts.iter().enumerate() {
            match &part.content {
                Some(part::Content::Text(text)) => {
                    if text.trim().is_empty() {
                        return Err(A2AError::ValidationError {
                            field: format!("parts[{}].text", idx),
                            message: "Text content cannot be empty".to_string(),
                        });
                    }
                }
                Some(part::Content::Data(data)) => {
                    if data.kind.is_none() {
                        return Err(A2AError::ValidationError {
                            field: format!("parts[{}].data", idx),
                            message: "Data content cannot be empty".to_string(),
                        });
                    }
                }
                Some(part::Content::Raw(bytes)) => {
                    if bytes.is_empty() {
                        return Err(A2AError::ValidationError {
                            field: format!("parts[{}].file", idx),
                            message: "File cannot be empty".to_string(),
                        });
                    }

                    // Validate mime type for receipts
                    if !part.media_type.is_empty() {
                        let supported_types = [
                            "image/jpeg",
                            "image/png",
                            "image/gif",
                            "image/webp",
                            "application/pdf",
                            "image/heic",
                            "image/heif",
                        ];
                        if !supported_types.contains(&part.media_type.as_str()) {
                            return Err(A2AError::ContentTypeNotSupported(format!(
                                "Unsupported file type '{}'. Supported types: {}",
                                part.media_type,
                                supported_types.join(", ")
                            )));
                        }
                    }
                }
                Some(part::Content::Url(uri)) => {
                    if uri.is_empty() {
                        return Err(A2AError::ValidationError {
                            field: format!("parts[{}].file", idx),
                            message: "File URL cannot be empty".to_string(),
                        });
                    }

                    // Validate mime type for receipts
                    if !part.media_type.is_empty() {
                        let supported_types = [
                            "image/jpeg",
                            "image/png",
                            "image/gif",
                            "image/webp",
                            "application/pdf",
                            "image/heic",
                            "image/heif",
                        ];
                        if !supported_types.contains(&part.media_type.as_str()) {
                            return Err(A2AError::ContentTypeNotSupported(format!(
                                "Unsupported file type '{}'. Supported types: {}",
                                part.media_type,
                                supported_types.join(", ")
                            )));
                        }
                    }
                }
                None => {
                    return Err(A2AError::ValidationError {
                        field: format!("parts[{}]", idx),
                        message: "Part content is required".to_string(),
                    });
                }
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tool_call_streaming_tests {
    use super::*;
    use a2a_agents_common::llm::{
        LlmError, LlmProvider, LlmResponse, LlmStreamEvent, ToolCall as LlmToolCall,
    };
    use a2a_rs::adapter::{InMemoryStreamingHandler, InMemoryTaskStorage};
    use a2a_rs::port::{AsyncStreamingHandler, NoopPushNotifier, UpdateEvent};
    use futures::StreamExt;
    use futures::stream::{self, BoxStream};

    /// A fake provider that streams a tool call as a partial chunk followed by the
    /// finalized call — the exact shape the handler turns into `tool-call`
    /// metadata on broadcast artifact updates.
    struct ToolCallProvider;

    #[async_trait]
    impl LlmProvider for ToolCallProvider {
        async fn chat_completion(&self, _request: LlmRequest) -> Result<LlmResponse, LlmError> {
            Err(LlmError::ProviderError("unused in this test".to_string()))
        }

        async fn chat_completion_stream(
            &self,
            _request: LlmRequest,
        ) -> Result<BoxStream<'static, Result<LlmStreamEvent, LlmError>>, LlmError> {
            let events = vec![
                Ok(LlmStreamEvent::ToolCallChunk {
                    id: "call_1".to_string(),
                    name: Some("create_reimbursement".to_string()),
                    arguments: "{\"amount\":".to_string(),
                }),
                Ok(LlmStreamEvent::ToolCall(LlmToolCall {
                    id: "call_1".to_string(),
                    name: "create_reimbursement".to_string(),
                    arguments: "{\"amount\":50}".to_string(),
                })),
            ];
            Ok(stream::iter(events).boxed())
        }
    }

    /// End-to-end: a tool call streamed by the provider becomes broadcast artifact
    /// updates whose `tool-call` metadata reaches a `combined_update_stream`
    /// subscriber — the same stream the SSE transport serializes to clients. The
    /// accumulator and `tool_call_metadata` builder are unit-tested in isolation;
    /// this covers the broadcast wiring between them, previously only
    /// compile-checked.
    #[tokio::test]
    async fn tool_call_metadata_reaches_stream_subscribers() {
        let streaming = InMemoryStreamingHandler::new();
        let handler = ReimbursementHandler::with_llm(
            InMemoryTaskStorage::new(),
            streaming.clone(),
            NoopPushNotifier,
            Some(Arc::new(ToolCallProvider) as Arc<dyn LlmProvider>),
        );

        let task_id = "task-tc";
        let message = Message::user_text("reimburse me".to_string(), "m1".to_string());

        // Drive the AI streaming path directly. The final JSON parse may fail (tool
        // arguments aren't a ReimbursementResponse), but the artifact broadcasts —
        // the wiring under test — happen during the stream, before that parse.
        let _ = handler
            .process_with_ai(task_id, "reimburse me", None, &message)
            .await;

        // Replay everything buffered for the task (id > 0): the SSE source.
        let mut updates = streaming
            .combined_update_stream(task_id, Some(0))
            .await
            .unwrap();

        let mut tool_call_updates = Vec::new();
        while let Ok(Some(Ok(seq))) =
            tokio::time::timeout(std::time::Duration::from_secs(2), updates.next()).await
        {
            if let UpdateEvent::ArtifactUpdate(ev) = seq.event {
                let is_tool_call = ev
                    .metadata
                    .as_ref()
                    .and_then(|m| m.get("type"))
                    .and_then(Value::as_str)
                    == Some("tool-call");
                if is_tool_call {
                    tool_call_updates.push(ev);
                }
            }
        }

        // Both the partial chunk and the finalized call surfaced as tool-call
        // metadata, carrying the call id and tool name end-to-end.
        assert_eq!(
            tool_call_updates.len(),
            2,
            "expected partial + final tool-call artifact updates"
        );
        for ev in &tool_call_updates {
            let meta = ev.metadata.as_ref().unwrap();
            assert_eq!(meta.get("call_id").and_then(Value::as_str), Some("call_1"));
            assert_eq!(
                meta.get("tool_name").and_then(Value::as_str),
                Some("create_reimbursement")
            );
        }
        let partials: Vec<bool> = tool_call_updates
            .iter()
            .map(|ev| {
                ev.metadata
                    .as_ref()
                    .unwrap()
                    .get("partial")
                    .and_then(Value::as_bool)
                    .unwrap()
            })
            .collect();
        assert_eq!(
            partials,
            vec![true, false],
            "partial streamed chunk, then the finalized non-partial call"
        );
    }
}