magi-code 0.80.1

//! Source-only, memory-only session registry for the application resource profile.
use super::super::persistent::valid_id;
use super::*;
use std::sync::Weak;

pub(in crate::service) const RESOURCE_METHODS: [&str; 4] = [
    "session.resources.register",
    "session.resources.list",
    "session.resources.unregister",
    "session.resources.clear",
];
const MAX_SNAPSHOTS: usize = 256;
const MAX_BYTES: usize = 8 * 1024 * 1024;
const MAX_SNAPSHOT_BYTES: usize = 32768;
const REGISTRY_IDLE_TIMEOUT: Duration = Duration::from_secs(15 * 60);

#[derive(Clone, Deserialize, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
enum Resource {
    Tool {
        contract: Arc<Contract>,
    },
    Skill {
        name: String,
        description: String,
        content: String,
    },
}
impl Resource {
    fn name(&self) -> &str {
        match self {
            Self::Tool { contract } => &contract.name,
            Self::Skill { name, .. } => name,
        }
    }
}
struct Snapshot {
    revision: String,
    resources: BTreeMap<String, Resource>,
    sources: BTreeMap<String, (Value, Arc<str>)>,
    bytes: usize,
    skill: Option<Arc<str>>,
}
impl Snapshot {
    fn descriptors(&self) -> Vec<Value> {
        let mut entries: Vec<_> = self
            .sources
            .values()
            .map(|(descriptor, _)| descriptor.clone())
            .collect();
        entries.sort_by(|left, right| {
            (left["kind"].as_str(), left["name"].as_str())
                .cmp(&(right["kind"].as_str(), right["name"].as_str()))
        });
        entries
    }
}
#[derive(Default)]
pub(in crate::service) struct Registry {
    sessions: HashMap<String, Arc<Snapshot>>,
    retained: Vec<Weak<Snapshot>>,
    idle_since: HashMap<String, Instant>,
    source_capture_reserved: bool,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Register {
    #[serde(deserialize_with = "Deserialize::deserialize")]
    expected_revision: Option<String>,
    resource: Resource,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Unregister {
    expected_revision: String,
    resource_id: String,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Clear {
    expected_registry_revision: String,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Selection {
    expected_registry_revision: String,
    #[serde(deserialize_with = "Deserialize::deserialize")]
    executor_id: Option<String>,
    #[serde(deserialize_with = "Deserialize::deserialize")]
    executor_generation: Option<u64>,
    tool_authorizations: Vec<ToolAuthorization>,
}

#[derive(Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
struct ToolAuthorization {
    resource_id: String,
    revision_id: String,
    revision_sha256: String,
}
#[derive(Clone)]
pub(in crate::service) struct CapturedTurn {
    turns: Vec<InlineTurn>,
    pub skill: Option<Arc<str>>,
    snapshot: Arc<Snapshot>,
    accepted_controller: Option<Value>,
}
impl CapturedTurn {
    pub fn validate_executor(
        &self,
        state: &ApplicationState,
        connection: &str,
    ) -> Result<(), Code> {
        for turn in &self.turns {
            state.authorize(connection, &turn.executor_id, turn.executor_generation)?;
        }
        Ok(())
    }
    pub fn validate_catalog(
        &self,
        runtime: &super::super::runtime::ServiceRuntime,
    ) -> Result<(), Code> {
        for name in self.snapshot.resources.keys() {
            validate_catalog_name(name, runtime)?;
        }
        Ok(())
    }
    /// Called only by the coordinator after preparation and authority revalidation.
    pub fn accept_controller(mut self, attribution: Value) -> Self {
        self.accepted_controller = Some(attribution);
        self
    }
    pub fn manifest(&self) -> Value {
        let snapshot = &self.snapshot;
        let mut manifest = json!({"profile":PROFILE,"registry_revision":snapshot.revision,
                    "resources":snapshot.descriptors(),
                    "tools":self.turns.iter().map(InlineTurn::manifest).collect::<Vec<_>>() });
        if let Some(controller) = &self.accepted_controller {
            manifest["accepted_controller"] = controller.clone();
            manifest["accepted_executor"] = self.turns.first().map_or(Value::Null, |turn| {
                        json!({"executor_id":turn.executor_id,"executor_generation":turn.executor_generation})
                    });
        }
        manifest
    }
    pub fn tools(
        &self,
        state: SharedApplicationState,
        session: String,
        turn: String,
    ) -> ApplicationTools {
        let tools: Vec<_> = self
            .turns
            .iter()
            .map(|captured| captured.tools(Arc::clone(&state), session.clone(), turn.clone()))
            .collect();
        let snapshot = self.snapshot.clone();
        ApplicationTools {
            definitions: tools
                .iter()
                .flat_map(|tool| tool.definitions.clone())
                .collect(),
            manifest: self.manifest(),
            callback: Arc::new(move |name, arguments, context| {
                // Keep the exact accepted resource bytes charged until the worker drops its tools.
                let _retained = &snapshot;
                let tool = tools
                    .iter()
                    .find(|tool| tool.handles(name))
                    .expect("dispatch checks application tool name");
                (tool.callback)(name, arguments, context)
            }),
        }
    }
}

pub(super) fn validate_catalog_name(
    name: &str,
    runtime: &super::super::runtime::ServiceRuntime,
) -> Result<(), Code> {
    if runtime
        .discovered_skill_names
        .iter()
        .any(|candidate| candidate.eq_ignore_ascii_case(name))
        || crate::config::disabled_tool_names_from_settings(&runtime.settings)
            .iter()
            .any(|candidate| candidate.eq_ignore_ascii_case(name))
        || crate::config::disabled_skill_names_from_settings(&runtime.settings)
            .iter()
            .any(|candidate| candidate.eq_ignore_ascii_case(name))
    {
        return Err(Code::InvalidPayload);
    }
    Ok(())
}
pub(super) fn validate_resource_name(name: &str) -> Result<(), Code> {
    if name.is_empty()
        || name.len() > 48
        || !name.as_bytes()[0].is_ascii_lowercase()
        || !name
            .bytes()
            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == b'_')
        || name.starts_with("mcp_")
        || crate::tools::ToolCapability::from_dispatch_name(name).is_some()
    {
        return Err(Code::InvalidPayload);
    }
    Ok(())
}
impl Registry {
    fn validate_expected_resource_revision(
        &self,
        session: &str,
        payload: &Value,
    ) -> Result<(), Code> {
        let expected = payload
            .get("expected_revision")
            .ok_or(Code::InvalidPayload)?;
        if !expected.is_null() && !expected.as_str().is_some_and(valid_id) {
            return Err(Code::InvalidPayload);
        }
        let name = payload["resource"]["name"]
            .as_str()
            .or_else(|| payload["resource"]["contract"]["name"].as_str())
            .ok_or(Code::InvalidPayload)?;
        let previous = self
            .sessions
            .get(session)
            .and_then(|snapshot| snapshot.sources.get(name));
        match previous {
            Some((descriptor, _))
                if descriptor["revision_id"] == *expected
                    && descriptor["kind"] == payload["resource"]["kind"] =>
            {
                Ok(())
            }
            None if expected.is_null() => Ok(()),
            _ => Err(Code::Conflict),
        }
    }
    fn ensure_session(&mut self, session: &str) -> Result<(), Code> {
        if self.sessions.contains_key(session) {
            return Ok(());
        }
        self.retained.retain(|snapshot| snapshot.strong_count() > 0);
        if self.sessions.len() >= 256 || self.retained.len() >= MAX_SNAPSHOTS {
            return Err(Code::LimitExceeded);
        }
        let snapshot = Arc::new(Snapshot {
            revision: uuid::Uuid::new_v4().to_string(),
            resources: BTreeMap::new(),
            sources: BTreeMap::new(),
            bytes: 0,
            skill: None,
        });
        self.retained.push(Arc::downgrade(&snapshot));
        self.sessions.insert(session.to_owned(), snapshot);
        Ok(())
    }
    /// Only the coordinator knows whether an older revision is pinned by work.
    pub fn expire_idle(&mut self, now: Instant, protected: impl Fn(&str) -> bool) {
        self.sessions.retain(|session, snapshot| {
            if protected(session) {
                self.idle_since.remove(session);
                return true;
            }
            if snapshot.resources.is_empty() {
                return false;
            }
            let since = self.idle_since.entry(session.clone()).or_insert(now);
            now.saturating_duration_since(*since) < REGISTRY_IDLE_TIMEOUT
        });
        self.idle_since
            .retain(|session, _| self.sessions.contains_key(session));
        self.retained.retain(|snapshot| snapshot.strong_count() > 0);
    }
    pub fn nonempty(&self, session: &str) -> bool {
        self.sessions
            .get(session)
            .is_some_and(|snapshot| !snapshot.resources.is_empty())
    }
    pub fn reserve_source_capture(&mut self, session: &str, payload: &Value) -> Result<(), Code> {
        if self.source_capture_reserved {
            return Err(Code::ConfigurationBusy);
        }
        self.validate_expected_resource_revision(session, payload)?;
        let snapshots: Vec<_> = self.retained.iter().filter_map(Weak::upgrade).collect();
        // Keep staging charged even if removals create new snapshots beside pinned revisions.
        if snapshots
            .iter()
            .map(|snapshot| snapshot.bytes)
            .sum::<usize>()
            + 131072
            > MAX_BYTES
            || snapshots
                .iter()
                .map(|snapshot| snapshot.resources.len())
                .sum::<usize>()
                >= 1024
            || snapshots.len() >= MAX_SNAPSHOTS
            || (!self.sessions.contains_key(session) && self.sessions.len() >= 256)
        {
            return Err(Code::LimitExceeded);
        }
        self.source_capture_reserved = true;
        Ok(())
    }
    pub fn release_source_capture(&mut self) {
        self.source_capture_reserved = false;
    }
    pub fn route(
        &mut self,
        session: &str,
        method: &str,
        payload: &Value,
        runtime: &super::super::runtime::ServiceRuntime,
    ) -> Result<Value, Code> {
        #[cfg(test)]
        if method == "session.resources.register" {
            return self.publish_source(session, prepare_source(payload, runtime)?, runtime,
                json!({"instance_id":"test-instance","connection_id":"test-connection","operation_id":"test-operation","grant_generation":1}));
        }
        self.route_captured(session, method, payload, runtime, None)
    }
    pub fn publish_source(
        &mut self,
        session: &str,
        prepared: PreparedSource,
        runtime: &super::super::runtime::ServiceRuntime,
        attribution: Value,
    ) -> Result<Value, Code> {
        let mut descriptor = prepared.descriptor;
        descriptor["registered_by"] = attribution;
        self.route_captured(
            session,
            "session.resources.register",
            &prepared.payload,
            runtime,
            Some((descriptor, prepared.content)),
        )
    }
    fn route_captured(
        &mut self,
        session: &str,
        method: &str,
        payload: &Value,
        runtime: &super::super::runtime::ServiceRuntime,
        captured: Option<(Value, Arc<str>)>,
    ) -> Result<Value, Code> {
        if method == "session.resources.list" {
            if payload != &json!({}) {
                return Err(Code::InvalidPayload);
            }
            self.ensure_session(session)?;
            let snapshot = &self.sessions[session];
            return Ok(
                json!({"registry_revision":snapshot.revision,"entries":snapshot.descriptors()}),
            );
        }
        let mut resources = self
            .sessions
            .get(session)
            .map(|s| s.resources.clone())
            .unwrap_or_default();
        let mut sources = self
            .sessions
            .get(session)
            .map(|snapshot| snapshot.sources.clone())
            .unwrap_or_default();
        let mut registered_entry = None;
        let mut response_details = json!({});
        match method {
            "session.resources.register" => {
                let params: Register = decode(payload)?;
                self.validate_expected_resource_revision(session, payload)?;
                let _expected_revision = params.expected_revision;
                validate_resource_name(params.resource.name())?;
                validate_catalog_name(params.resource.name(), runtime)?;
                match &params.resource {
                    Resource::Tool { contract } => validate_contract(contract)?,
                    Resource::Skill {
                        description,
                        content,
                        ..
                    } => {
                        if description.len() > 512
                            || content.len() > 16384
                            || content.contains('\0')
                        {
                            return Err(Code::InvalidPayload);
                        }
                    }
                }
                let name = params.resource.name().to_owned();
                let (mut descriptor, content) = captured.ok_or(Code::InvalidPayload)?;
                if let Some((previous, _)) = sources.get(&name) {
                    descriptor["resource_id"] = previous["resource_id"].clone();
                }
                registered_entry = Some(descriptor.clone());
                sources.insert(name, (descriptor, content));
                resources.insert(params.resource.name().to_owned(), params.resource);
            }
            "session.resources.unregister" => {
                let params: Unregister = decode(payload)?;
                if !valid_id(&params.resource_id) || !valid_id(&params.expected_revision) {
                    return Err(Code::InvalidPayload);
                }
                let name = sources
                    .iter()
                    .find(|(_, (descriptor, _))| descriptor["resource_id"] == params.resource_id)
                    .map(|(name, _)| name.clone())
                    .ok_or(Code::ResourceNotFound)?;
                if sources[&name].0["revision_id"] != params.expected_revision {
                    return Err(Code::Conflict);
                }
                resources.remove(&name);
                sources.remove(&name);
                response_details = json!({"resource_id":params.resource_id,"removed_revision":params.expected_revision});
            }
            "session.resources.clear" => {
                let params: Clear = decode(payload)?;
                if !valid_id(&params.expected_registry_revision) {
                    return Err(Code::InvalidPayload);
                }
                let snapshot = self.sessions.get(session).ok_or(Code::Conflict)?;
                if snapshot.revision != params.expected_registry_revision {
                    return Err(Code::Conflict);
                }
                if resources.is_empty() {
                    return Ok(json!({"registry_revision":snapshot.revision,"removed_count":0}));
                }
                response_details = json!({"removed_count":resources.len()});
                resources.clear();
                sources.clear();
            }
            _ => return Err(Code::UnsupportedOperation),
        }
        let bytes = serde_json::to_vec(&resources)
            .map_err(|_| Code::InternalError)?
            .len()
            + sources
                .values()
                .map(|(descriptor, content)| descriptor.to_string().len() + content.len())
                .sum::<usize>();
        let tools = resources
            .values()
            .filter(|r| matches!(r, Resource::Tool { .. }))
            .count();
        let skill_bytes: usize = resources
            .values()
            .map(|r| match r {
                Resource::Skill { content, .. } => content.len() + 1,
                _ => 0,
            })
            .sum();
        if tools > 16
            || resources.len() - tools > 16
            || bytes > MAX_SNAPSHOT_BYTES
            || skill_bytes > 16384
        {
            return Err(Code::LimitExceeded);
        }
        // Build instructions once per revision, not once per tool or callback.
        let skill = resources
            .values()
            .filter_map(|resource| match resource {
                Resource::Skill { content, .. } => Some(content.as_str()),
                _ => None,
            })
            .collect::<Vec<_>>()
            .join("\n");
        let skill: Option<Arc<str>> = (!skill.is_empty()).then(|| Arc::from(skill));
        // Serialized resources and their combined instruction allocation are both retained.
        let bytes = bytes + skill.as_ref().map_or(0, |skill| skill.len());
        self.retained.retain(|entry| entry.strong_count() > 0);
        let retained_bytes: usize = self
            .retained
            .iter()
            .filter_map(Weak::upgrade)
            .map(|s| s.bytes)
            .sum();
        let retained_resources: usize = self
            .retained
            .iter()
            .filter_map(Weak::upgrade)
            .map(|s| s.resources.len())
            .sum();
        // Reserve replacement peak as well as snapshots pinned by admission/active workers.
        if (!resources.is_empty()
            && (self.retained.len() >= MAX_SNAPSHOTS
                || retained_bytes + bytes + usize::from(self.source_capture_reserved) * 131072
                    > MAX_BYTES
                || retained_resources
                    + resources.len()
                    + usize::from(self.source_capture_reserved)
                    > 1024))
            || (!self.sessions.contains_key(session) && self.sessions.len() >= 256)
        {
            return Err(Code::LimitExceeded);
        }
        let snapshot = Arc::new(Snapshot {
            revision: uuid::Uuid::new_v4().to_string(),
            resources,
            sources,
            bytes,
            skill,
        });
        let mut response = response_details;
        response["registry_revision"] = json!(snapshot.revision);
        if let Some(entry) = registered_entry {
            response["entry"] = entry;
        }
        self.retained.push(Arc::downgrade(&snapshot));
        self.sessions.insert(session.to_owned(), snapshot);
        Ok(response)
    }
}
impl ApplicationState {
    pub fn capture_session(
        &self,
        connection: &str,
        session: &str,
        payload: &Value,
    ) -> Result<CapturedTurn, Code> {
        let selection: Selection = decode(payload)?;
        let snapshot = self
            .registry
            .sessions
            .get(session)
            .ok_or(Code::InvalidPayload)?;
        if selection.expected_registry_revision != snapshot.revision {
            return Err(Code::InvalidPayload);
        }
        let expected: Vec<ToolAuthorization> = snapshot
            .sources
            .values()
            .filter(|(descriptor, _)| descriptor["kind"] == "tool")
            .map(|(descriptor, _)| ToolAuthorization {
                resource_id: descriptor["resource_id"]
                    .as_str()
                    .expect("resource identity")
                    .into(),
                revision_id: descriptor["revision_id"]
                    .as_str()
                    .expect("revision identity")
                    .into(),
                revision_sha256: descriptor["revision_sha256"]
                    .as_str()
                    .expect("revision hash")
                    .into(),
            })
            .collect();
        {
            let supplied = &selection.tool_authorizations;
            if supplied.len() != expected.len()
                || expected.iter().any(|entry| {
                    supplied
                        .iter()
                        .filter(|candidate| *candidate == entry)
                        .count()
                        != 1
                })
            {
                return Err(Code::InvalidPayload);
            }
        }
        let skill = snapshot.skill.clone();
        let contracts: Vec<_> = snapshot
            .resources
            .values()
            .filter_map(|resource| match resource {
                Resource::Tool { contract } => Some(contract),
                _ => None,
            })
            .collect();
        let mut turns = Vec::new();
        if contracts.is_empty() {
            if selection.executor_id.is_some() || selection.executor_generation.is_some() {
                return Err(Code::InvalidPayload);
            }
        } else {
            let executor_id = selection.executor_id.ok_or(Code::InvalidPayload)?;
            let generation = selection.executor_generation.ok_or(Code::InvalidPayload)?;
            self.authorize(connection, &executor_id, generation)?;
            for contract in contracts {
                turns.push(InlineTurn {
                    resource_identity: snapshot.sources.get(&contract.name).map(
                        |(descriptor, _)| ResourceIdentity {
                            resource_id: descriptor["resource_id"]
                                .as_str()
                                .expect("resource id")
                                .into(),
                            revision_id: descriptor["revision_id"]
                                .as_str()
                                .expect("revision id")
                                .into(),
                            revision_sha256: descriptor["revision_sha256"]
                                .as_str()
                                .expect("revision hash")
                                .into(),
                        },
                    ),
                    executor_id: executor_id.clone(),
                    executor_generation: generation,
                    tool: Arc::clone(contract),
                    skill: skill.clone(),
                });
            }
        }
        Ok(CapturedTurn {
            turns,
            skill,
            snapshot: Arc::clone(snapshot),
            accepted_controller: None,
        })
    }
}