agntcy-shadi-cli 0.1.5

use std::fs;
use std::path::PathBuf;
use std::sync::Arc;
use std::thread;
use std::time::Duration;

use a2a::event::StreamResponse;
use a2a::*;
use a2a_client::A2AClient;
use a2a_server::{
    AgentExecutor, DefaultRequestHandler, InMemoryTaskStore, RequestHandler,
    ServiceParams as A2AServiceParams,
};
use agent_secrets::{AgentSecretAccess, AgentVerifier, SecretResult, SessionContext};
use async_trait::async_trait;
use futures::stream::BoxStream;
use shadi_a2a::{A2AChannelBuilder, A2AGroupChannelBuilder, SLIM_SRC_METADATA_KEY, SlimRpcHandler};
use slim_bindings::{Name, Service};
use slim_rpc::Server;
use tokio::runtime::Builder as TokioRuntimeBuilder;
use tokio::sync::Notify;

use crate::cli_types::{SlimA2ACollaborateArgs, SlimA2AEchoPeerArgs, SlimA2ASendArgs};
use crate::slim_shell::{
    build_client_config_for_endpoint, build_server_config_for_endpoint, format_slim_error,
    parse_name, resolve_client_tls_material_for_agent, resolve_server_tls_material,
    resolve_slim_auth,
};

const DEFAULT_SLIM_ENDPOINT: &str = "127.0.0.1:47357";
pub(crate) const SHELL_A2A_ECHO_PEER_USAGE: &str =
    "usage: /slim a2a-echo-peer [--endpoint HOST:PORT] [--agent-id ID] [--listen-timeout SECONDS] [--ready-file PATH] [--start-local-node]";
pub(crate) const SHELL_A2A_SEND_USAGE: &str =
    "usage: /slim a2a-send [--endpoint HOST:PORT] [--agent-id ID] [--peer-agent-id ID] [--destination NAME] [--message TEXT...] [--stream] [--timeout SECONDS] [--session-id ID]";
pub(crate) const SHELL_A2A_COLLABORATE_USAGE: &str =
    "usage: /slim a2a-collaborate <peer1,peer2,...> [--endpoint HOST:PORT] [--agent-id ID] [--timeout SECONDS] [--message TEXT...]";

struct VerifiedSessionVerifier;

impl AgentVerifier for VerifiedSessionVerifier {
    fn verify(&self, session: &SessionContext) -> SecretResult<()> {
        AgentSecretAccess::require_verified(session)
    }
}

struct SlimA2AExecutor {
    agent_name: String,
}

#[async_trait]
impl AgentExecutor for SlimA2AExecutor {
    fn execute(
        &self,
        ctx: a2a_server::ExecutorContext,
    ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
        let input = ctx
            .message
            .as_ref()
            .map(readable_message_text)
            .unwrap_or_else(|| "(no request message)".to_string());
        let response = Message {
            message_id: new_message_id(),
            context_id: Some(ctx.context_id.clone()),
            task_id: Some(ctx.task_id.clone()),
            role: Role::Agent,
            parts: vec![Part::text(format!("echo:{}:{}", self.agent_name, input))],
            metadata: None,
            extensions: None,
            reference_task_ids: None,
        };
        let history = ctx.message.clone().map(|message| vec![message]);

        Box::pin(futures::stream::iter(vec![
            Ok(StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
                task_id: ctx.task_id.clone(),
                context_id: ctx.context_id.clone(),
                status: TaskStatus {
                    state: TaskState::Working,
                    message: None,
                    timestamp: None,
                },
                metadata: None,
            })),
            Ok(StreamResponse::Task(Task {
                id: ctx.task_id,
                context_id: ctx.context_id,
                status: TaskStatus {
                    state: TaskState::Completed,
                    message: Some(response),
                    timestamp: None,
                },
                artifacts: None,
                history,
                metadata: None,
            })),
        ]))
    }

    fn cancel(
        &self,
        ctx: a2a_server::ExecutorContext,
    ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
        Box::pin(futures::stream::once(async move {
            Ok(StreamResponse::Task(Task {
                id: ctx.task_id,
                context_id: ctx.context_id,
                status: TaskStatus {
                    state: TaskState::Canceled,
                    message: None,
                    timestamp: None,
                },
                artifacts: None,
                history: None,
                metadata: None,
            }))
        }))
    }
}

struct SlimA2AHandler {
    inner: DefaultRequestHandler,
    card: AgentCard,
    request_seen: Arc<Notify>,
}

impl SlimA2AHandler {
    fn new(agent_name: String, target: String, request_seen: Arc<Notify>) -> Self {
        Self {
            inner: DefaultRequestHandler::new(
                SlimA2AExecutor {
                    agent_name: agent_name.clone(),
                },
                InMemoryTaskStore::new(),
            ),
            card: a2a_agent_card(&agent_name, &target),
            request_seen,
        }
    }
}

#[async_trait]
impl RequestHandler for SlimA2AHandler {
    async fn send_message(
        &self,
        params: &A2AServiceParams,
        req: SendMessageRequest,
    ) -> Result<SendMessageResponse, A2AError> {
        let result = self.inner.send_message(params, req).await;
        if result.is_ok() {
            self.request_seen.notify_waiters();
        }
        result
    }

    async fn send_streaming_message(
        &self,
        params: &A2AServiceParams,
        req: SendMessageRequest,
    ) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
        let result = self.inner.send_streaming_message(params, req).await;
        if result.is_ok() {
            self.request_seen.notify_waiters();
        }
        result
    }

    async fn get_task(
        &self,
        params: &A2AServiceParams,
        req: GetTaskRequest,
    ) -> Result<Task, A2AError> {
        self.inner.get_task(params, req).await
    }

    async fn list_tasks(
        &self,
        params: &A2AServiceParams,
        req: ListTasksRequest,
    ) -> Result<ListTasksResponse, A2AError> {
        self.inner.list_tasks(params, req).await
    }

    async fn cancel_task(
        &self,
        params: &A2AServiceParams,
        req: CancelTaskRequest,
    ) -> Result<Task, A2AError> {
        self.inner.cancel_task(params, req).await
    }

    async fn subscribe_to_task(
        &self,
        params: &A2AServiceParams,
        req: SubscribeToTaskRequest,
    ) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
        self.inner.subscribe_to_task(params, req).await
    }

    async fn create_push_config(
        &self,
        params: &A2AServiceParams,
        req: TaskPushNotificationConfig,
    ) -> Result<TaskPushNotificationConfig, A2AError> {
        self.inner.create_push_config(params, req).await
    }

    async fn get_push_config(
        &self,
        params: &A2AServiceParams,
        req: GetTaskPushNotificationConfigRequest,
    ) -> Result<TaskPushNotificationConfig, A2AError> {
        self.inner.get_push_config(params, req).await
    }

    async fn list_push_configs(
        &self,
        params: &A2AServiceParams,
        req: ListTaskPushNotificationConfigsRequest,
    ) -> Result<ListTaskPushNotificationConfigsResponse, A2AError> {
        self.inner.list_push_configs(params, req).await
    }

    async fn delete_push_config(
        &self,
        params: &A2AServiceParams,
        req: DeleteTaskPushNotificationConfigRequest,
    ) -> Result<(), A2AError> {
        self.inner.delete_push_config(params, req).await
    }

    async fn get_extended_agent_card(
        &self,
        _params: &A2AServiceParams,
        _req: GetExtendedAgentCardRequest,
    ) -> Result<AgentCard, A2AError> {
        Ok(self.card.clone())
    }
}

pub(crate) fn run_a2a_echo_peer(args: SlimA2AEchoPeerArgs) -> Result<(), String> {
    let endpoint = resolve_endpoint(args.endpoint.as_deref());
    let auth = resolve_slim_auth(&args.agent_id)?;
    let peer_name = slim_name(&args.agent_id);
    let client_tls = resolve_client_tls_material_for_agent(Some(&args.agent_id))?;
    let server_tls = resolve_server_tls_material()?;

    let node_service = if args.start_local_node {
        let service = Service::new(format!("shadictl-a2a-node-{}", std::process::id()));
        service
            .run_server(build_server_config_for_endpoint(&endpoint, &server_tls))
            .map_err(format_slim_error)?;
        thread::sleep(Duration::from_millis(300));
        Some(service)
    } else {
        None
    };

    let service = Service::new(format!("shadictl-a2a-peer-{}", std::process::id()));
    let connection_id = service
        .connect(build_client_config_for_endpoint(&endpoint, &client_tls))
        .map_err(format_slim_error)?;
    let peer_name_ref = Arc::new(parse_name(&peer_name)?);
    let app = shadi_identity::create_app(&service, peer_name_ref.clone(), &auth)
        .map_err(format_slim_error)?;
    app.subscribe(peer_name_ref.clone(), Some(connection_id))
        .map_err(format_slim_error)?;

    let server = Arc::new(Server::new_with_shared_rx_and_connection(
        app.inner(),
        app.name().as_slim_name(),
        None,
        app.notification_receiver(),
        Some(slim_bindings::get_runtime()),
    ));
    let request_seen = Arc::new(Notify::new());
    let handler = Arc::new(SlimA2AHandler::new(
        peer_name.clone(),
        format!("slimrpc://{}", peer_name),
        request_seen.clone(),
    ));
    SlimRpcHandler::new(handler).register(server.as_ref());

    let ready_file = args.ready_file.clone();
    let wait_seconds = args.listen_timeout_seconds;
    let endpoint_label = endpoint.clone();
    let peer_label = peer_name.clone();
    let runtime = TokioRuntimeBuilder::new_current_thread()
        .enable_all()
        .build()
        .map_err(|err| format!("failed to create tokio runtime: {}", err))?;

    let serve_result = runtime.block_on(async move {
        let server_task = {
            let server = server.clone();
            tokio::spawn(async move {
                server
                    .serve()
                    .await
                    .map_err(|err| format!("A2A SLIMRPC server failed: {}", err))
            })
        };

        tokio::time::sleep(Duration::from_millis(300)).await;
        if let Some(ready_file) = &ready_file {
            fs::write(ready_file, b"ready")
                .map_err(|err| format!("failed to write {}: {}", ready_file.display(), err))?;
        }

        println!("[shadictl a2a-peer] ready as {} on {}", peer_label, endpoint_label);

        let request_result = tokio::time::timeout(
            Duration::from_secs(wait_seconds),
            request_seen.notified(),
        )
        .await;

        if request_result.is_ok() {
            tokio::time::sleep(Duration::from_millis(300)).await;
        }

        server.shutdown().await;
        let server_status = server_task
            .await
            .map_err(|err| format!("failed to join A2A SLIMRPC server task: {}", err))?;
        server_status?;

        request_result.map_err(|_| {
            format!(
                "timed out waiting for A2A request after {}s",
                wait_seconds
            )
        })?;

        Ok::<(), String>(())
    });

    let _ = app.unsubscribe(peer_name_ref, Some(connection_id));
    let _ = service.disconnect(connection_id);
    let _ = service.shutdown();

    if let Some(service) = node_service {
        let _ = service.stop_server(endpoint.clone());
        let _ = service.shutdown();
    }

    serve_result
}

pub(crate) fn run_a2a_send(args: SlimA2ASendArgs) -> Result<(), String> {
    let detail = run_a2a_send_once(&args)?;
    println!("{}", detail);
    Ok(())
}

pub(crate) fn run_a2a_collaborate(args: SlimA2ACollaborateArgs) -> Result<(), String> {
    let detail = run_a2a_collaborate_once(&args)?;
    println!("{}", detail);
    Ok(())
}

pub(crate) fn parse_shell_a2a_echo_peer_args(
    args: &[&str],
) -> Result<SlimA2AEchoPeerArgs, String> {
    let mut parsed = SlimA2AEchoPeerArgs {
        endpoint: None,
        agent_id: "secops-a".to_string(),
        listen_timeout_seconds: 20,
        ready_file: None,
        start_local_node: false,
    };

    let mut index = 0;
    while index < args.len() {
        match args[index] {
            "--endpoint" => {
                parsed.endpoint = Some(next_value(args, &mut index, SHELL_A2A_ECHO_PEER_USAGE)?.to_string());
            }
            "--agent-id" => {
                parsed.agent_id = next_value(args, &mut index, SHELL_A2A_ECHO_PEER_USAGE)?.to_string();
            }
            "--listen-timeout" | "--listen-timeout-seconds" => {
                let value = next_value(args, &mut index, SHELL_A2A_ECHO_PEER_USAGE)?;
                parsed.listen_timeout_seconds = value
                    .parse::<u64>()
                    .map_err(|_| format!("invalid timeout value: {value}"))?;
            }
            "--ready-file" => {
                parsed.ready_file = Some(PathBuf::from(next_value(
                    args,
                    &mut index,
                    SHELL_A2A_ECHO_PEER_USAGE,
                )?));
            }
            "--start-local-node" => {
                parsed.start_local_node = true;
            }
            _ => return Err(SHELL_A2A_ECHO_PEER_USAGE.to_string()),
        }
        index += 1;
    }

    Ok(parsed)
}

pub(crate) fn parse_shell_a2a_send_args(args: &[&str]) -> Result<SlimA2ASendArgs, String> {
    let mut parsed = SlimA2ASendArgs {
        endpoint: None,
        agent_id: "avatar".to_string(),
        peer_agent_id: "secops-a".to_string(),
        destination: None,
        message: "hello from SHADI A2A".to_string(),
        stream: false,
        timeout_seconds: 20,
        session_id: "shadictl-a2a-session".to_string(),
    };

    let mut index = 0;
    while index < args.len() {
        match args[index] {
            "--endpoint" => {
                parsed.endpoint = Some(next_value(args, &mut index, SHELL_A2A_SEND_USAGE)?.to_string());
            }
            "--agent-id" => {
                parsed.agent_id = next_value(args, &mut index, SHELL_A2A_SEND_USAGE)?.to_string();
            }
            "--peer-agent-id" => {
                parsed.peer_agent_id = next_value(args, &mut index, SHELL_A2A_SEND_USAGE)?.to_string();
            }
            "--destination" => {
                parsed.destination = Some(next_value(args, &mut index, SHELL_A2A_SEND_USAGE)?.to_string());
            }
            "--message" => {
                let (message, next_index) = collect_message_value(args, index + 1)?;
                parsed.message = message;
                index = next_index;
                continue;
            }
            "--stream" => {
                parsed.stream = true;
            }
            "--timeout" | "--timeout-seconds" => {
                let value = next_value(args, &mut index, SHELL_A2A_SEND_USAGE)?;
                parsed.timeout_seconds = value
                    .parse::<u64>()
                    .map_err(|_| format!("invalid timeout value: {value}"))?;
            }
            "--session-id" => {
                parsed.session_id = next_value(args, &mut index, SHELL_A2A_SEND_USAGE)?.to_string();
            }
            _ => return Err(SHELL_A2A_SEND_USAGE.to_string()),
        }
        index += 1;
    }

    Ok(parsed)
}

/// `--message` value collection reuses `collect_message_value` (and its
/// `is_send_flag` stop-list), since every flag `a2a-collaborate` accepts
/// (`--endpoint`, `--agent-id`, `--timeout[-seconds]`) is already in that list.
pub(crate) fn parse_shell_a2a_collaborate_args(
    args: &[&str],
) -> Result<SlimA2ACollaborateArgs, String> {
    let (peer_agent_ids, rest) = args
        .split_first()
        .ok_or_else(|| SHELL_A2A_COLLABORATE_USAGE.to_string())?;

    let mut parsed = SlimA2ACollaborateArgs {
        endpoint: None,
        agent_id: "avatar".to_string(),
        peer_agent_ids: peer_agent_ids.to_string(),
        message: "hello from SHADI A2A".to_string(),
        timeout_seconds: 20,
    };

    let mut index = 0;
    while index < rest.len() {
        match rest[index] {
            "--endpoint" => {
                parsed.endpoint =
                    Some(next_value(rest, &mut index, SHELL_A2A_COLLABORATE_USAGE)?.to_string());
            }
            "--agent-id" => {
                parsed.agent_id =
                    next_value(rest, &mut index, SHELL_A2A_COLLABORATE_USAGE)?.to_string();
            }
            "--message" => {
                let (message, next_index) = collect_message_value(rest, index + 1)?;
                parsed.message = message;
                index = next_index;
                continue;
            }
            "--timeout" | "--timeout-seconds" => {
                let value = next_value(rest, &mut index, SHELL_A2A_COLLABORATE_USAGE)?;
                parsed.timeout_seconds = value
                    .parse::<u64>()
                    .map_err(|_| format!("invalid timeout value: {value}"))?;
            }
            _ => return Err(SHELL_A2A_COLLABORATE_USAGE.to_string()),
        }
        index += 1;
    }

    Ok(parsed)
}

fn run_a2a_send_once(args: &SlimA2ASendArgs) -> Result<String, String> {
    let endpoint = resolve_endpoint(args.endpoint.as_deref());
    let auth = resolve_slim_auth(&args.agent_id)?;
    let client_tls = resolve_client_tls_material_for_agent(Some(&args.agent_id))?;
    let local_name = slim_name(&args.agent_id);
    let destination = args
        .destination
        .clone()
        .unwrap_or_else(|| slim_name(&args.peer_agent_id));

    let service = Service::new(format!("shadictl-a2a-client-{}", std::process::id()));
    let connection_id = service
        .connect(build_client_config_for_endpoint(&endpoint, &client_tls))
        .map_err(format_slim_error)?;
    let local_name_ref = Arc::new(parse_name(&local_name)?);
    let remote_name_ref = Arc::new(parse_name(&destination)?);
    let app = shadi_identity::create_app(&service, local_name_ref.clone(), &auth)
        .map_err(format_slim_error)?;
    app.subscribe(local_name_ref.clone(), Some(connection_id))
        .map_err(format_slim_error)?;

    let runtime = TokioRuntimeBuilder::new_current_thread()
        .enable_all()
        .build()
        .map_err(|err| format!("failed to create tokio runtime: {}", err))?;

    let mut session = SessionContext::new(&args.agent_id, &args.session_id);
    session.verified = true;
    let verifier: Arc<dyn AgentVerifier> = Arc::new(VerifiedSessionVerifier);
    // slim_rpc::Channel captures `tokio::runtime::Handle::current()` at construction,
    // so build the transport inside the runtime context.
    let channel = {
        let _enter = runtime.enter();
        A2AChannelBuilder::new(app.clone(), remote_name_ref, verifier, session)
            .connection_id(connection_id)
            .build()
    };
    let client = A2AClient::new(Box::new(channel));
    let request = SendMessageRequest {
        message: Message::new(Role::User, vec![Part::text(args.message.clone())]),
        configuration: None,
        metadata: None,
        tenant: None,
    };

    let response_detail = runtime
        .block_on(async {
            tokio::time::timeout(Duration::from_secs(args.timeout_seconds), async {
                if args.stream {
                    use futures::StreamExt;

                    let stream = client.send_streaming_message(&request).await?;
                    let events = stream.collect::<Vec<_>>().await;
                    client.destroy().await?;
                    Ok::<String, A2AError>(describe_a2a_stream(&events))
                } else {
                    let response = client.send_message(&request).await?;
                    client.destroy().await?;
                    Ok::<String, A2AError>(describe_a2a_response(&response))
                }
            })
            .await
            .map_err(|_| {
                A2AError::internal(format!(
                    "timed out waiting for A2A response after {}s",
                    args.timeout_seconds
                ))
            })?
        })
        .map_err(|err| format!("failed to send A2A message: {}", err))?;

    let _ = app.unsubscribe(local_name_ref, Some(connection_id));
    let _ = service.disconnect(connection_id);
    let _ = service.shutdown();

    Ok(format!(
        "sent {:?} to {} via {} and received {}",
        args.message,
        destination,
        local_name,
        response_detail
    ))
}

/// Broadcast one A2A `Message` to a SLIM group of `peer_agent_ids` via the
/// `Collaborate` operation, while also listening for whatever the *other*
/// members independently broadcast to this agent — a group is many-to-many, so
/// every member plays both roles (see the SLIMRPC collaborative channel
/// extension spec). Runs for `--timeout-seconds` before shutting down.
fn run_a2a_collaborate_once(args: &SlimA2ACollaborateArgs) -> Result<String, String> {
    let endpoint = resolve_endpoint(args.endpoint.as_deref());
    let auth = resolve_slim_auth(&args.agent_id)?;
    let client_tls = resolve_client_tls_material_for_agent(Some(&args.agent_id))?;
    let local_name = slim_name(&args.agent_id);

    let peer_names: Vec<Arc<Name>> = args
        .peer_agent_ids
        .split(',')
        .map(str::trim)
        .filter(|id| !id.is_empty())
        .map(|id| parse_name(&slim_name(id)).map(Arc::new))
        .collect::<Result<_, _>>()?;
    if peer_names.is_empty() {
        return Err("--peer-agent-ids must list at least one peer".to_string());
    }

    let service = Service::new(format!("shadictl-a2a-collab-{}", std::process::id()));
    let connection_id = service
        .connect(build_client_config_for_endpoint(&endpoint, &client_tls))
        .map_err(format_slim_error)?;
    let local_name_ref = Arc::new(parse_name(&local_name)?);
    let app = shadi_identity::create_app(&service, local_name_ref.clone(), &auth)
        .map_err(format_slim_error)?;
    app.subscribe(local_name_ref.clone(), Some(connection_id))
        .map_err(format_slim_error)?;

    let runtime = TokioRuntimeBuilder::new_current_thread()
        .enable_all()
        .build()
        .map_err(|err| format!("failed to create tokio runtime: {}", err))?;

    let received: Arc<std::sync::Mutex<Vec<String>>> = Arc::new(std::sync::Mutex::new(Vec::new()));
    let received_for_handler = received.clone();

    let mut session = SessionContext::new(&args.agent_id, "shadictl-a2a-collaborate");
    session.verified = true;
    let verifier: Arc<dyn AgentVerifier> = Arc::new(VerifiedSessionVerifier);

    // Both the listener's Server and the initiator's group Channel capture
    // `tokio::runtime::Handle::current()` at construction, so build them inside
    // this runtime's context (same fix as the P2P a2a-send path).
    let (server, channel) = {
        let _enter = runtime.enter();
        let server = Arc::new(Server::new_with_shared_rx_and_connection(
            app.inner(),
            app.name().as_slim_name(),
            None,
            app.notification_receiver(),
            Some(runtime.handle().clone()),
        ));
        shadi_a2a::register_collaborate(server.as_ref(), move |message| {
            // `register_collaborate`'s listener path has no per-message sender
            // attribution today (see a2a-slimrpc's Collaborate limitations) — only
            // `collaborate()`'s own client-observer stream gets `slim-src` via
            // MulticastItem. The message content itself states the sender, so
            // just surface the text.
            let text = readable_message_text(&message);
            received_for_handler.lock().unwrap().push(text);
        });

        let channel = A2AGroupChannelBuilder::new(app.clone(), peer_names.clone(), verifier, session)
            .connection_id(connection_id)
            .build()
            .map_err(|err| format!("failed to build group channel: {err}"))?;

        (server, channel)
    };

    let intro = Message::new(Role::Agent, vec![Part::text(args.message.clone())]);
    let outbound = futures::stream::once(async move { intro });
    let received_for_replies = received.clone();

    runtime.block_on(async {
        let server_task = {
            let server = server.clone();
            tokio::spawn(async move {
                let _ = server.serve().await;
            })
        };
        // Give every group member's own listener time to finish subscribing
        // before broadcasting — `Channel::new_group`'s self-invite dance has no
        // built-in retry, so an invite that lands before the peer is ready can be
        // missed. Every member typically starts around the same time (as in the
        // demo), so this is a generous, shared settle window, not per-peer.
        tokio::time::sleep(Duration::from_secs(2)).await;

        // Unlike the listener above, this is *our own* initiated session, so
        // replies here (if any participant isn't listen-only) carry proper
        // `slim-src` attribution via MulticastItem — surface it.
        let replies = channel
            .collaborate(outbound, Some(Duration::from_secs(args.timeout_seconds)))
            .map_err(|err| format!("failed to broadcast message: {err}"))?;
        futures::pin_mut!(replies);
        while let Some(reply) = futures::StreamExt::next(&mut replies).await {
            if let Ok(message) = reply {
                let sender = message
                    .metadata
                    .as_ref()
                    .and_then(|metadata| metadata.get(SLIM_SRC_METADATA_KEY))
                    .and_then(|value| value.as_str())
                    .unwrap_or("<unknown>")
                    .to_string();
                let text = readable_message_text(&message);
                received_for_replies
                    .lock()
                    .unwrap()
                    .push(format!("{sender}: {text}"));
            }
        }

        // Keep listening for other members' independently-initiated broadcasts
        // for the rest of the timeout window.
        tokio::time::sleep(Duration::from_secs(args.timeout_seconds)).await;

        server.shutdown().await;
        let _ = server_task.await;
        Ok::<(), String>(())
    })?;

    let _ = app.unsubscribe(local_name_ref, Some(connection_id));
    let _ = service.disconnect(connection_id);
    let _ = service.shutdown();

    let received = received.lock().unwrap();
    Ok(if received.is_empty() {
        format!(
            "broadcast {:?} to {} peer(s); received nothing within {}s",
            args.message,
            peer_names.len(),
            args.timeout_seconds
        )
    } else {
        format!(
            "broadcast {:?} to {} peer(s); received:\n  {}",
            args.message,
            peer_names.len(),
            received.join("\n  ")
        )
    })
}

fn resolve_endpoint(override_value: Option<&str>) -> String {
    override_value
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(str::to_owned)
        .or_else(|| {
            std::env::var("SLIM_ENDPOINT")
                .ok()
                .filter(|value| !value.trim().is_empty())
        })
        .unwrap_or_else(|| DEFAULT_SLIM_ENDPOINT.to_string())
}

fn next_value<'a>(args: &'a [&'a str], index: &mut usize, usage: &str) -> Result<&'a str, String> {
    if *index + 1 >= args.len() {
        return Err(usage.to_string());
    }
    *index += 1;
    Ok(args[*index])
}

fn collect_message_value(args: &[&str], start: usize) -> Result<(String, usize), String> {
    if start >= args.len() {
        return Err(SHELL_A2A_SEND_USAGE.to_string());
    }

    let mut end = start;
    while end < args.len() && !is_send_flag(args[end]) {
        end += 1;
    }

    if end == start {
        return Err(SHELL_A2A_SEND_USAGE.to_string());
    }

    Ok((args[start..end].join(" "), end))
}

fn is_send_flag(token: &str) -> bool {
    matches!(
        token,
        "--endpoint"
            | "--agent-id"
            | "--peer-agent-id"
            | "--destination"
            | "--message"
            | "--stream"
            | "--timeout"
            | "--timeout-seconds"
            | "--session-id"
    )
}

fn slim_name(agent_id: &str) -> String {
    if agent_id.contains('/') {
        agent_id.to_string()
    } else {
        format!("agntcy/shadi/{}", agent_id)
    }
}

fn readable_message_text(message: &Message) -> String {
    let text = message
        .parts
        .iter()
        .filter_map(Part::as_text)
        .collect::<Vec<_>>()
        .join(" ");
    if text.is_empty() {
        "(no text parts)".to_string()
    } else {
        text
    }
}

fn describe_a2a_response(response: &SendMessageResponse) -> String {
    match response {
        SendMessageResponse::Message(message) => {
            format!("message {:?}", readable_message_text(message))
        }
        SendMessageResponse::Task(task) => {
            let detail = task
                .status
                .message
                .as_ref()
                .map(readable_message_text)
                .unwrap_or_else(|| format!("state {:?}", task.status.state));
            format!("task {} ({})", task.id, detail)
        }
    }
}

fn describe_a2a_stream(events: &[Result<StreamResponse, A2AError>]) -> String {
    let mut descriptions = Vec::new();
    for event in events {
        match event {
            Ok(StreamResponse::StatusUpdate(update)) => {
                descriptions.push(format!("status {:?}", update.status.state));
            }
            Ok(StreamResponse::Task(task)) => descriptions.push(format!(
                "task {} ({})",
                task.id,
                task.status
                    .message
                    .as_ref()
                    .map(readable_message_text)
                    .unwrap_or_else(|| format!("state {:?}", task.status.state))
            )),
            Ok(StreamResponse::Message(message)) => {
                descriptions.push(format!("message {:?}", readable_message_text(message)));
            }
            Ok(StreamResponse::ArtifactUpdate(update)) => {
                descriptions.push(format!("artifact {}", update.task_id));
            }
            Err(error) => descriptions.push(format!("error {}", error)),
        }
    }
    format!("stream [{}]", descriptions.join(", "))
}

fn a2a_agent_card(agent_name: &str, target: &str) -> AgentCard {
    AgentCard {
        name: format!("SHADI CLI A2A Peer ({})", agent_name),
        description: "shadictl A2A peer exposed over SLIMRPC".to_string(),
        version: env!("CARGO_PKG_VERSION").to_string(),
        supported_interfaces: vec![AgentInterface::new(
            target.to_string(),
            TRANSPORT_PROTOCOL_SLIMRPC,
        )],
        capabilities: AgentCapabilities {
            streaming: Some(true),
            push_notifications: Some(false),
            extensions: None,
            extended_agent_card: Some(true),
        },
        default_input_modes: vec!["text/plain".to_string()],
        default_output_modes: vec!["text/plain".to_string()],
        skills: Vec::new(),
        provider: None,
        documentation_url: None,
        icon_url: None,
        security_schemes: None,
        security_requirements: None,
        signatures: None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use agent_secrets::SecretError;
    use futures::StreamExt;
    use std::collections::HashMap;

    fn send_request(text: &str) -> SendMessageRequest {
        SendMessageRequest {
            message: Message::new(Role::User, vec![Part::text(text)]),
            configuration: None,
            metadata: None,
            tenant: None,
        }
    }

    fn executor_context(message: Option<Message>) -> a2a_server::ExecutorContext {
        a2a_server::ExecutorContext {
            message,
            task_id: "task-1".to_string(),
            stored_task: None,
            context_id: "context-1".to_string(),
            metadata: None,
            user: None,
            service_params: HashMap::new(),
            tenant: None,
        }
    }

    #[test]
    fn slim_name_canonicalizes_bare_agent_ids() {
        assert_eq!(slim_name("avatar"), "agntcy/shadi/avatar");
        assert_eq!(slim_name("agntcy/other/agent"), "agntcy/other/agent");
    }

    #[test]
    fn describe_a2a_response_formats_message_and_task_variants() {
        let message = SendMessageResponse::Message(Message::new(
            Role::Agent,
            vec![Part::text("hello from shadictl")],
        ));
        assert_eq!(
            describe_a2a_response(&message),
            "message \"hello from shadictl\""
        );

        let task = SendMessageResponse::Task(Task {
            id: "task-1".to_string(),
            context_id: "context-1".to_string(),
            status: TaskStatus {
                state: TaskState::Completed,
                message: Some(Message::new(Role::Agent, vec![Part::text("done")])),
                timestamp: None,
            },
            artifacts: None,
            history: None,
            metadata: None,
        });
        assert_eq!(describe_a2a_response(&task), "task task-1 (done)");
    }

    #[test]
    fn describe_a2a_stream_formats_status_and_task_events() {
        let events = vec![
            Ok(StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
                task_id: "task-1".to_string(),
                context_id: "context-1".to_string(),
                status: TaskStatus {
                    state: TaskState::Working,
                    message: None,
                    timestamp: None,
                },
                metadata: None,
            })),
            Ok(StreamResponse::Task(Task {
                id: "task-1".to_string(),
                context_id: "context-1".to_string(),
                status: TaskStatus {
                    state: TaskState::Completed,
                    message: Some(Message::new(Role::Agent, vec![Part::text("stream done")])),
                    timestamp: None,
                },
                artifacts: None,
                history: None,
                metadata: None,
            })),
        ];

        assert_eq!(
            describe_a2a_stream(&events),
            "stream [status Working, task task-1 (stream done)]"
        );
    }

    #[test]
    fn describe_a2a_stream_formats_message_artifact_and_error_variants() {
        let events = vec![
            Ok(StreamResponse::Message(Message::new(
                Role::Agent,
                vec![Part::text("direct reply")],
            ))),
            Ok(StreamResponse::ArtifactUpdate(a2a::event::TaskArtifactUpdateEvent {
                task_id: "task-1".to_string(),
                context_id: "context-1".to_string(),
                artifact: Artifact {
                    artifact_id: "artifact-1".to_string(),
                    name: None,
                    description: None,
                    parts: vec![Part::text("artifact body")],
                    metadata: None,
                    extensions: None,
                },
                append: Some(true),
                last_chunk: Some(false),
                metadata: None,
            })),
            Err(A2AError::internal("stream failed")),
        ];

        assert_eq!(
            describe_a2a_stream(&events),
            "stream [message \"direct reply\", artifact task-1, error stream failed]"
        );
    }

    #[test]
    fn readable_message_text_returns_placeholder_when_text_is_missing() {
        let message = Message::new(Role::User, vec![]);

        assert_eq!(readable_message_text(&message), "(no text parts)");
    }

    #[test]
    fn resolve_endpoint_prefers_override_then_env_then_default() {
        let _guard = crate::lock_test_env();
        std::env::set_var("SLIM_ENDPOINT", "192.0.2.8:4555");

        assert_eq!(
            resolve_endpoint(Some("198.51.100.5:4777")),
            "198.51.100.5:4777"
        );
        assert_eq!(resolve_endpoint(Some("   ")), "192.0.2.8:4555");

        std::env::remove_var("SLIM_ENDPOINT");
        assert_eq!(resolve_endpoint(None), DEFAULT_SLIM_ENDPOINT);
    }

    #[test]
    fn collect_message_value_stops_at_next_flag() {
        let (message, next_index) = collect_message_value(
            &["hello", "from", "avatar", "--stream", "--timeout", "7"],
            0,
        )
        .expect("collect message");

        assert_eq!(message, "hello from avatar");
        assert_eq!(next_index, 3);
    }

    #[test]
    fn a2a_agent_card_exposes_slimrpc_interface() {
        let card = a2a_agent_card("agntcy/shadi/secops-a", "slimrpc://agntcy/shadi/secops-a");

        assert_eq!(card.name, "SHADI CLI A2A Peer (agntcy/shadi/secops-a)");
        assert_eq!(card.supported_interfaces.len(), 1);
        assert_eq!(
            card.supported_interfaces[0].protocol_binding,
            TRANSPORT_PROTOCOL_SLIMRPC
        );
        assert_eq!(card.capabilities.streaming, Some(true));
    }

    #[test]
    fn slim_a2a_executor_emits_status_and_completed_task() {
        let executor = SlimA2AExecutor {
            agent_name: "agntcy/shadi/secops-a".to_string(),
        };
        let events = futures::executor::block_on(async {
            executor
                .execute(executor_context(Some(Message::new(
                    Role::User,
                    vec![Part::text("hello executor")],
                ))))
                .collect::<Vec<_>>()
                .await
        });

        assert_eq!(events.len(), 2);
        match &events[0] {
            Ok(StreamResponse::StatusUpdate(update)) => {
                assert_eq!(update.status.state, TaskState::Working);
            }
            other => panic!("unexpected first executor event: {other:?}"),
        }
        match &events[1] {
            Ok(StreamResponse::Task(task)) => {
                assert_eq!(task.id, "task-1");
                assert_eq!(
                    readable_message_text(task.status.message.as_ref().expect("task message")),
                    "echo:agntcy/shadi/secops-a:hello executor"
                );
                assert_eq!(task.history.as_ref().expect("history").len(), 1);
            }
            other => panic!("unexpected second executor event: {other:?}"),
        }
    }

    #[test]
    fn slim_a2a_executor_cancel_marks_task_canceled() {
        let executor = SlimA2AExecutor {
            agent_name: "agntcy/shadi/secops-a".to_string(),
        };
        let events = futures::executor::block_on(async {
            executor
                .cancel(executor_context(None))
                .collect::<Vec<_>>()
                .await
        });

        assert_eq!(events.len(), 1);
        match &events[0] {
            Ok(StreamResponse::Task(task)) => {
                assert_eq!(task.id, "task-1");
                assert_eq!(task.status.state, TaskState::Canceled);
            }
            other => panic!("unexpected cancel event: {other:?}"),
        }
    }

    #[tokio::test]
    async fn slim_a2a_handler_routes_requests_and_push_configs() {
        let request_seen = Arc::new(Notify::new());
        let handler = SlimA2AHandler::new(
            "agntcy/shadi/secops-a".to_string(),
            "slimrpc://agntcy/shadi/secops-a".to_string(),
            request_seen.clone(),
        );
        let params: A2AServiceParams = HashMap::new();

        let send_wait = request_seen.notified();
        let response = handler
            .send_message(&params, send_request("hello handler"))
            .await
            .expect("send message through handler");
        tokio::time::timeout(Duration::from_secs(1), send_wait)
            .await
            .expect("send_message should notify waiters");

        let task = match response {
            SendMessageResponse::Task(task) => task,
            other => panic!("unexpected send_message response: {other:?}"),
        };

        let send_stream_wait = request_seen.notified();
        let stream_events = handler
            .send_streaming_message(&params, send_request("hello stream"))
            .await
            .expect("send streaming message")
            .collect::<Vec<_>>()
            .await;
        tokio::time::timeout(Duration::from_secs(1), send_stream_wait)
            .await
            .expect("send_streaming_message should notify waiters");
        assert_eq!(stream_events.len(), 2);

        let fetched = handler
            .get_task(
                &params,
                GetTaskRequest {
                    id: task.id.clone(),
                    history_length: Some(1),
                    tenant: None,
                },
            )
            .await
            .expect("get task");
        assert_eq!(fetched.id, task.id);

        let listed = handler
            .list_tasks(
                &params,
                ListTasksRequest {
                    context_id: Some(task.context_id.clone()),
                    status: None,
                    page_size: Some(10),
                    page_token: None,
                    history_length: Some(1),
                    status_timestamp_after: None,
                    include_artifacts: Some(false),
                    tenant: None,
                },
            )
            .await
            .expect("list tasks");
        assert!(listed.tasks.iter().any(|entry| entry.id == task.id));

        let push_config_err = handler
            .create_push_config(
                &params,
                TaskPushNotificationConfig {
                    task_id: task.id.clone(),
                    tenant: None,
                    url: "https://example.invalid/hook".to_string(),
                    id: Some("cfg-1".to_string()),
                    token: None,
                    authentication: None,
                },
            )
            .await
            .expect_err("push config should not be supported");
        assert!(push_config_err.message.contains("not supported"));

        let fetched_push_err = handler
            .get_push_config(
                &params,
                GetTaskPushNotificationConfigRequest {
                    task_id: task.id.clone(),
                    id: "cfg-1".to_string(),
                    tenant: None,
                },
            )
            .await
            .expect_err("get push config should not be supported");
        assert!(fetched_push_err.message.contains("not supported"));

        let listed_push_err = handler
            .list_push_configs(
                &params,
                ListTaskPushNotificationConfigsRequest {
                    task_id: task.id.clone(),
                    page_size: Some(10),
                    page_token: None,
                    tenant: None,
                },
            )
            .await
            .expect_err("list push configs should not be supported");
        assert!(listed_push_err.message.contains("not supported"));

        let delete_push_err = handler
            .delete_push_config(
                &params,
                DeleteTaskPushNotificationConfigRequest {
                    task_id: task.id.clone(),
                    id: "cfg-1".to_string(),
                    tenant: None,
                },
            )
            .await
            .expect_err("delete push config should not be supported");
        assert!(delete_push_err.message.contains("not supported"));

        let card = handler
            .get_extended_agent_card(&params, GetExtendedAgentCardRequest { tenant: None })
            .await
            .expect("extended agent card");
        assert_eq!(card.supported_interfaces[0].url, "slimrpc://agntcy/shadi/secops-a");
    }

    #[test]
    fn verified_session_verifier_requires_verified_session() {
        let verifier = VerifiedSessionVerifier;
        let unverified = SessionContext::new("avatar", "session-1");
        let err = verifier.verify(&unverified).unwrap_err();
        assert!(matches!(err, SecretError::NotAuthorized));

        let mut verified = SessionContext::new("avatar", "session-2");
        verified.verified = true;
        verifier.verify(&verified).expect("verified session should pass");
    }

    #[test]
    fn parse_shell_a2a_send_args_supports_multiword_message_and_flags() {
        let parsed = parse_shell_a2a_send_args(&[
            "--agent-id",
            "avatar",
            "--peer-agent-id",
            "secops-a",
            "--message",
            "hello",
            "from",
            "shell",
            "--stream",
            "--timeout",
            "7",
            "--session-id",
            "shell-session",
        ])
        .expect("parse shell args");

        assert_eq!(parsed.agent_id, "avatar");
        assert_eq!(parsed.peer_agent_id, "secops-a");
        assert_eq!(parsed.message, "hello from shell");
        assert!(parsed.stream);
        assert_eq!(parsed.timeout_seconds, 7);
        assert_eq!(parsed.session_id, "shell-session");
    }

    #[test]
    fn parse_shell_a2a_echo_peer_args_supports_overrides() {
        let parsed = parse_shell_a2a_echo_peer_args(&[
            "--endpoint",
            "127.0.0.1:48555",
            "--agent-id",
            "secops-a",
            "--listen-timeout",
            "12",
            "--ready-file",
            "/tmp/ready.flag",
            "--start-local-node",
        ])
        .expect("parse shell peer args");

        assert_eq!(parsed.endpoint.as_deref(), Some("127.0.0.1:48555"));
        assert_eq!(parsed.agent_id, "secops-a");
        assert_eq!(parsed.listen_timeout_seconds, 12);
        assert_eq!(parsed.ready_file, Some(PathBuf::from("/tmp/ready.flag")));
        assert!(parsed.start_local_node);
    }

    #[test]
    fn parse_shell_a2a_send_args_rejects_missing_message_value() {
        let err = parse_shell_a2a_send_args(&["--message"]).unwrap_err();
        assert_eq!(err, SHELL_A2A_SEND_USAGE);
    }

    #[test]
    fn parse_shell_a2a_echo_peer_args_rejects_missing_value() {
        let err = parse_shell_a2a_echo_peer_args(&["--ready-file"]).unwrap_err();
        assert_eq!(err, SHELL_A2A_ECHO_PEER_USAGE);
    }
}