superdoc-sdk 2.25.0

use crate::{
    generated::{
        operations::{DocApi, DocApiPreset},
        types::{DocCloseBoundParams, DocCloseResult, DocOpenParams, DocOpenResult},
    },
    invoke::Invoker,
    runtime::{
        contract, policy,
        process::{CallState, HostMode, HostOptions, Transport},
        routing, trim_js_whitespace,
    },
    Error, InvokeOptions, Result,
};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use serde_json::{json, Value};
use std::{
    collections::{BTreeMap, BTreeSet},
    path::Path,
    sync::{
        atomic::{AtomicU64, AtomicU8, Ordering},
        Arc, Mutex,
    },
    time::{Duration, SystemTime, UNIX_EPOCH},
};
use tokio::sync::{broadcast, watch, Notify};

#[derive(Clone)]
pub struct CollaborationAuth {
    token: String,
}
impl CollaborationAuth {
    pub fn token(token: impl Into<String>) -> Result<Self> {
        let token = token.into();
        if token.is_empty() {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "collaboration token must not be empty",
            ));
        }
        Ok(Self { token })
    }
    fn json(&self) -> Value {
        json!({"type":"token","token":self.token})
    }
}
impl std::fmt::Debug for CollaborationAuth {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("CollaborationAuth([redacted])")
    }
}

#[derive(Clone, Debug, Default)]
pub struct OpenOptions {
    pub invoke: InvokeOptions,
    pub collaboration_auth: Option<CollaborationAuth>,
}

#[derive(Clone, Default)]
pub struct ClientOptions {
    pub default_change_mode: Option<String>,
    pub user_name: Option<String>,
    pub user_email: Option<String>,
    /// Opt in to negotiated document RPC on the CLI host. Collaboration and
    /// other ineligible opens keep their CLI route before any request is sent.
    pub document_rpc: bool,
}

struct ClientInner {
    transport: Transport,
    options: ClientOptions,
    opening: Mutex<BTreeSet<String>>,
    lifecycle: Mutex<ClientLifecycle>,
    idle: Notify,
    documents: Mutex<BTreeMap<String, Arc<DocumentPersistence>>>,
}

#[derive(Default)]
struct ClientLifecycle {
    active: usize,
    stopping: bool,
    managed_documents: bool,
    opened_documents: bool,
    shutdown: Option<watch::Receiver<Option<Result<()>>>>,
}

struct ActiveCall(Arc<ClientInner>);

impl Drop for ActiveCall {
    fn drop(&mut self) {
        let mut lifecycle = self.0.lifecycle.lock().expect("SDK lifecycle lock");
        lifecycle.active -= 1;
        if lifecycle.active == 0 {
            self.0.idle.notify_one();
        }
    }
}

#[derive(Default)]
struct DirtyDocument {
    confirmed: bool,
    pending: usize,
    generation: u64,
}

struct DocumentPersistence {
    state: Arc<AtomicU8>,
    dirty: Mutex<DirtyDocument>,
}

impl DocumentPersistence {
    fn save_generation(&self) -> Option<u64> {
        let dirty = self.dirty.lock().expect("SDK dirty document lock");
        (dirty.pending == 0).then_some(dirty.generation)
    }

    fn saved(&self, generation: u64) {
        let mut dirty = self.dirty.lock().expect("SDK dirty document lock");
        if dirty.generation == generation {
            dirty.confirmed = false;
        }
    }
}

struct MutationGuard {
    document: Arc<DocumentPersistence>,
    transmission: Arc<CallState>,
    settled: bool,
}

impl MutationGuard {
    fn new(
        document: Arc<DocumentPersistence>,
        transmission: Arc<CallState>,
        transport: &Transport,
    ) -> Self {
        transport.retain_state();
        {
            let mut dirty = document.dirty.lock().expect("SDK dirty document lock");
            dirty.pending += 1;
            dirty.generation += 1;
        }
        Self {
            document,
            transmission,
            settled: false,
        }
    }

    fn settle(&mut self, applied: bool) {
        let mut dirty = self.document.dirty.lock().expect("SDK dirty document lock");
        dirty.pending -= 1;
        dirty.confirmed |= applied;
        self.settled = true;
    }
}

impl Drop for MutationGuard {
    fn drop(&mut self) {
        if !self.settled {
            self.settle(!self.transmission.cancel_before_send());
        }
    }
}

#[derive(Clone)]
pub struct Client(Arc<ClientInner>);

struct OpenGuard {
    client: Client,
    id: String,
}
impl Drop for OpenGuard {
    fn drop(&mut self) {
        self.client
            .0
            .opening
            .lock()
            .expect("SDK opening lock")
            .remove(&self.id);
    }
}

impl Client {
    pub async fn connect(host: HostOptions) -> Result<Self> {
        Self::connect_with_options(host, ClientOptions::default()).await
    }

    pub async fn connect_with_options(host: HostOptions, options: ClientOptions) -> Result<Self> {
        if options
            .default_change_mode
            .as_deref()
            .is_some_and(|mode| !["direct", "tracked"].contains(&mode))
        {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "default change mode must be direct or tracked",
            ));
        }
        if options
            .user_name
            .as_deref()
            .is_some_and(|name| trim_js_whitespace(name).is_empty())
            || (options.user_email.is_some() && options.user_name.is_none())
        {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "user identity requires a non-empty name",
            ));
        }
        Ok(Self(Arc::new(ClientInner {
            transport: Transport::connect(host).await?,
            options,
            opening: Mutex::new(BTreeSet::new()),
            lifecycle: Mutex::new(ClientLifecycle::default()),
            idle: Notify::new(),
            documents: Mutex::new(BTreeMap::new()),
        })))
    }

    /// Local contract introspection requires no running host.
    pub fn describe() -> Value {
        contract()["introspection"]["overview"].clone()
    }

    pub fn describe_command(lookup: &str) -> Result<Value> {
        let index = contract()["introspection"]["lookup"]
            [trim_js_whitespace(lookup).to_lowercase()]
        .as_u64()
        .ok_or_else(|| {
            Error::new(
                "OPERATION_NOT_FOUND",
                "no contract operation matches the requested lookup",
            )
        })?;
        Ok(contract()["introspection"]["overview"]["operations"][index as usize].clone())
    }

    pub fn subscribe(&self) -> Notifications {
        Notifications(self.0.transport.subscribe())
    }
    pub fn state_directory(&self) -> &Path {
        self.0.transport.state_directory()
    }

    pub async fn open(&self, params: &DocOpenParams) -> Result<Document> {
        self.open_with_options(params, OpenOptions::default()).await
    }

    pub async fn open_with_options(
        &self,
        params: &DocOpenParams,
        options: OpenOptions,
    ) -> Result<Document> {
        let _active = self.begin_call()?;
        let mut params = serde_json::to_value(params)
            .map_err(|_| Error::new("INVALID_ARGUMENT", "open parameters could not be encoded"))?;
        let session_id = match params.get("sessionId") {
            Some(value) => value
                .as_str()
                .filter(|id| !id.is_empty())
                .ok_or_else(|| Error::new("INVALID_ARGUMENT", "session ID must be non-empty"))?
                .to_owned(),
            None => fresh_session_id(),
        };
        {
            let mut opening = self.0.opening.lock().expect("SDK opening lock");
            if self.0.transport.has_session(&session_id) || !opening.insert(session_id.clone()) {
                let mut error = Error::new(
                    "SESSION_ALREADY_OPEN",
                    "session ID is already open or opening",
                );
                error.details = Some(Box::new(json!({"sessionId":session_id})));
                return Err(error);
            }
        }
        let _guard = OpenGuard {
            client: self.clone(),
            id: session_id.clone(),
        };
        params["sessionId"] = json!(session_id);
        if options.collaboration_auth.is_some()
            && self.0.transport.options().mode == HostMode::Document
        {
            return Err(Self::unsupported_collaboration_auth(HostMode::Document));
        }
        let features = self.0.transport.features();
        let structured = options.collaboration_auth.is_none()
            && routing::can_open_document(&params, &options.invoke)
            && (self.0.transport.options().mode == HostMode::Document
                || self.0.options.document_rpc)
            && policy()["documentRpcFeatures"]
                .as_array()
                .is_some_and(|required| {
                    required
                        .iter()
                        .all(|feature| features.contains(feature.as_str().unwrap_or("")))
                });
        if self.0.transport.options().mode == HostMode::Document && !structured {
            return Err(Error::new(
                "DOCUMENT_RPC_INPUT_UNSUPPORTED",
                "standalone document host cannot support these open parameters",
            ));
        }
        let result = if structured {
            if self.0.options.default_change_mode.is_some() {
                self.require_feature("defaultChangeMode")?;
            }
            if self.0.options.user_name.is_some() {
                self.require_feature("documentAuthor")?;
            }
            let mut request = json!({"sessionId":session_id,"path":params["doc"]});
            if let Some(mode) = &self.0.options.default_change_mode {
                request["defaultChangeMode"] = json!(mode);
            }
            if let Some(name) = &self.0.options.user_name {
                let mut author = json!({"name":name});
                if let Some(email) = &self.0.options.user_email {
                    author["email"] = json!(email);
                }
                request["author"] = author;
            }
            self.rpc("document.open", request, &options.invoke, true, true, true)
                .await
                .and_then(|response| {
                    routing::map_lifecycle(
                        "doc.open",
                        response,
                        &session_id,
                        params["doc"].as_str(),
                    )
                })
        } else {
            if let Some(name) = &self.0.options.user_name {
                if params.get("userName").is_none() {
                    params["userName"] = json!(name);
                }
            }
            if let Some(email) = &self.0.options.user_email {
                if params.get("userEmail").is_none() {
                    params["userEmail"] = json!(email);
                }
            }
            self.invoke_cli(
                "doc.open",
                params,
                &options.invoke,
                options.collaboration_auth.as_ref(),
            )
            .await
        };
        let result = match result {
            Ok(value) => serde_json::from_value::<DocOpenResult>(value).map_err(Error::decode),
            Err(mut error) => {
                error.operation_id = Some("doc.open".into());
                if ["TIMEOUT", "HOST_DISCONNECTED", "HOST_PROTOCOL_ERROR"]
                    .contains(&error.code.as_str())
                {
                    let _ = self.0.transport.shutdown().await;
                    error.outcome_unknown = true;
                }
                return Err(error);
            }
        };
        let result = match result {
            Ok(result) if result.context_id == session_id && result.active => result,
            invalid => {
                // An invalid open result cannot leave an unreachable native
                // session alive. Retiring also invalidates any sibling handles.
                let _ = self.0.transport.shutdown().await;
                let mut error = match invalid {
                    Err(error) => error,
                    Ok(_) => Error::new(
                        "HOST_PROTOCOL_ERROR",
                        "host returned an invalid open result",
                    ),
                };
                error.operation_id = Some("doc.open".into());
                error.outcome_unknown = true;
                return Err(error);
            }
        };
        let state = self.0.transport.session_state(&session_id);
        if state.load(Ordering::Acquire) != 0 {
            return Err(Error::new(
                "DOCUMENT_CLOSED",
                "session closed while it was opening",
            ));
        }
        self.0
            .lifecycle
            .lock()
            .expect("SDK lifecycle lock")
            .opened_documents = true;
        let persistence = structured.then(|| {
            let persistence = Arc::new(DocumentPersistence {
                state: state.clone(),
                dirty: Mutex::new(DirtyDocument::default()),
            });
            self.0
                .documents
                .lock()
                .expect("SDK document lock")
                .insert(session_id.clone(), persistence.clone());
            self.0
                .lifecycle
                .lock()
                .expect("SDK lifecycle lock")
                .managed_documents = true;
            persistence
        });
        Ok(Document(Arc::new(DocumentInner {
            client: self.clone(),
            session_id,
            state,
            structured,
            open_result: result,
            persistence,
        })))
    }

    pub async fn shutdown(&self) -> Result<()> {
        let mut completion = {
            let mut lifecycle = self.0.lifecycle.lock().expect("SDK lifecycle lock");
            if let Some(completion) = &lifecycle.shutdown {
                completion.clone()
            } else {
                let (result, completion) = watch::channel(None);
                lifecycle.stopping = true;
                lifecycle.shutdown = Some(completion.clone());
                let client = self.clone();
                tokio::spawn(async move {
                    let outcome = client.shutdown_owned().await;
                    if outcome.as_ref().is_err_and(|error| !error.outcome_unknown)
                        && !client.0.transport.is_closed()
                    {
                        let mut lifecycle = client.0.lifecycle.lock().expect("SDK lifecycle lock");
                        lifecycle.stopping = false;
                        lifecycle.shutdown = None;
                    }
                    result.send_replace(Some(outcome));
                });
                completion
            }
        };
        loop {
            if let Some(result) = completion.borrow().clone() {
                return result;
            }
            completion.changed().await.map_err(|_| {
                Error::new(
                    "HOST_SHUTDOWN_FAILED",
                    "managed shutdown owner stopped before confirming completion",
                )
            })?;
        }
    }

    fn begin_call(&self) -> Result<ActiveCall> {
        let mut lifecycle = self.0.lifecycle.lock().expect("SDK lifecycle lock");
        if lifecycle.stopping {
            return Err(Error::new("HOST_DISCONNECTED", "host is disposing"));
        }
        lifecycle.active += 1;
        Ok(ActiveCall(self.0.clone()))
    }

    async fn shutdown_owned(&self) -> Result<()> {
        loop {
            let idle = self.0.idle.notified();
            if self.0.lifecycle.lock().expect("SDK lifecycle lock").active == 0 {
                break;
            }
            idle.await;
        }
        self.0.transport.drain().await?;
        if self.0.transport.is_closed() {
            self.0.transport.shutdown().await?;
            let unsaved = self
                .0
                .documents
                .lock()
                .expect("SDK document lock")
                .values()
                .any(|document| {
                    let dirty = document.dirty.lock().expect("SDK dirty document lock");
                    dirty.confirmed || dirty.pending != 0
                });
            if unsaved {
                let mut error = Error::new(
                    "HOST_DISCONNECTED",
                    "host disconnected before dirty documents could be saved",
                );
                error.details = Some(Box::new(json!({
                    "path": self.state_directory(),
                    "sourceSaveConfirmed": false
                })));
                return Err(error);
            }
            return Ok(());
        }
        if !self
            .0
            .lifecycle
            .lock()
            .expect("SDK lifecycle lock")
            .managed_documents
        {
            if !self
                .0
                .lifecycle
                .lock()
                .expect("SDK lifecycle lock")
                .opened_documents
            {
                return self.0.transport.shutdown().await;
            }
            return self.0.transport.shutdown_checked().await.map_err(|error| {
                let mut failure = Error::new(
                    "HOST_SHUTDOWN_FAILED",
                    "host shutdown failed; its state directory has been retained",
                );
                failure.outcome_unknown = error.outcome_unknown;
                failure.exit_code = error.exit_code;
                failure.details = Some(Box::new(json!({
                    "path": self.state_directory(),
                    "exitCode": failure.exit_code.as_deref(),
                    "cause": {"code": error.code, "details": error.details}
                })));
                failure
            });
        }
        let documents: Vec<_> = self
            .0
            .documents
            .lock()
            .expect("SDK document lock")
            .iter()
            .map(|(id, state)| (id.clone(), state.clone()))
            .collect();
        for (id, document) in &documents {
            if document.state.load(Ordering::Acquire) != 0 {
                continue;
            }
            let dirty = {
                let dirty = document.dirty.lock().expect("SDK dirty document lock");
                dirty.confirmed || dirty.pending != 0
            };
            if !dirty {
                continue;
            }
            self.0.transport.retain_state();
            if !self
                .0
                .transport
                .features()
                .contains(policy()["features"]["sourceSave"].as_str().unwrap_or(""))
            {
                let mut error = Error::new(
                    "DOCUMENT_RPC_INPUT_UNSUPPORTED",
                    "host cannot persist dirty source documents",
                );
                error.details = Some(Box::new(
                    json!({"feature":policy()["features"]["sourceSave"]}),
                ));
                return Err(error);
            }
            let generation = document.save_generation();
            let response = self
                .0
                .transport
                .request(
                    "document.save",
                    json!({"sessionId":id}),
                    None,
                    self.0.transport.options().request_timeout.saturating_add(
                        Duration::from_millis(
                            policy()["watchdogHeadroomMs"].as_u64().unwrap_or(5000),
                        ),
                    ),
                    true,
                    false,
                )
                .await;
            let result = match response {
                Err(error) => Err(error),
                Ok(response) if response["inPlace"] != true => {
                    let mut error =
                        Error::new("HOST_PROTOCOL_ERROR", "host did not confirm a source save");
                    error.outcome_unknown = true;
                    Err(error)
                }
                Ok(response) => {
                    routing::map_lifecycle("doc.save", response, id, None).map_err(|mut error| {
                        error.outcome_unknown = true;
                        error
                    })
                }
            };
            if let Err(error) = result {
                if error.outcome_unknown {
                    let _ =
                        document
                            .state
                            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| {
                                (state < 2).then_some(3)
                            });
                    let _ = self.0.transport.retire_preserving_state().await;
                }
                return Err(error);
            }
            if let Some(generation) = generation {
                document.saved(generation);
            }
        }
        let active_documents: Vec<_> = documents
            .iter()
            .filter(|(_, document)| document.state.load(Ordering::Acquire) == 0)
            .map(|(_, document)| document.state.clone())
            .collect();
        self.0.transport.shutdown_checked().await?;
        for state in active_documents {
            let _ = state.fetch_update(Ordering::AcqRel, Ordering::Acquire, |value| {
                matches!(value, 0 | 3).then_some(2)
            });
        }
        Ok(())
    }

    fn remove_document(&self, id: &str, state: &Arc<AtomicU8>) {
        let mut documents = self.0.documents.lock().expect("SDK document lock");
        if documents
            .get(id)
            .is_some_and(|entry| Arc::ptr_eq(&entry.state, state))
        {
            documents.remove(id);
        }
    }

    pub fn presets(&self) -> DocApiPreset<'_, Self> {
        DocApiPreset::new(self)
    }

    fn require_feature(&self, key: &str) -> Result<()> {
        if !self
            .0
            .transport
            .features()
            .contains(policy()["features"][key].as_str().unwrap_or(""))
        {
            return Err(Error::new(
                "CAPABILITY_UNSUPPORTED",
                "connected host lacks the requested optional capability",
            ));
        }
        Ok(())
    }

    fn unsupported_collaboration_auth(mode: HostMode) -> Error {
        let mut error = Error::new(
            "CAPABILITY_UNSUPPORTED",
            if mode == HostMode::Document {
                "Per-open collaboration authentication is not supported by the standalone document host."
            } else {
                "The connected CLI host does not support per-open collaboration authentication."
            },
        );
        error.details = Some(Box::new(json!({
            "feature": policy()["features"]["collaborationAuthPerOpen"],
        })));
        error.operation_id = Some("doc.open".into());
        error
    }

    async fn rpc(
        &self,
        method: &str,
        params: Value,
        options: &InvokeOptions,
        mutates: bool,
        open: bool,
        supports_deadline: bool,
    ) -> Result<Value> {
        self.rpc_tracked(
            method,
            params,
            options,
            mutates,
            open,
            supports_deadline,
            None,
        )
        .await
    }

    #[allow(clippy::too_many_arguments)]
    async fn rpc_tracked(
        &self,
        method: &str,
        params: Value,
        options: &InvokeOptions,
        mutates: bool,
        open: bool,
        supports_deadline: bool,
        state: Option<Arc<CallState>>,
    ) -> Result<Value> {
        let effective = options
            .timeout
            .unwrap_or(self.0.transport.options().request_timeout);
        let headroom =
            Duration::from_millis(policy()["watchdogHeadroomMs"].as_u64().unwrap_or(5000));
        if options.timeout.is_some() {
            self.require_feature("requestTimeout")?;
            if !supports_deadline {
                return Err(Error::new(
                    "DOCUMENT_RPC_INPUT_UNSUPPORTED",
                    "this native operation does not support an explicit request deadline",
                ));
            }
        }
        let deadline = (supports_deadline
            && self.0.transport.features().contains(
                policy()["features"]["requestTimeout"]
                    .as_str()
                    .unwrap_or(""),
            ))
        .then_some(effective);
        self.0
            .transport
            .request_tracked(
                method,
                params,
                deadline,
                effective.saturating_add(headroom),
                mutates,
                open,
                state.unwrap_or_default(),
            )
            .await
    }

    fn reject_unsupported_dispatch_preview(
        &self,
        id: &str,
        options: &InvokeOptions,
        full: bool,
    ) -> Result<()> {
        if self.0.transport.options().mode == HostMode::Cli
            && id == "doc.preset.dispatch"
            && options.dry_run == Some(true)
            && !full
        {
            let mut error = Error::new(
                "CAPABILITY_UNSUPPORTED",
                "The core preset dispatcher does not support dryRun previews.",
            );
            error.operation_id = Some(id.into());
            return Err(error);
        }
        Ok(())
    }

    async fn invoke_cli(
        &self,
        id: &str,
        params: Value,
        options: &InvokeOptions,
        auth: Option<&CollaborationAuth>,
    ) -> Result<Value> {
        self.invoke_cli_tracked(id, params, options, auth, None)
            .await
    }

    async fn invoke_cli_tracked(
        &self,
        id: &str,
        params: Value,
        options: &InvokeOptions,
        auth: Option<&CollaborationAuth>,
        state: Option<Arc<CallState>>,
    ) -> Result<Value> {
        if self.0.transport.options().mode != HostMode::Cli {
            if [
                "doc.preset.list",
                "doc.preset.getCatalog",
                "doc.preset.getTools",
                "doc.preset.getSystemPrompt",
                "doc.preset.getMcpPrompt",
            ]
            .contains(&id)
            {
                if !self
                    .0
                    .transport
                    .features()
                    .contains("document.preset.full.v1")
                {
                    let mut error = Error::new(
                        "CAPABILITY_UNSUPPORTED",
                        "host cannot provide preset metadata",
                    );
                    error.details = Some(Box::new(json!({"feature":"document.preset.full.v1"})));
                    return Err(error);
                }
                if auth.is_some() {
                    return Err(Error::new(
                        "DOCUMENT_RPC_INPUT_UNSUPPORTED",
                        "preset metadata does not accept authentication",
                    ));
                }
                if options
                    .stdin_bytes
                    .as_ref()
                    .is_some_and(|bytes| bytes.len() > 32 * 1024 * 1024)
                {
                    return Err(Error::new(
                        "PAYLOAD_TOO_LARGE",
                        "stdin exceeds the host's 32 MiB limit",
                    ));
                }
                validate_json(&params)?;
                let operation = operation(id)?;
                let params = Value::Object(routing::normalize_aliases(operation, &params)?);
                routing::argv(
                    operation,
                    &params,
                    options,
                    self.0.options.default_change_mode.as_deref(),
                )?;
                return self
                    .rpc_tracked(
                        "preset.metadata",
                        json!({"operationId":id,"input":params}),
                        options,
                        false,
                        false,
                        true,
                        state,
                    )
                    .await;
            }
            return Err(Error::new(
                "DOCUMENT_RPC_INPUT_UNSUPPORTED",
                "operation requires the CLI host",
            ));
        }
        let operation = operation(id)?;
        if auth.is_some() {
            if id != "doc.open" {
                return Err(Error::new(
                    "INVALID_ARGUMENT",
                    "authentication is scoped to document open",
                ));
            }
            self.require_feature("collaborationAuthPerOpen")
                .map_err(|_| Self::unsupported_collaboration_auth(HostMode::Cli))?;
        }
        if options
            .stdin_bytes
            .as_ref()
            .is_some_and(|bytes| bytes.len() > 32 * 1024 * 1024)
        {
            return Err(Error::new(
                "PAYLOAD_TOO_LARGE",
                "stdin exceeds the host's 32 MiB limit",
            ));
        }
        validate_json(&params)?;
        let mut request = json!({"argv":routing::argv(operation,&params,options,self.0.options.default_change_mode.as_deref())?,"stdinBase64":options.stdin_bytes.as_ref().map(|bytes|BASE64.encode(bytes)).unwrap_or_default()});
        if let Some(auth) = auth {
            request["collaborationAuth"] = auth.json();
        }
        let response = self
            .rpc_tracked(
                "cli.invoke",
                request,
                options,
                operation["mutates"] == true,
                id == "doc.open",
                true,
                state,
            )
            .await?;
        routing::map_cli(operation, response).map_err(|mut error| {
            error.outcome_unknown = operation["mutates"] == true;
            error
        })
    }

    pub async fn get_tools(&self, provider: &str) -> Result<Value> {
        let _active = self.begin_call()?;
        self.invoke_cli(
            "doc.preset.getTools",
            json!({"provider":provider,"preset":"core"}),
            &InvokeOptions::default(),
            None,
        )
        .await
    }

    pub async fn get_system_prompt(&self) -> Result<Value> {
        let _active = self.begin_call()?;
        self.invoke_cli(
            "doc.preset.getSystemPrompt",
            json!({"preset":"core"}),
            &InvokeOptions::default(),
            None,
        )
        .await
    }
}

impl Invoker for Client {
    async fn invoke(&self, id: &str, params: Value, options: InvokeOptions) -> Result<Value> {
        let _active = self.begin_call()?;
        if ![
            "doc.preset.list",
            "doc.preset.getCatalog",
            "doc.preset.getTools",
            "doc.preset.getSystemPrompt",
            "doc.preset.getMcpPrompt",
        ]
        .contains(&id)
            || operation(id)?["requiresDocumentContext"] != false
            || operation(id)?["mutates"] != false
        {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "client only invokes sessionless preset getters",
            ));
        }
        self.invoke_cli(id, params, &options, None).await
    }
}

fn operation(id: &str) -> Result<&'static Value> {
    contract()["operations"].get(id).ok_or_else(|| {
        Error::new(
            "OPERATION_NOT_FOUND",
            "operation is not in the generated contract",
        )
    })
}

fn fresh_session_id() -> String {
    static NEXT: AtomicU64 = AtomicU64::new(1);
    format!(
        "rust-{}-{}-{}",
        std::process::id(),
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos(),
        NEXT.fetch_add(1, Ordering::Relaxed)
    )
}

fn validate_json(value: &Value) -> Result<()> {
    let mut stack = vec![(value, 0)];
    while let Some((value, depth)) = stack.pop() {
        if depth > 128 {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "request JSON exceeds the supported nesting depth",
            ));
        }
        match value {
            Value::Number(number)
                if number.is_i64()
                    && number.as_i64().is_some_and(|n| {
                        !(-crate::wire::Integer::MAX..=crate::wire::Integer::MAX).contains(&n)
                    }) =>
            {
                return Err(Error::new(
                    "INVALID_ARGUMENT",
                    "JSON integer cannot be represented exactly by the host",
                ))
            }
            Value::Number(number)
                if number.is_u64()
                    && number
                        .as_u64()
                        .is_some_and(|n| n > crate::wire::Integer::MAX as u64) =>
            {
                return Err(Error::new(
                    "INVALID_ARGUMENT",
                    "JSON integer cannot be represented exactly by the host",
                ))
            }
            Value::Array(values) => stack.extend(values.iter().map(|value| (value, depth + 1))),
            Value::Object(values) => stack.extend(values.values().map(|value| (value, depth + 1))),
            _ => {}
        }
    }
    Ok(())
}

struct DocumentInner {
    client: Client,
    session_id: String,
    // Active, closing, closed, or outcome unknown. All clones share this state.
    state: Arc<AtomicU8>,
    structured: bool,
    open_result: DocOpenResult,
    persistence: Option<Arc<DocumentPersistence>>,
}

impl Drop for DocumentInner {
    fn drop(&mut self) {
        if self.state.load(Ordering::Acquire) >= 2 {
            self.client.remove_document(&self.session_id, &self.state);
            self.client
                .0
                .transport
                .remove_session(&self.session_id, &self.state);
        }
    }
}

#[derive(Clone)]
pub struct Document(Arc<DocumentInner>);

struct CloseGuard {
    state: Arc<AtomicU8>,
    transmission: Arc<CallState>,
}
impl Drop for CloseGuard {
    fn drop(&mut self) {
        let _ = self.state.compare_exchange(
            1,
            if self.transmission.cancel_before_send() {
                0
            } else {
                3
            },
            Ordering::AcqRel,
            Ordering::Acquire,
        );
    }
}

impl Document {
    pub fn session_id(&self) -> &str {
        &self.0.session_id
    }
    pub fn open_result(&self) -> &DocOpenResult {
        &self.0.open_result
    }
    pub fn api(&self) -> DocApi<'_, Self> {
        DocApi::new(self)
    }

    pub async fn close(&self, params: &DocCloseBoundParams) -> Result<DocCloseResult> {
        self.api().close(params).await
    }

    pub async fn replace_file(
        &self,
        path: impl AsRef<Path>,
        options: InvokeOptions,
    ) -> Result<Value> {
        self.active()?;
        let _active = self.0.client.begin_call()?;
        self.replace_path(path.as_ref(), options, None).await
    }

    /// The transport owns staged bytes until the native reply or confirmed exit,
    /// including when the caller drops this future after transmission.
    pub async fn replace_bytes(&self, bytes: &[u8], options: InvokeOptions) -> Result<Value> {
        use std::io::Write;
        if bytes.len() > 32 * 1024 * 1024 {
            return Err(Error::new(
                "PAYLOAD_TOO_LARGE",
                "replacement exceeds 32 MiB",
            ));
        }
        self.active()?;
        let _active = self.0.client.begin_call()?;
        let mut file = tempfile::NamedTempFile::new_in(self.0.client.state_directory())
            .map_err(|_| Error::new("IO_ERROR", "could not stage replacement bytes"))?;
        file.write_all(bytes)
            .map_err(|_| Error::new("IO_ERROR", "could not stage replacement bytes"))?;
        let asset = file.into_temp_path();
        let path = asset.to_path_buf();
        self.replace_path(&path, options, Some(CallState::with_asset(asset)))
            .await
    }

    async fn replace_path(
        &self,
        path: &Path,
        options: InvokeOptions,
        state: Option<Arc<CallState>>,
    ) -> Result<Value> {
        self.active()?;
        self.0.client.require_feature("replaceFile")?;
        let path = path
            .to_str()
            .ok_or_else(|| Error::new("INVALID_ARGUMENT", "replacement path must be UTF-8"))?;
        let transmission = state.unwrap_or_default();
        let mut mutation = self.0.persistence.as_ref().map(|document| {
            MutationGuard::new(
                document.clone(),
                transmission.clone(),
                &self.0.client.0.transport,
            )
        });
        let response = self
            .0
            .client
            .rpc_tracked(
                "document.replaceFile",
                json!({"sessionId":self.0.session_id,"path":path}),
                &options,
                true,
                false,
                false,
                Some(transmission),
            )
            .await;
        if let Some(mutation) = &mut mutation {
            mutation.settle(
                response
                    .as_ref()
                    .map_or_else(|error| error.outcome_unknown, |_| true),
            );
        }
        let value = response?;
        if value["sessionId"] != self.0.session_id
            || value["replaced"] != true
            || !value["byteLength"].is_number()
            || !value["revision"].is_string()
        {
            let mut error = Error::new(
                "HOST_PROTOCOL_ERROR",
                "host returned an invalid replacement result",
            );
            error.outcome_unknown = true;
            return Err(error);
        }
        Ok(
            json!({"contextId":self.0.session_id,"runtime":"v2","replaced":true,"document":{"byteLength":value["byteLength"],"revision":value["revision"]}}),
        )
    }

    pub(crate) async fn check_core_revision(
        &self,
        tool_name: &str,
        args: &Value,
        options: &InvokeOptions,
        configured: Option<&crate::generated::types::DocPresetGetToolsParams>,
    ) -> Result<Option<Value>> {
        let Some(expected) = options.expected_revision.as_deref() else {
            return Ok(None);
        };
        let params = serde_json::from_value(
            serde_json::json!({"offset":crate::wire::Integer::MAX,"limit":1}),
        )
        .map_err(Error::decode)?;
        let page = self.api().blocks().list(&params).await?;
        if page.revision.is_empty() || page.revision == "unknown" {
            let cli_reader = !self.0.structured
                && self.0.open_result.session_type == "local"
                && self
                    .0
                    .client
                    .0
                    .transport
                    .features()
                    .contains("cli.invoke.localExecutionFacts");
            return Err(if cli_reader {
                Error::new(
                    "HOST_PROTOCOL_ERROR",
                    "CLI returned an invalid execution revision.",
                )
            } else {
                Error::new(
                    "EVIDENCE_UNAVAILABLE",
                    "Required document revision is unavailable.",
                )
            });
        }
        if expected == page.revision {
            return Ok(None);
        }
        if tool_name == "superdoc_perform_action" {
            return crate::agent::stale_action_revision_receipt(args, &page.revision, configured)
                .map(Some);
        }
        Err(Error::new(
            "REVISION_CONFLICT",
            crate::agent::stale_revision_message(tool_name),
        ))
    }

    pub async fn dispatch_tool(
        &self,
        name: &str,
        mut args: Value,
        mut options: InvokeOptions,
    ) -> Result<Value> {
        self.active()?;
        crate::agent::prepare_core_dispatch(name, &mut args, &mut options, None)?;
        self.0.client.reject_unsupported_dispatch_preview(
            "doc.preset.dispatch",
            &options,
            false,
        )?;
        if let Some(receipt) = self
            .check_core_revision(name, &args, &options, None)
            .await?
        {
            return Ok(json!({ "result": receipt }));
        }
        let _active = self.0.client.begin_call()?;
        let transmission = Arc::new(CallState::default());
        let mut mutation = self.0.persistence.as_ref().map(|document| {
            MutationGuard::new(
                document.clone(),
                transmission.clone(),
                &self.0.client.0.transport,
            )
        });
        let result = self
            .0
            .client
            .invoke_cli_tracked(
                "doc.preset.dispatch",
                json!({"sessionId":self.0.session_id,"toolName":name,"args":args,"preset":"core"}),
                &options,
                None,
                Some(transmission),
            )
            .await;
        if let Some(mutation) = &mut mutation {
            mutation.settle(result.as_ref().map_or_else(
                |error| error.outcome_unknown,
                |value| value["success"] != false,
            ));
        }
        result
    }

    fn active(&self) -> Result<()> {
        match self.0.state.load(Ordering::Acquire) {
            0 if !self.0.client.0.transport.is_closed() => Ok(()),
            3 => Err(Error::new(
                "DOCUMENT_STATE_UNKNOWN",
                "document state is uncertain; use a new client to reopen verified saved content",
            )),
            _ => Err(Error::new(
                "DOCUMENT_CLOSED",
                "document handle is closed or closing",
            )),
        }
    }
}

impl Invoker for Document {
    fn handle_decode_error(&self, error: Error) -> Error {
        if error.outcome_unknown {
            let _ = self
                .0
                .state
                .fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| {
                    (state < 2).then_some(3)
                });
        }
        error
    }

    async fn invoke(
        &self,
        id: &str,
        mut params: Value,
        mut options: InvokeOptions,
    ) -> Result<Value> {
        self.active().map_err(|mut error| {
            if error.code == "DOCUMENT_CLOSED" && operation(id).is_ok() {
                error.message =
                    format!("Document handle is closed or closing; cannot invoke {id}.");
                error.operation_id = Some(id.into());
                error.details = Some(Box::new(json!({
                    "sessionId": self.0.session_id,
                    "operationId": id
                })));
            }
            error
        })?;
        let _active = self.0.client.begin_call()?;
        let operation = operation(id)?;
        if operation["sdkSurface"] != "document" {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "operation is not on the public document surface",
            ));
        }
        self.0.client.reject_unsupported_dispatch_preview(
            id,
            &options,
            params.get("preset") == Some(&json!("full")),
        )?;
        let binds_session = operation["params"]
            .as_array()
            .is_some_and(|params| params.iter().any(|param| param["name"] == "sessionId"));
        if binds_session {
            params["sessionId"] = json!(self.0.session_id);
        }
        let closing = id == "doc.close";
        let transmission = closing.then(|| Arc::new(CallState::default()));
        let _close = if closing {
            self.0
                .state
                .compare_exchange(0, 1, Ordering::AcqRel, Ordering::Acquire)
                .map_err(|_| Error::new("DOCUMENT_CLOSED", "document is already closing"))?;
            Some(CloseGuard {
                state: self.0.state.clone(),
                transmission: transmission.clone().expect("close transmission state"),
            })
        } else {
            None
        };
        if options.change_mode.is_none()
            && operation["params"]
                .as_array()
                .is_some_and(|params| params.iter().any(|param| param["name"] == "changeMode"))
        {
            options.change_mode = self.0.client.0.options.default_change_mode.clone();
        }
        let preserves_review =
            params.get("mode").is_none() || params["mode"] == "review-preserving";
        let save_generation = self
            .0
            .persistence
            .as_ref()
            .filter(|_| id == "doc.save")
            .and_then(|document| document.save_generation());
        let mut mutation = None;
        let mut observed_full_change = None;
        let result: Result<Value> = async {
            if self.0.structured && binds_session {
                let (method, request) = routing::document_request(
                    &self.0.session_id,
                    operation,
                    &params,
                    &options,
                    self.0.client.0.transport.features(),
                )?;
                let mutates = (method == "document.invoke" || method == "document.preset.dispatch")
                    && operation["mutates"] == true
                    && request["options"]["dryRun"] != true;
                let transmission = transmission.clone().unwrap_or_default();
                if mutates {
                    mutation = self.0.persistence.as_ref().map(|document| {
                        MutationGuard::new(
                            document.clone(),
                            transmission.clone(),
                            &self.0.client.0.transport,
                        )
                    });
                }
                let safe_deadline = method != "document.save"
                    && (operation["mutates"] == false
                        || operation["idempotency"] == "idempotent"
                        || request["options"]["dryRun"] == true
                        || method != "document.invoke");
                validate_json(&request)?;
                let response = self
                    .0
                    .client
                    .rpc_tracked(
                        &method,
                        request,
                        &options,
                        operation["mutates"] == true,
                        false,
                        safe_deadline,
                        Some(transmission.clone()),
                    )
                    .await;
                let response = if method == "document.preset.dispatch" {
                    observed_full_change = transmission.session_revision_changed();
                    match response {
                        Ok(mut response) => {
                            let valid = response.as_object().is_some_and(|object| {
                                let present = object.get("resultPresent").and_then(Value::as_bool);
                                observed_full_change.is_some()
                                    && present.is_some()
                                    && object.contains_key("result") == present.unwrap_or(false)
                                    && object.len() == if present == Some(true) { 3 } else { 2 }
                            });
                            if !valid {
                                let mut error = Error::new(
                                    "HOST_PROTOCOL_ERROR",
                                    "Full preset response has an invalid completion envelope",
                                );
                                error.outcome_unknown = true;
                                Err(error)
                            } else {
                                if let Some(object) = response.as_object_mut() {
                                    object.remove("sessionRevisionChanged");
                                    object.remove("resultPresent");
                                }
                                Ok(response)
                            }
                        }
                        Err(mut error) => {
                            if observed_full_change.is_none() && transmission.transmission_started()
                            {
                                error.outcome_unknown = true;
                            }
                            Err(error)
                        }
                    }
                } else {
                    response
                };
                response.and_then(|response| {
                    if id == "doc.save" || closing {
                        routing::map_lifecycle(id, response, &self.0.session_id, None).map_err(
                            |mut error| {
                                error.outcome_unknown = operation["mutates"] == true;
                                error
                            },
                        )
                    } else {
                        Ok(response)
                    }
                })
            } else {
                self.0
                    .client
                    .invoke_cli_tracked(id, params, &options, None, transmission)
                    .await
            }
        }
        .await;
        let result = result.and_then(|value| routing::map_domain_result(operation, value));
        if let Some(mutation) = &mut mutation {
            mutation.settle(result.as_ref().map_or_else(
                |error| error.outcome_unknown || observed_full_change.unwrap_or(false),
                |value| observed_full_change.unwrap_or_else(|| value["success"] != false),
            ));
        }
        match result {
            Ok(value) => {
                if closing {
                    if value["closed"] != true || value["contextId"] != self.0.session_id {
                        let mut error = Error::new(
                            "HOST_PROTOCOL_ERROR",
                            "close was not confirmed by the host",
                        );
                        error.outcome_unknown = true;
                        return Err(error);
                    }
                    self.0.state.store(2, Ordering::Release);
                    self.0
                        .client
                        .remove_document(&self.0.session_id, &self.0.state);
                    self.0
                        .client
                        .0
                        .transport
                        .remove_session(&self.0.session_id, &self.0.state);
                } else if id == "doc.save" && preserves_review {
                    if let (Some(document), Some(generation)) =
                        (&self.0.persistence, save_generation)
                    {
                        document.saved(generation);
                    }
                }
                Ok(value)
            }
            Err(mut error) => {
                error.operation_id = Some(id.into());
                if closing && !error.outcome_unknown && !self.0.client.0.transport.is_closed() {
                    let _ =
                        self.0
                            .state
                            .compare_exchange(1, 0, Ordering::AcqRel, Ordering::Acquire);
                }
                Err(error)
            }
        }
    }
}

/// Lag is explicit; callers must resynchronize rather than assume every remote
/// change notification was observed. Session invalidation runs in the reader's
/// owner independently of this public notification stream.
pub struct Notifications(broadcast::Receiver<Value>);
impl Notifications {
    pub async fn recv(&mut self) -> Result<Value> {
        self.0.recv().await.map_err(|error| match error {
            broadcast::error::RecvError::Lagged(_) => Error::new(
                "NOTIFICATIONS_LAGGED",
                "resynchronize document state before using notification-derived caches",
            ),
            broadcast::error::RecvError::Closed => {
                Error::new("HOST_DISCONNECTED", "notification stream closed")
            }
        })
    }
}