superdoc_sdk/

agent.rs

use crate::{
    generated::types::{
        DocCapabilitiesGetBoundParams, DocCapabilitiesGetResult, DocPresetDispatchBoundParams,
        DocPresetDispatchResult, DocPresetGetSystemPromptParams, DocPresetGetSystemPromptResult,
        DocPresetGetToolsParams, DocPresetGetToolsResult,
    },
    wire::Presence,
    Document, Error, InvokeOptions, Result,
};
use serde::Deserialize;
use serde_json::Value;
use std::sync::OnceLock;

pub(crate) fn prepare_core_dispatch(
    tool_name: &str,
    args: &mut Value,
    options: &mut InvokeOptions,
    configured: Option<&DocPresetGetToolsParams>,
) -> Result<()> {
    if !args.is_object() {
        let mut error = Error::new("INVALID_ARGUMENT", "tool arguments must be an object");
        error.details = Some(Box::new(serde_json::json!({"toolName": tool_name})));
        return Err(error);
    }
    let authority = action_policy()?;
    validate_core_policy(tool_name, args, configured, authority)?;
    let Some(tool) = authority
        .preflight
        .tools
        .iter()
        .find(|tool| tool.name == tool_name)
    else {
        let mut error = Error::new("TOOL_DISPATCH_NOT_FOUND", "Unknown Core tool");
        error.host_message = Some(format!("Unknown tool: {tool_name}").into());
        error.details = Some(Box::new(
            serde_json::json!({"toolName":tool_name,"preset":"core"}),
        ));
        return Err(error);
    };
    validate_core_arguments(tool, args, &authority.preflight)?;
    if options.dry_run == Some(true) {
        let mut error = Error::new(
            "CAPABILITY_UNSUPPORTED",
            "the core preset dispatcher does not support dryRun previews",
        );
        error.operation_id = Some("doc.preset.dispatch".into());
        return Err(error);
    }
    if options
        .change_mode
        .as_deref()
        .is_some_and(|mode| !["direct", "tracked"].contains(&mode))
    {
        return Err(Error::new(
            "INVALID_ARGUMENT",
            "core dispatch change mode must be direct or tracked",
        ));
    }
    let explicit = (tool_name == "superdoc_perform_action")
        .then(|| args.get("expectedRevision"))
        .flatten();
    let revision = match explicit {
        Some(value) => value.as_str(),
        None => options.expected_revision.as_deref(),
    };
    if revision.is_some() || explicit.is_some() {
        let revision = revision.filter(|revision| !revision.is_empty());
        let Some(revision) = revision else {
            let mut error = Error::new(
                "INVALID_ARGUMENT",
                "expectedRevision must be a non-empty live revision or a non-negative safe integer.",
            );
            error.details = Some(Box::new(
                serde_json::json!({"preset":"core","option":"expectedRevision"}),
            ));
            return Err(error);
        };
        options.expected_revision = Some(revision.to_owned());
    }
    Ok(())
}

pub(crate) fn stale_revision_message(action: &str) -> String {
    format!("{action}: expectedRevision is stale. Inspect this live document again before editing.")
}

pub(crate) fn stale_action_revision_receipt(
    args: &Value,
    live_revision: &str,
    configured: Option<&DocPresetGetToolsParams>,
) -> Result<Value> {
    let action = args.get("action").and_then(Value::as_str);
    let intent = action.unwrap_or("superdoc_perform_action");
    let mut receipt = serde_json::json!({
        "status": "failed",
        "intent": intent,
        "preSnapshot": { "revision": live_revision },
        "selectedTargets": [],
        "executedOperations": [],
        "verification": [],
        "errors": [{
            "code": "REVISION_CONFLICT",
            "message": stale_revision_message(intent),
            "recovery": { "kind": "reinspect" },
        }],
    });
    let enforcing = configured.is_some_and(
        |configured| matches!(&configured.enforce_change_mode, Presence::Present(mode) if mode == "tracked"),
    );
    if enforcing {
        let requested_tracked = args.get("changeMode").and_then(Value::as_str) == Some("tracked");
        let actions = &action_policy()?.actions;
        let disposition = action.and_then(|action| {
            actions
                .iter()
                .find(|entry| entry.name == action)
                .map(|entry| entry.disposition.as_str())
        });
        let (change_mode, forced) = match disposition {
            Some("tracked") => ("tracked", !requested_tracked),
            Some(_) => ("direct", requested_tracked),
            None => (
                if requested_tracked {
                    "tracked"
                } else {
                    "direct"
                },
                false,
            ),
        };
        receipt["changeMode"] = Value::from(change_mode);
        receipt["changeModeForced"] = Value::from(forced);
    }
    Ok(receipt)
}

fn invalid_core_arguments(message: String, details: Value) -> Error {
    let mut error = Error::new("INVALID_ARGUMENT", "Core tool arguments are invalid");
    error.host_message = Some(message.into());
    error.details = Some(Box::new(details));
    error
}

fn trim_core_csv_name<'a>(name: &'a str, authority: &CorePreflight) -> &'a str {
    name.trim_matches(|character: char| authority.trim_code_points.contains(&(character as u32)))
}

fn csv_names<'a>(value: &'a Presence<String>, authority: &CorePreflight) -> Vec<&'a str> {
    match value {
        Presence::Missing => Vec::new(),
        Presence::Present(value) => value
            .split(',')
            .map(|name| trim_core_csv_name(name, authority))
            .filter(|name| !name.is_empty())
            .collect(),
    }
}

fn validate_core_policy(
    tool_name: &str,
    args: &Value,
    configured: Option<&DocPresetGetToolsParams>,
    authority: &ActionPolicy,
) -> Result<()> {
    let Some(configured) = configured else {
        return Ok(());
    };
    let enforcing = match &configured.enforce_change_mode {
        Presence::Missing => false,
        Presence::Present(mode) if mode == "tracked" => true,
        Presence::Present(mode) => {
            return Err(invalid_core_arguments(
                format!(
                    "enforceChangeMode: only \"tracked\" is supported; got {}.",
                    serde_json::json!(mode)
                ),
                serde_json::json!({"enforceChangeMode":mode,"supported":["tracked"]}),
            ))
        }
    };
    let allow = csv_names(&configured.allow_direct_actions, &authority.preflight);
    if !configured.allow_direct_actions.is_missing() && !enforcing {
        return Err(invalid_core_arguments(
            "allowDirectActions requires enforceChangeMode: \"tracked\"; without enforcement every action already runs as requested.".into(),
            serde_json::json!({"allowDirectActions":allow})));
    }
    if enforcing
        && authority
            .preflight
            .enforcement_side_doors
            .iter()
            .any(|name| name == tool_name)
    {
        return Err(invalid_core_arguments(
            format!("{tool_name} bypasses the tracked-changes policy and is refused under enforceChangeMode: \"tracked\"."),
            serde_json::json!({"toolName":tool_name,"reason":"enforce-change-mode"})));
    }
    if tool_name == "superdoc_perform_action" {
        if let Some(action) = args.get("action").and_then(Value::as_str) {
            if csv_names(&configured.exclude_actions, &authority.preflight).contains(&action) {
                return Err(invalid_core_arguments(
                    format!("Action {action} is excluded by configuration."),
                    serde_json::json!({"toolName":tool_name,"action":action,"excluded":true,"reason":"excluded"}),
                ));
            }
            if enforcing
                && !allow.contains(&"*")
                && !allow.contains(&action)
                && authority
                    .actions
                    .iter()
                    .any(|entry| entry.name == action && entry.disposition == "direct")
            {
                return Err(invalid_core_arguments(
                    format!("Action {action} edits the document directly and is refused under enforceChangeMode: \"tracked\". Pass it in allowDirectActions to permit it."),
                    serde_json::json!({"toolName":tool_name,"action":action,"excluded":true,"reason":"direct-only"})));
            }
        }
    }
    Ok(())
}

fn meaningful_argument_key(key: &str, authority: &CorePreflight) -> bool {
    key.chars().any(|character| {
        let point = character as u32;
        let at = authority
            .letter_number_ranges
            .partition_point(|range| range[1] < point);
        authority
            .letter_number_ranges
            .get(at)
            .is_some_and(|range| range[0] <= point)
    })
}

fn validate_core_arguments(tool: &CoreTool, args: &Value, authority: &CorePreflight) -> Result<()> {
    let unknown = args
        .as_object()
        .expect("Core arguments were checked")
        .keys()
        .filter(|name| !tool.known_keys.contains(name) && meaningful_argument_key(name, authority))
        .cloned()
        .collect::<Vec<_>>();
    if !unknown.is_empty() {
        let retired = unknown
            .iter()
            .filter(|name| authority.retired_arguments.contains_key(*name))
            .cloned()
            .collect::<Vec<_>>();
        let message = if retired.is_empty() {
            format!(
                "Unknown argument(s) for {}: {}",
                tool.name,
                unknown.join(", ")
            )
        } else {
            retired
                .iter()
                .map(|name| authority.retired_arguments[name].as_str())
                .collect::<Vec<_>>()
                .join(" ")
        };
        let mut details = serde_json::json!({"toolName":tool.name,"unknownKeys":unknown,"knownKeys":tool.known_keys});
        if !retired.is_empty() {
            details["retiredKeys"] = serde_json::json!(retired);
        }
        return Err(invalid_core_arguments(message, details));
    }
    let missing = tool
        .required_keys
        .iter()
        .filter(|name| args.get(name.as_str()).is_none_or(Value::is_null))
        .collect::<Vec<_>>();
    if !missing.is_empty() {
        return Err(invalid_core_arguments(
            format!(
                "Missing required argument(s) for {}: {}",
                tool.name,
                missing
                    .iter()
                    .map(|name| name.as_str())
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
            serde_json::json!({"toolName":tool.name,"missingKeys":missing}),
        ));
    }
    Ok(())
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct CoreTool {
    name: String,
    known_keys: Vec<String>,
    required_keys: Vec<String>,
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct CorePreflight {
    catalog_sha256: String,
    dispatch_sha256: String,
    action_source_sha256: String,
    tools: Vec<CoreTool>,
    enforcement_side_doors: Vec<String>,
    retired_arguments: std::collections::BTreeMap<String, String>,
    unicode_version: String,
    character_data_sha256: String,
    letter_number_ranges: Vec<[u32; 2]>,
    trim_code_points: Vec<u32>,
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ActionPolicy {
    version: u32,
    descriptor_sha256: String,
    contract_hash: String,
    actions: Vec<Action>,
    preflight: CorePreflight,
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Action {
    name: String,
    operations: Vec<String>,
    disposition: String,
}

fn action_policy() -> Result<&'static ActionPolicy> {
    static POLICY: OnceLock<std::result::Result<ActionPolicy, String>> = OnceLock::new();
    let policy = POLICY.get_or_init(|| {
        serde_json::from_str(include_str!("../assets/action-policy.json"))
            .map_err(|error| error.to_string())
    });
    let policy = policy
        .as_ref()
        .map_err(|_| Error::new("HOST_PROTOCOL_ERROR", "generated action policy is invalid"))?;
    if policy.version != 2
        || policy.descriptor_sha256.len() != 64
        || policy.contract_hash != crate::generated::CONTRACT_HASH
        || policy.preflight.catalog_sha256.len() != 64
        || policy.preflight.dispatch_sha256.len() != 64
        || policy.preflight.action_source_sha256.len() != 64
        || policy.preflight.tools.len() != 6
        || policy.preflight.unicode_version.is_empty()
        || policy.preflight.character_data_sha256.len() != 64
        || policy.preflight.letter_number_ranges.is_empty()
    {
        return Err(Error::new(
            "HOST_PROTOCOL_ERROR",
            "generated action policy does not match the SDK contract",
        ));
    }
    Ok(policy)
}

fn unavailable_from_operations(operations: &Value) -> Result<Vec<String>> {
    Ok(action_policy()?
        .actions
        .iter()
        .filter(|action| {
            action.disposition == "tracked"
                && action.operations.iter().any(|id| {
                    let capability = &operations[id];
                    capability["available"] != true
                        || (capability["tracked"] != true
                            && capability["trackedSupport"] != "conditional")
                })
        })
        .map(|action| action.name.clone())
        .collect())
}

fn unavailable_tracked_actions(capabilities: &DocCapabilitiesGetResult) -> Result<Vec<String>> {
    unavailable_from_operations(
        &serde_json::to_value(&capabilities.operations).map_err(Error::decode)?,
    )
}

fn merge_exclusions(
    configured: &Presence<String>,
    unavailable: &[String],
    authority: &CorePreflight,
) -> Presence<String> {
    if unavailable.is_empty() {
        return configured.clone();
    }
    let mut names = Vec::<String>::new();
    let configured = match configured {
        Presence::Present(value) => value
            .split(',')
            .map(|name| trim_core_csv_name(name, authority))
            .filter(|value| !value.is_empty())
            .map(str::to_owned)
            .collect::<Vec<_>>(),
        Presence::Missing => Vec::new(),
    };
    for name in configured.into_iter().chain(unavailable.iter().cloned()) {
        if !names.contains(&name) {
            names.push(name);
        }
    }
    Presence::Present(names.join(","))
}

#[derive(Clone)]
pub struct CoreToolkit {
    document: Document,
    policy: DocPresetGetToolsParams,
    tools: DocPresetGetToolsResult,
    system_prompt: DocPresetGetSystemPromptResult,
    capabilities: Option<Box<DocCapabilitiesGetResult>>,
}

impl CoreToolkit {
    pub fn tools(&self) -> &DocPresetGetToolsResult {
        &self.tools
    }

    pub fn system_prompt(&self) -> &DocPresetGetSystemPromptResult {
        &self.system_prompt
    }

    pub fn prompt(&self) -> &str {
        &self.system_prompt.prompt
    }

    pub fn policy(&self) -> &DocPresetGetToolsParams {
        &self.policy
    }

    pub fn capabilities_snapshot(&self) -> Option<&DocCapabilitiesGetResult> {
        self.capabilities.as_deref()
    }

    pub async fn dispatch(
        &self,
        params: &DocPresetDispatchBoundParams,
    ) -> Result<DocPresetDispatchResult> {
        self.dispatch_with_options(params, InvokeOptions::default())
            .await
    }

    pub async fn dispatch_with_options(
        &self,
        params: &DocPresetDispatchBoundParams,
        mut options: InvokeOptions,
    ) -> Result<DocPresetDispatchResult> {
        let mut params = params.clone();
        let mut args = match &params.args {
            Presence::Missing => serde_json::json!({}),
            Presence::Present(args) => args.clone(),
        };
        if params.tool_name == "superdoc_perform_action" && args.get("expectedRevision").is_some() {
            params.expected_revision = Presence::Missing;
        } else if !params.expected_revision.is_missing() {
            options.expected_revision = None;
        }
        prepare_core_dispatch(
            &params.tool_name,
            &mut args,
            &mut options,
            Some(&self.policy),
        )?;
        if let Presence::Present(revision) = &params.expected_revision {
            let invalid = || {
                Error::new("INVALID_ARGUMENT", "expectedRevision must be a non-empty live revision or a non-negative safe integer.")
            };
            let revision = match serde_json::to_value(revision).map_err(Error::decode)? {
                Value::String(value) if !value.is_empty() => value,
                Value::Number(value) => {
                    let value = value.as_f64().ok_or_else(invalid)?;
                    if value < 0.0
                        || value.fract() != 0.0
                        || value > crate::wire::Integer::MAX as f64
                    {
                        return Err(invalid());
                    }
                    if value == 0.0 {
                        "0".into()
                    } else {
                        value.to_string()
                    }
                }
                _ => return Err(invalid()),
            };
            options.expected_revision = Some(revision);
        }
        if let Some(receipt) = self
            .document
            .check_core_revision(&params.tool_name, &args, &options, Some(&self.policy))
            .await?
        {
            return Ok(DocPresetDispatchResult::new().with_result(receipt));
        }
        params.args = Presence::Present(args);
        params.preset = self.policy.preset.clone();
        params.exclude_actions = self.policy.exclude_actions.clone();
        params.enforce_change_mode = self.policy.enforce_change_mode.clone();
        params.allow_direct_actions = self.policy.allow_direct_actions.clone();
        self.document
            .api()
            .preset()
            .dispatch_with_options(&params, options)
            .await
    }
}

#[derive(Clone)]
pub struct FullToolkit {
    document: Document,
    policy: DocPresetGetToolsParams,
    tools: DocPresetGetToolsResult,
    system_prompt: DocPresetGetSystemPromptResult,
}

impl FullToolkit {
    pub fn tools(&self) -> &DocPresetGetToolsResult {
        &self.tools
    }

    pub fn system_prompt(&self) -> &DocPresetGetSystemPromptResult {
        &self.system_prompt
    }

    pub fn prompt(&self) -> &str {
        &self.system_prompt.prompt
    }

    pub fn policy(&self) -> &DocPresetGetToolsParams {
        &self.policy
    }

    pub async fn dispatch(
        &self,
        params: &DocPresetDispatchBoundParams,
    ) -> Result<DocPresetDispatchResult> {
        self.dispatch_with_options(params, InvokeOptions::default())
            .await
    }

    pub async fn dispatch_with_options(
        &self,
        params: &DocPresetDispatchBoundParams,
        options: InvokeOptions,
    ) -> Result<DocPresetDispatchResult> {
        let mut params = params.clone();
        params.preset = self.policy.preset.clone();
        params.exclude_actions = self.policy.exclude_actions.clone();
        params.enforce_change_mode = self.policy.enforce_change_mode.clone();
        params.allow_direct_actions = self.policy.allow_direct_actions.clone();
        self.document
            .api()
            .preset()
            .dispatch_with_options(&params, options)
            .await
    }
}

impl Document {
    pub async fn full_toolkit(&self, input: &DocPresetGetToolsParams) -> Result<FullToolkit> {
        let mut policy = input.clone();
        if let Presence::Present(preset) = &policy.preset {
            if preset != "full" {
                return Err(Error::new(
                    "INVALID_ARGUMENT",
                    "full_toolkit requires the full preset",
                ));
            }
        }
        policy.preset = Presence::Present("full".to_owned());
        let tools = self.api().preset().get_tools(&policy).await?;
        let mut prompt_params = DocPresetGetSystemPromptParams::new();
        prompt_params.preset = policy.preset.clone();
        prompt_params.exclude_actions = policy.exclude_actions.clone();
        prompt_params.enforce_change_mode = policy.enforce_change_mode.clone();
        prompt_params.allow_direct_actions = policy.allow_direct_actions.clone();
        let system_prompt = self
            .api()
            .preset()
            .get_system_prompt(&prompt_params)
            .await?;
        Ok(FullToolkit {
            document: self.clone(),
            policy,
            tools,
            system_prompt,
        })
    }
}

impl Document {
    pub async fn core_toolkit(&self, input: &DocPresetGetToolsParams) -> Result<CoreToolkit> {
        let mut policy = input.clone();
        if let Presence::Present(preset) = &policy.preset {
            if preset != "core" {
                return Err(Error::new(
                    "INVALID_ARGUMENT",
                    "core_toolkit requires the core preset",
                ));
            }
        }
        policy.preset = Presence::Present("core".to_owned());
        let capabilities = if matches!(&policy.enforce_change_mode, Presence::Present(mode) if mode == "tracked")
        {
            Some(Box::new(
                self.api()
                    .capabilities()
                    .get(&DocCapabilitiesGetBoundParams::new())
                    .await?,
            ))
        } else {
            None
        };
        if let Some(capabilities) = &capabilities {
            policy.exclude_actions = merge_exclusions(
                &policy.exclude_actions,
                &unavailable_tracked_actions(capabilities)?,
                &action_policy()?.preflight,
            );
        }
        let tools = self.api().preset().get_tools(&policy).await?;
        let mut prompt_params = DocPresetGetSystemPromptParams::new();
        prompt_params.preset = policy.preset.clone();
        prompt_params.exclude_actions = policy.exclude_actions.clone();
        prompt_params.enforce_change_mode = policy.enforce_change_mode.clone();
        prompt_params.allow_direct_actions = policy.allow_direct_actions.clone();
        let system_prompt = self
            .api()
            .preset()
            .get_system_prompt(&prompt_params)
            .await?;
        Ok(CoreToolkit {
            document: self.clone(),
            policy,
            tools,
            system_prompt,
            capabilities,
        })
    }
}

#[cfg(test)]
#[path = "agent_tests.rs"]
mod tests;