magi-code 0.80.1

//! Application-executed callbacks for the negotiated application resource profile.
#[path = "application_source.rs"]
mod source;
pub(super) use source::{PreparedSource, prepare_source};
#[path = "application_evidence.rs"]
mod evidence;
#[path = "application_registry.rs"]
mod registry;
pub(super) use evidence::EVIDENCE_METHOD;
#[cfg(test)]
#[path = "application_tests.rs"]
mod tests;
use super::protocol::ServiceErrorCode as Code;
use crate::{
    output::ToolDispatchContext,
    tools::{ToolResult, ToolResultDisplay, application::ApplicationTools},
};
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
pub(super) use registry::{CapturedTurn, RESOURCE_METHODS};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use std::{
    collections::{BTreeMap, HashMap, VecDeque},
    io::Read,
    sync::{Arc, Mutex},
    time::{Duration, Instant},
};
use subtle::ConstantTimeEq;

pub(super) const PROFILE: &str = "application_resources_v1";
const MAX_CALLS: usize = 32;
pub(super) const METHODS: [&str; 7] = [
    "executor.create",
    "executor.confirm",
    "executor.resume",
    "executor.rotate",
    "executor.rotate_confirm",
    "executor.release",
    "turn.application_tool_result",
];

#[derive(Clone, Serialize)]
pub(super) struct InlineTurn {
    pub executor_id: String,
    pub executor_generation: u64,
    pub tool: Arc<Contract>,
    pub skill: Option<Arc<str>>,
    #[serde(skip)]
    pub resource_identity: Option<ResourceIdentity>,
}
#[derive(Clone)]
pub(super) struct ResourceIdentity {
    pub resource_id: String,
    pub revision_id: String,
    pub revision_sha256: String,
}
impl ResourceIdentity {
    fn add_to(&self, value: &mut Value) {
        value["resource_id"] = json!(self.resource_id);
        value["revision_id"] = json!(self.revision_id);
        value["revision_sha256"] = json!(self.revision_sha256);
    }
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub(super) struct Contract {
    name: String,
    description: String,
    input_schema: Schema,
    output_schema: Schema,
    #[serde(default = "default_timeout")]
    timeout_ms: u64,
}
fn default_timeout() -> u64 {
    120_000
}

#[derive(Clone, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
struct Schema {
    #[serde(rename = "type")]
    kind: String,
    properties: BTreeMap<String, Property>,
    required: Vec<String>,
    #[serde(rename = "additionalProperties")]
    additional: bool,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(tag = "type", deny_unknown_fields)]
enum Property {
    #[serde(rename = "string")]
    String {
        #[serde(rename = "maxLength")]
        max_length: u64,
    },
    #[serde(rename = "boolean")]
    Boolean,
    #[serde(rename = "integer")]
    Integer { minimum: i64, maximum: i64 },
}
impl Schema {
    fn validate(&self) -> Result<(), Code> {
        let names: std::collections::BTreeSet<_> = self.required.iter().collect();
        if self.kind != "object"
            || self.additional
            || self.properties.len() > 32
            || names.len() != self.required.len()
            || names.len() != self.properties.len()
            || self
                .properties
                .keys()
                .any(|name| !names.contains(name) || name.len() > 48)
            || self.properties.values().any(|property| match property {
                Property::String { max_length } => !(1..=4096).contains(max_length),
                Property::Boolean => false,
                Property::Integer { minimum, maximum } => {
                    minimum > maximum
                        || *minimum < -9_007_199_254_740_991
                        || *maximum > 9_007_199_254_740_991
                }
            })
        {
            return Err(Code::InvalidPayload);
        }
        Ok(())
    }
    fn accepts(&self, value: &Value) -> bool {
        let Some(object) = value.as_object() else {
            return false;
        };
        object.len() == self.properties.len()
            && value.to_string().len() <= 16_384
            && self.properties.iter().all(|(name, property)| {
                let value = &value[name];
                match property {
                    Property::String { max_length } => value
                        .as_str()
                        .is_some_and(|s| s.chars().count() as u64 <= *max_length),
                    Property::Boolean => value.is_boolean(),
                    Property::Integer { minimum, maximum } => value
                        .as_i64()
                        .is_some_and(|n| (*minimum..=*maximum).contains(&n)),
                }
            })
    }
}

struct Executor {
    verifier: [u8; 32],
    pending_rotation: Option<PendingRotation>,
    connection: Option<String>,
    generation: u64,
    confirmed: bool,
    expires: Instant,
}
struct PendingRotation {
    verifier: [u8; 32],
    expires: Instant,
}
struct Call {
    executor: String,
    session: String,
    turn: String,
    provider_call_id: String,
    revision: String,
    resource_identity: Option<ResourceIdentity>,
    deadline: Instant,
    dispatched: bool,
    outcome: Option<ToolResult>,
    output_schema: Schema,
    cancellation: crate::cancellation::AgentCancellation,
}
pub(super) struct CallbackEvent {
    pub connection: String,
    pub session: String,
    pub name: &'static str,
    pub payload: Value,
}
#[derive(Default)]
pub(super) struct ApplicationState {
    executors: HashMap<String, Executor>,
    calls: HashMap<String, Call>,
    events: VecDeque<CallbackEvent>,
    turns: HashMap<String, String>,
    pub registry: registry::Registry,
    evidence: HashMap<String, evidence::TurnEvidence>,
}
pub(super) type SharedApplicationState = Arc<Mutex<ApplicationState>>;

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Proof {
    executor_id: String,
    resume_secret: String,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Attachment {
    executor_id: String,
    executor_generation: u64,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct RotationConfirmation {
    executor_id: String,
    executor_generation: u64,
    resume_secret: String,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ResultSubmission {
    executor_id: String,
    executor_generation: u64,
    turn_id: String,
    call_id: String,
    revision_sha256: String,
    resource_id: Option<String>,
    revision_id: Option<String>,
    outcome: Outcome,
}
#[derive(Deserialize)]
#[serde(tag = "status", deny_unknown_fields)]
enum Outcome {
    #[serde(rename = "success")]
    Success { value: Value },
    #[serde(rename = "error")]
    Error { code: ErrorCode, message: String },
}
#[derive(Deserialize)]
#[serde(rename_all = "snake_case")]
enum ErrorCode {
    ExecutionFailed,
    PermissionDenied,
}
fn decode<T: serde::de::DeserializeOwned>(value: &Value) -> Result<T, Code> {
    serde_json::from_value(value.clone()).map_err(|_| Code::InvalidPayload)
}
fn generate_resume_secret() -> Result<String, Code> {
    let mut bytes = [0u8; 32];
    std::fs::File::open("/dev/urandom")
        .and_then(|mut file| file.read_exact(&mut bytes))
        .map_err(|_| Code::InternalError)?;
    Ok(URL_SAFE_NO_PAD.encode(bytes))
}
impl ApplicationState {
    pub fn route(
        &mut self,
        connection: &str,
        method: &str,
        session: Option<&str>,
        payload: &Value,
        now: Instant,
    ) -> Result<Value, Code> {
        self.expire(now);
        match method {
            "executor.create" => {
                if payload != &json!({}) {
                    return Err(Code::InvalidPayload);
                }
                if self.executors.len() >= 256
                    || self
                        .executors
                        .values()
                        .any(|e| e.connection.as_deref() == Some(connection))
                {
                    return Err(Code::LimitExceeded);
                }
                let secret = generate_resume_secret()?;
                let id = uuid::Uuid::new_v4().to_string();
                self.executors.insert(
                    id.clone(),
                    Executor {
                        verifier: Sha256::digest(secret.as_bytes()).into(),
                        pending_rotation: None,
                        connection: Some(connection.into()),
                        generation: 1,
                        confirmed: false,
                        expires: now + Duration::from_secs(60),
                    },
                );
                Ok(json!({"executor_id":id,"executor_generation":1,"resume_secret":secret}))
            }
            "executor.confirm" | "executor.resume" => {
                self.authenticate(connection, method, payload, now)
            }
            "executor.rotate" => self.rotate(connection, payload, now),
            "executor.rotate_confirm" => self.confirm_rotation(connection, payload),
            "executor.release" => {
                let params: Attachment = decode(payload)?;
                self.authorize(connection, &params.executor_id, params.executor_generation)?;
                self.executors.remove(&params.executor_id);
                Ok(json!({"status":"released"}))
            }
            "turn.application_tool_result" => self.submit_result(connection, session, payload, now),
            _ => Err(Code::UnsupportedOperation),
        }
    }
    fn rotate(&mut self, connection: &str, payload: &Value, now: Instant) -> Result<Value, Code> {
        let params: Attachment = decode(payload)?;
        self.authorize(connection, &params.executor_id, params.executor_generation)?;
        let secret = generate_resume_secret()?;
        let executor = self
            .executors
            .get_mut(&params.executor_id)
            .ok_or(Code::StaleConnection)?;
        executor.pending_rotation = Some(PendingRotation {
            verifier: Sha256::digest(secret.as_bytes()).into(),
            expires: now + Duration::from_secs(60),
        });
        Ok(
            json!({"executor_id":params.executor_id,"executor_generation":executor.generation,
            "status":"rotation_pending","pending_timeout_ms":60000,"resume_secret":secret}),
        )
    }
    fn confirm_rotation(&mut self, connection: &str, payload: &Value) -> Result<Value, Code> {
        let params: RotationConfirmation = decode(payload)?;
        self.authorize(connection, &params.executor_id, params.executor_generation)?;
        let executor = self
            .executors
            .get_mut(&params.executor_id)
            .ok_or(Code::StaleConnection)?;
        let verifier: [u8; 32] = Sha256::digest(params.resume_secret.as_bytes()).into();
        let pending = executor
            .pending_rotation
            .as_ref()
            .ok_or(Code::StaleConnection)?;
        if !bool::from(pending.verifier.ct_eq(&verifier)) {
            return Err(Code::StaleConnection);
        }
        executor.verifier = verifier;
        executor.pending_rotation = None;
        Ok(
            json!({"executor_id":params.executor_id,"executor_generation":executor.generation,"status":"rotation_confirmed"}),
        )
    }
    fn authenticate(
        &mut self,
        connection: &str,
        method: &str,
        payload: &Value,
        now: Instant,
    ) -> Result<Value, Code> {
        let proof: Proof = decode(payload)?;
        if self
            .executors
            .iter()
            .any(|(id, e)| id != &proof.executor_id && e.connection.as_deref() == Some(connection))
        {
            return Err(Code::LimitExceeded);
        }
        let executor = self
            .executors
            .get_mut(&proof.executor_id)
            .ok_or(Code::StaleConnection)?;
        let verifier: [u8; 32] = Sha256::digest(proof.resume_secret.as_bytes()).into();
        let current_proof = bool::from(executor.verifier.ct_eq(&verifier));
        let pending_proof = method == "executor.resume"
            && executor
                .pending_rotation
                .as_ref()
                .is_some_and(|pending| bool::from(pending.verifier.ct_eq(&verifier)));
        if !current_proof && !pending_proof {
            return Err(Code::StaleConnection);
        }
        if method == "executor.confirm" {
            if executor.connection.as_deref() != Some(connection) {
                return Err(Code::StaleConnection);
            }
        } else {
            executor.generation = executor
                .generation
                .checked_add(1)
                .filter(|g| *g <= 9_007_199_254_740_991)
                .ok_or(Code::LimitExceeded)?;
            executor.connection = Some(connection.into());
            if pending_proof {
                executor.verifier = verifier;
            }
            executor.pending_rotation = None;
        }
        executor.confirmed = true;
        executor.expires = now + Duration::from_secs(900);
        let pending: Vec<_> = self
            .calls
            .iter()
            .filter(|(_, call)| call.executor == proof.executor_id && call.outcome.is_none())
            .map(|(id, call)| call.pending_evidence(id, now))
            .collect();
        Ok(
            json!({"executor_id":proof.executor_id,"executor_generation":executor.generation,"pending_calls":pending}),
        )
    }
    fn authorize(&self, connection: &str, id: &str, generation: u64) -> Result<(), Code> {
        if self.executors.get(id).is_some_and(|e| {
            e.confirmed && e.connection.as_deref() == Some(connection) && e.generation == generation
        }) {
            Ok(())
        } else {
            Err(Code::StaleConnection)
        }
    }
    fn submit_result(
        &mut self,
        connection: &str,
        session: Option<&str>,
        payload: &Value,
        now: Instant,
    ) -> Result<Value, Code> {
        let params: ResultSubmission = decode(payload)?;
        self.authorize(connection, &params.executor_id, params.executor_generation)?;
        let call = self
            .calls
            .get_mut(&params.call_id)
            .ok_or(Code::UnknownTurn)?;
        let expected_ids = call
            .resource_identity
            .as_ref()
            .map(|identity| (identity.resource_id.as_str(), identity.revision_id.as_str()));
        if params
            .resource_id
            .as_deref()
            .zip(params.revision_id.as_deref())
            != expected_ids
            || params.resource_id.is_some() != params.revision_id.is_some()
        {
            return Err(Code::InvalidPayload);
        }
        if call.executor != params.executor_id
            || Some(call.session.as_str()) != session
            || call.turn != params.turn_id
            || call.revision != params.revision_sha256
        {
            return Err(Code::InvalidPayload);
        }
        if call.outcome.is_some()
            || !call.dispatched
            || now >= call.deadline
            || call.cancellation.is_canceled()
        {
            return Err(Code::UnknownTurn);
        }
        let (success, content) = match params.outcome {
            Outcome::Success { value } => {
                if !call.output_schema.accepts(&value) {
                    return Err(Code::InvalidPayload);
                }
                (true, value.to_string())
            }
            Outcome::Error { code, message } => {
                if message.len() > 512 {
                    return Err(Code::InvalidPayload);
                }
                let label = match code {
                    ErrorCode::ExecutionFailed => "execution_failed",
                    ErrorCode::PermissionDenied => "permission_denied",
                };
                (false, format!("{label}: {message}"))
            }
        };
        call.outcome = Some(ToolResult {
            tool_name: String::new(),
            success,
            content,
            metadata: json!({"application_call_id":params.call_id,"external_effects":"reported"}),
            display: ToolResultDisplay::default(),
        });
        self.update_call_evidence(
            &params.call_id,
            if success { "succeeded" } else { "failed" },
            "reported",
        );
        Ok(json!({"call_id":params.call_id,"status":"accepted"}))
    }
    pub fn disconnect(&mut self, connection: &str, now: Instant) {
        for executor in self.executors.values_mut() {
            if executor.connection.as_deref() == Some(connection) {
                executor.connection = None;
                if executor.confirmed {
                    executor.expires = now + Duration::from_secs(900);
                }
            }
        }
    }
    fn expire(&mut self, now: Instant) {
        self.executors.retain(|id, executor| {
            if executor
                .pending_rotation
                .as_ref()
                .is_some_and(|pending| now >= pending.expires)
            {
                executor.pending_rotation = None;
            }
            (executor.confirmed && executor.connection.is_some())
                || now < executor.expires
                || self.turns.values().any(|executor| executor == id)
        });
    }
    pub fn drain_events(&mut self) -> Vec<CallbackEvent> {
        let now = Instant::now();
        self.expire(now);
        self.events
            .drain(..)
            .filter(|event| {
                if event.name != "turn.application_tool_call" {
                    return true;
                }
                self.executors
                    .get(event.payload["executor_id"].as_str().unwrap_or_default())
                    .is_some_and(|executor| {
                        executor.connection.as_deref() == Some(&event.connection)
                            && event.payload["executor_generation"] == executor.generation
                    })
                    && self
                        .calls
                        .get(event.payload["call_id"].as_str().unwrap_or_default())
                        .is_some_and(|call| {
                            call.outcome.is_none()
                                && now < call.deadline
                                && !call.cancellation.is_canceled()
                        })
            })
            .collect()
    }
    pub fn finish_turn(&mut self, turn: &str) {
        self.calls.retain(|_, call| call.turn != turn);
        self.turns.remove(turn);
    }
}

fn validate_contract(tool: &Contract) -> Result<(), Code> {
    registry::validate_resource_name(&tool.name)?;
    if !(1_000..=600_000).contains(&tool.timeout_ms) || tool.description.len() > 512 {
        return Err(Code::InvalidPayload);
    }
    tool.input_schema.validate()?;
    tool.output_schema.validate()
}
impl InlineTurn {
    pub fn manifest(&self) -> Value {
        let mut manifest = json!({"profile":PROFILE,"executor_id":self.executor_id,"executor_generation":self.executor_generation,
            "name":self.tool.name,"revision_sha256":crate::hex::lower_hex(Sha256::digest(serde_json::to_vec(self).expect("serializable contract")))});
        if let Some(identity) = &self.resource_identity {
            identity.add_to(&mut manifest);
        }
        manifest
    }
    pub fn tools(
        &self,
        state: SharedApplicationState,
        session: String,
        turn: String,
    ) -> ApplicationTools {
        state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .turns
            .insert(turn.clone(), self.executor_id.clone());
        let captured = self.clone();
        ApplicationTools {
            definitions: vec![
                json!({"type":"function","name":self.tool.name,"description":self.tool.description,"parameters":self.tool.input_schema}),
            ],
            manifest: self.manifest(),
            callback: Arc::new(move |name, arguments, context| {
                captured.execute(&state, &session, &turn, name, arguments, context)
            }),
        }
    }
    fn execute(
        &self,
        state: &SharedApplicationState,
        session: &str,
        turn: &str,
        name: &str,
        arguments: Value,
        context: &ToolDispatchContext,
    ) -> ToolResult {
        let failure = |reason: &str, dispatched: bool| ToolResult {
            tool_name: name.into(),
            success: false,
            content: reason.into(),
            metadata: json!({"application_call_id":context.application_call_id,"application_fatal":true,"external_effects":if dispatched {"uncertain"} else {"not_dispatched"}}),
            display: ToolResultDisplay::default(),
        };
        if !self.tool.input_schema.accepts(&arguments) {
            return failure("invalid application tool arguments", false);
        }
        let Some(provider_call_id) = &context.provider_call_id else {
            return failure("missing provider call identity", false);
        };
        if provider_call_id.len() > 128 {
            return failure("invalid provider call identity", false);
        }
        let Some(id) = context.application_call_id.clone() else {
            return failure("missing persisted application call identity", false);
        };
        let revision = self.manifest()["revision_sha256"]
            .as_str()
            .unwrap_or_default()
            .to_owned();
        {
            let mut state = state.lock().unwrap_or_else(|e| e.into_inner());
            if state.calls.len() >= MAX_CALLS
                || state
                    .calls
                    .values()
                    .any(|call| call.turn == turn && call.provider_call_id == *provider_call_id)
            {
                return failure("application call limit or duplicate identity", false);
            }
            let mut evidence = json!({
                "call_id":id,"provider_call_id":provider_call_id,"name":name,
                "executor_id":self.executor_id,"executor_generation":self.executor_generation,
                "revision_sha256":revision,"status":"pending","external_effects":"not_dispatched"
            });
            if let Some(identity) = &self.resource_identity {
                identity.add_to(&mut evidence);
            }
            if state
                .reserve_call_evidence(session, turn, &id, evidence)
                .is_err()
            {
                return failure("application evidence capacity exhausted", false);
            }
            let call = Call {
                resource_identity: self.resource_identity.clone(),
                executor: self.executor_id.clone(),
                session: session.into(),
                turn: turn.into(),
                provider_call_id: provider_call_id.clone(),
                revision: revision.clone(),
                deadline: Instant::now() + Duration::from_secs(60),
                dispatched: false,
                outcome: None,
                output_schema: self.tool.output_schema.clone(),
                cancellation: context.cancellation.clone(),
            };
            if state.reserve_resume_capacity(&id, &call).is_err() {
                state.update_call_evidence(&id, "failed", "not_dispatched");
                return failure("application resume capacity exhausted", false);
            }
            state.calls.insert(id.clone(), call);
        }
        loop {
            {
                let mut state = state.lock().unwrap_or_else(|e| e.into_inner());
                let now = Instant::now();
                let attachment = state
                    .executors
                    .get(&self.executor_id)
                    .filter(|e| e.confirmed)
                    .and_then(|e| e.connection.as_ref().map(|c| (c.clone(), e.generation)));
                // Reserve one cancellation event per admitted call; at most 64 events total.
                let queue_available = state.events.len() < MAX_CALLS;
                let Some(call) = state.calls.get_mut(&id) else {
                    return failure("application call closed", true);
                };
                if let Some(mut result) = call.outcome.clone() {
                    result.tool_name = name.into();
                    return result;
                }
                if context.cancellation.is_canceled() || now >= call.deadline {
                    let reason = if context.cancellation.is_canceled() {
                        "cancelled"
                    } else if call.dispatched {
                        "tool_timeout"
                    } else {
                        "executor_unavailable"
                    };
                    let result = failure(reason, call.dispatched);
                    call.outcome = Some(result.clone());
                    let effects = if call.dispatched {
                        "uncertain"
                    } else {
                        "not_dispatched"
                    };
                    let dispatched = call.dispatched;
                    state.update_call_evidence(
                        &id,
                        if reason == "cancelled" {
                            "cancelled"
                        } else {
                            "failed"
                        },
                        effects,
                    );
                    if dispatched && let Some((connection, _)) = attachment {
                        state.events.push_back(CallbackEvent {
                            connection,
                            session: session.into(),
                            name: "turn.application_tool_cancel",
                            payload: json!({"turn_id":turn,"call_id":id,"reason":reason}),
                        });
                    }
                    return result;
                }
                if !call.dispatched
                    && queue_available
                    && let Some((connection, generation)) = attachment
                {
                    call.dispatched = true;
                    call.deadline = now + Duration::from_millis(self.tool.timeout_ms);
                    state.update_call_evidence(&id, "pending", "uncertain");
                    let mut payload = json!({"turn_id":turn,"call_id":id,"provider_call_id":provider_call_id,"executor_id":self.executor_id,"executor_generation":generation,"revision_sha256":revision,"name":name,"arguments":arguments,"timeout_ms":self.tool.timeout_ms});
                    if let Some(identity) = &self.resource_identity {
                        identity.add_to(&mut payload);
                    }
                    state.events.push_back(CallbackEvent {
                        connection,
                        session: session.into(),
                        name: "turn.application_tool_call",
                        payload,
                    });
                }
            }
            std::thread::sleep(Duration::from_millis(10));
        }
    }
}