magi-code 0.101.0

//! Bounded, child-owned reply routing. No worker or event-loop locks span execution.
use super::{MAX_ARGUMENT_BYTES, MAX_PENDING, MAX_RESULT_BYTES, Registry, bounded_depth};
use crate::{
    cancellation::AgentCancellation,
    tools::{ToolResult, ToolResultDisplay},
};
use crossbeam_channel::{Sender, bounded};
use serde::Deserialize;
use serde_json::{Value, json};
use std::{
    sync::Arc,
    time::{Duration, Instant},
};

type Emit = dyn Fn(&'static str, Value) -> anyhow::Result<()> + Send + Sync;

#[derive(Clone, PartialEq, Eq, serde::Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct Identity {
    instance_id: String,
    session_id: String,
    run_id: String,
    call_id: String,
    task_id: Option<String>,
    subagent_id: Option<String>,
}

// Explicit fields (rather than flatten) retain deny_unknown_fields/duplicate checks.
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct ResultParams {
    instance_id: String,
    session_id: String,
    run_id: String,
    call_id: String,
    task_id: Value,
    subagent_id: Value,
    status: String,
    text: String,
    #[serde(default, deserialize_with = "present_structured")]
    structured: Option<Box<serde_json::value::RawValue>>,
}

fn present_structured<'de, D: serde::Deserializer<'de>>(
    deserializer: D,
) -> Result<Option<Box<serde_json::value::RawValue>>, D::Error> {
    Box::<serde_json::value::RawValue>::deserialize(deserializer).map(Some)
}

pub(super) struct Pending {
    identity: Identity,
    deadline: Instant,
    cancellation: AgentCancellation,
    reply: Sender<Result<ToolResult, &'static str>>,
    name: String,
}

pub(crate) struct ClientToolRun {
    pub(crate) registry: Arc<Registry>,
    instance_id: String,
    session_id: String,
    run_id: String,
    emit: Arc<Emit>,
    shutting_down: std::sync::atomic::AtomicBool,
}
impl ClientToolRun {
    pub(crate) fn new(
        registry: Arc<Registry>,
        instance_id: String,
        session_id: String,
        run_id: String,
        emit: impl Fn(&'static str, Value) -> anyhow::Result<()> + Send + Sync + 'static,
    ) -> Self {
        Self {
            registry,
            instance_id,
            session_id,
            run_id,
            emit: Arc::new(emit),
            shutting_down: std::sync::atomic::AtomicBool::new(false),
        }
    }

    pub(crate) fn shutdown(&self) {
        self.shutting_down
            .store(true, std::sync::atomic::Ordering::SeqCst);
    }

    pub(crate) fn call(
        &self,
        name: &str,
        arguments: Value,
        cancellation: &AgentCancellation,
        task: Option<&(String, String)>,
    ) -> anyhow::Result<ToolResult> {
        cancellation.check()?;
        let tool = self
            .registry
            .tools
            .iter()
            .find(|t| t.name == name)
            .ok_or_else(|| anyhow::anyhow!("client_tool_unavailable"))?;
        anyhow::ensure!(
            arguments.is_object() && serde_json::to_vec(&arguments)?.len() <= MAX_ARGUMENT_BYTES,
            "client_tool_payload_limit"
        );
        let identity = Identity {
            instance_id: self.instance_id.clone(),
            session_id: self.session_id.clone(),
            run_id: self.run_id.clone(),
            call_id: uuid::Uuid::new_v4().to_string(),
            task_id: task.map(|t| t.0.clone()),
            subagent_id: task.map(|t| t.1.clone()),
        };
        let mut data = event_identity(&identity);
        data["name"] = json!(name);
        data["arguments"] = arguments;
        data["timeout_ms"] = json!(tool.timeout_ms);
        // Exact frame shape and worst possible sequence width. LF counts on output.
        let frame = json!({"protocol":"magi-code.local-agent","version":1,"type":"event",
            "instance_id":self.instance_id,"seq":u64::MAX,"event":"client.tool.call","data":data});
        anyhow::ensure!(
            serde_json::to_vec(&frame)?.len() < 65_536 && bounded_depth(&frame, 32),
            "client_tool_payload_limit"
        );
        let deadline = Instant::now() + Duration::from_millis(tool.timeout_ms);
        let (reply, receiver) = bounded(1);
        {
            let mut pending = self
                .registry
                .pending
                .lock()
                .unwrap_or_else(|e| e.into_inner());
            cancellation.check()?;
            anyhow::ensure!(pending.len() < MAX_PENDING, "client_tool_busy");
            pending.insert(
                identity.call_id.clone(),
                Pending {
                    identity: identity.clone(),
                    deadline,
                    cancellation: cancellation.clone(),
                    reply,
                    name: name.to_owned(),
                },
            );
        }
        let _pending_guard = PendingGuard {
            run: self,
            identity: &identity,
        };
        let result = (|| {
            (self.emit)("client.tool.call", data)?;
            loop {
                // Result acceptance and expiry/cancellation share the pending-map lock.
                // Poll the channel first: an already accepted completion won that race.
                if let Ok(result) = receiver.try_recv() {
                    return result.map_err(anyhow::Error::msg);
                }
                if cancellation.is_canceled() || Instant::now() >= deadline {
                    let removed = self
                        .registry
                        .pending
                        .lock()
                        .unwrap_or_else(|e| e.into_inner())
                        .remove(&identity.call_id);
                    if removed.is_some() {
                        let reason = if cancellation.is_canceled() {
                            if self.shutting_down.load(std::sync::atomic::Ordering::SeqCst) {
                                "shutdown"
                            } else {
                                "cancelled"
                            }
                        } else {
                            "timeout"
                        };
                        let mut data = event_identity(&identity);
                        data["reason"] = json!(reason);
                        (self.emit)("client.tool.cancel", data)?;
                        cancellation.check()?;
                        anyhow::bail!("client_tool_timeout");
                    }
                    // Another thread accepted a result and sent under the same lock.
                    return receiver
                        .recv()
                        .map_err(|_| anyhow::anyhow!("client_tool_unavailable"))?
                        .map_err(anyhow::Error::msg);
                }
                match receiver.recv_timeout(
                    Duration::from_millis(20)
                        .min(deadline.saturating_duration_since(Instant::now())),
                ) {
                    Ok(result) => return result.map_err(anyhow::Error::msg),
                    Err(crossbeam_channel::RecvTimeoutError::Timeout) => {}
                    Err(_) => anyhow::bail!("client_tool_unavailable"),
                }
            }
        })();
        // Guard removes interrupted/panicking calls; completed calls already left the map.
        result
    }
}

struct PendingGuard<'a> {
    run: &'a ClientToolRun,
    identity: &'a Identity,
}
impl Drop for PendingGuard<'_> {
    fn drop(&mut self) {
        let removed = self
            .run
            .registry
            .pending
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .remove(&self.identity.call_id);
        if removed.is_some() {
            let mut data = event_identity(self.identity);
            data["reason"] = json!("run_finished");
            let _ = (self.run.emit)("client.tool.cancel", data);
        }
    }
}

fn event_identity(identity: &Identity) -> Value {
    json!({"session_id":identity.session_id,"run_id":identity.run_id,"call_id":identity.call_id,
        "task_id":identity.task_id,"subagent_id":identity.subagent_id})
}
fn optional_id(value: Value) -> Result<Option<String>, &'static str> {
    match value {
        Value::Null => Ok(None),
        Value::String(s)
            if !s.is_empty()
                && s.len() <= 1024
                && s.is_ascii()
                && !s.chars().any(char::is_control) =>
        {
            Ok(Some(s))
        }
        _ => Err("invalid_params"),
    }
}
impl Registry {
    pub(crate) fn answer(&self, params: ResultParams) -> Result<(), &'static str> {
        let structured = params
            .structured
            .as_ref()
            .map(|raw| serde_json::from_str::<Value>(raw.get()).map_err(|_| "invalid_params"))
            .transpose()?;
        if !matches!(params.status.as_str(), "success" | "error")
            || params
                .text
                .len()
                .saturating_add(structured.as_ref().map_or(0, |v| v.to_string().len()))
                > MAX_RESULT_BYTES
        {
            return Err("invalid_params");
        }
        let identity = Identity {
            instance_id: params.instance_id,
            session_id: params.session_id,
            run_id: params.run_id,
            call_id: params.call_id,
            task_id: optional_id(params.task_id)?,
            subagent_id: optional_id(params.subagent_id)?,
        };
        let mut pending = self.pending.lock().unwrap_or_else(|e| e.into_inner());
        let call = pending
            .get(&identity.call_id)
            .ok_or("stale_client_tool_call")?;
        if call.identity != identity
            || Instant::now() >= call.deadline
            || call.cancellation.is_canceled()
        {
            return Err("stale_client_tool_call");
        }
        let text = crate::output::sanitize_display_text(&params.text);
        // Protect each JSON string without making redaction invalidate JSON syntax.
        let mut structured = structured;
        if let Some(value) = &mut structured {
            protect_json(value)?;
        }
        let mut content = text.clone();
        if let Some(value) = &structured {
            content.push_str("\nstructured: ");
            content.push_str(&value.to_string());
        }
        let mut metadata = json!({});
        if let Some(value) = structured {
            metadata["structured_result"] = json!({"text":text,"structured":value});
        }
        let result = ToolResult {
            tool_name: call.name.clone(),
            success: params.status == "success",
            content,
            metadata,
            display: ToolResultDisplay::default(),
        };
        // Redaction can expand short secrets; never silently truncate structured data.
        if result.content.len() + result.metadata.to_string().len() > 2 * MAX_RESULT_BYTES {
            return Err("invalid_params");
        }
        if Instant::now() >= call.deadline || call.cancellation.is_canceled() {
            return Err("stale_client_tool_call");
        }
        let call = pending
            .remove(&identity.call_id)
            .expect("matched pending call");
        call.reply
            .try_send(Ok(result))
            .map_err(|_| "stale_client_tool_call")
    }
}
fn protect_json(value: &mut Value) -> Result<(), &'static str> {
    match value {
        Value::String(text) => *text = crate::output::sanitize_display_text(text),
        Value::Array(items) => {
            for item in items {
                protect_json(item)?;
            }
        }
        Value::Object(items) => {
            let old = std::mem::take(items);
            for (key, mut value) in old {
                if crate::output::is_credential_like_key(&key) {
                    value = json!("[REDACTED]");
                } else {
                    protect_json(&mut value)?;
                }
                if items
                    .insert(crate::output::sanitize_display_text(&key), value)
                    .is_some()
                {
                    return Err("invalid_params");
                }
            }
        }
        _ => {}
    }
    Ok(())
}