magi-code 0.77.1

Repository-aware CLI coding agent for terminal work
Documentation
use super::{
    protocol::{ServiceErrorCode, ServiceRequest},
    runtime::ServiceRuntime,
};
use crate::sessions::{Session, SessionEvent, SessionEventKind, SessionWriterLease};
use serde::Deserialize;
use serde_json::{Value, json};
use std::{collections::HashMap, sync::Arc};

pub(super) const MAX_ATTACHMENTS: usize = 16;
use crate::sessions::FRONTEND_PAGE_LIMIT as MAX_PAGE_ITEMS;

enum AttachmentKind {
    ImplicitTurn,
    ExplicitOpen,
    DaemonClaim,
}

struct Attachment {
    session: Session,
    _writer: SessionWriterLease,
    explicit: bool,
    closing: bool,
}

pub(super) struct ServiceSessions {
    runtime: Arc<ServiceRuntime>,
    attachments: HashMap<String, Attachment>,
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct PageParams {
    #[serde(default)]
    after: Option<String>,
    #[serde(default = "default_limit")]
    limit: usize,
}
fn default_limit() -> usize {
    MAX_PAGE_ITEMS
}

impl ServiceSessions {
    pub(super) fn new(runtime: Arc<ServiceRuntime>) -> Self {
        Self {
            runtime,
            attachments: HashMap::new(),
        }
    }

    pub(super) fn for_turn(&mut self, id: Option<&str>) -> Result<Session, ServiceErrorCode> {
        self.attach(id, AttachmentKind::ImplicitTurn)
    }

    fn attach(
        &mut self,
        id: Option<&str>,
        kind: AttachmentKind,
    ) -> Result<Session, ServiceErrorCode> {
        let explicit = !matches!(kind, AttachmentKind::ImplicitTurn);
        if let Some(id) = id {
            validate_id(id)?;
            if let Some(attachment) = self.attachments.get_mut(id) {
                if attachment.closing {
                    return Err(ServiceErrorCode::SessionBusy);
                }
                attachment.explicit |= explicit;
                return Ok(attachment.session.clone());
            }
        }
        if self.attachments.len() >= MAX_ATTACHMENTS {
            return Err(ServiceErrorCode::LimitExceeded);
        }
        let session = match id {
            Some(id) => self.runtime.session_manager.open_existing(id),
            None => self.runtime.session_manager.create(),
        }
        .map_err(|_| ServiceErrorCode::SessionUnavailable)?;
        let writer = match kind {
            AttachmentKind::DaemonClaim => session.try_daemon_writer(),
            AttachmentKind::ImplicitTurn | AttachmentKind::ExplicitOpen => {
                session.try_frontend_writer()
            }
        }
        .map_err(|_| ServiceErrorCode::SessionUnavailable)?
        .ok_or(ServiceErrorCode::SessionBusy)?;
        let session = session
            .activate()
            .map_err(|_| ServiceErrorCode::SessionBusy)?;
        if id.is_none() {
            let event = SessionEvent::new_kind(
                SessionEventKind::Diagnostic,
                session.id().to_owned(),
                self.runtime.cwd.clone(),
                json!({"lifecycle": "created"}),
            );
            session
                .append_with_outcome(&event)
                .map_err(|_| ServiceErrorCode::SessionUnavailable)?;
        }
        self.attachments.insert(
            session.id().to_owned(),
            Attachment {
                session: session.clone(),
                _writer: writer,
                explicit,
                closing: false,
            },
        );
        Ok(session)
    }

    pub(super) fn operation(
        &mut self,
        request: &ServiceRequest,
        active: bool,
    ) -> Result<Value, ServiceErrorCode> {
        match request.method.as_str() {
            "session.list" => {
                if request.session_id.is_some() {
                    return Err(ServiceErrorCode::InvalidPayload);
                }
                let params = page_params(&request.payload)?;
                if let Some(after) = &params.after {
                    validate_id(after)?;
                }
                let page: crate::sessions::FrontendSessionList = self
                    .runtime
                    .session_manager
                    .frontend_list(params.after.as_deref(), params.limit)
                    .map_err(|error| {
                        if error
                            .downcast_ref::<crate::sessions::BoundedReadError>()
                            .is_some()
                        {
                            ServiceErrorCode::LimitExceeded
                        } else {
                            ServiceErrorCode::SessionUnavailable
                        }
                    })?;
                serde_json::to_value(page).map_err(|_| ServiceErrorCode::InternalError)
            }
            "session.create" => {
                empty_params(&request.payload)?;
                if request.session_id.is_some() {
                    return Err(ServiceErrorCode::InvalidPayload);
                }
                let session = self.attach(None, AttachmentKind::ExplicitOpen)?;
                Ok(json!({"session_id": session.id(), "status": "open"}))
            }
            "session.open" => {
                empty_params(&request.payload)?;
                let id = request_id(request)?;
                let session = self.attach(Some(id), AttachmentKind::ExplicitOpen)?;
                Ok(json!({"session_id": session.id(), "status": "open"}))
            }
            "session.replay" => {
                let id = request_id(request)?;
                let params = page_params(&request.payload)?;
                // Replay is read-only and does not acquire or require the writer attachment.
                let session = self
                    .runtime
                    .session_manager
                    .open_existing(id)
                    .map_err(|_| ServiceErrorCode::SessionUnavailable)?;
                let page: crate::sessions::FrontendReplayPage = session
                    .frontend_replay(params.after.as_deref(), params.limit)
                    .map_err(|error| {
                        if error
                            .downcast_ref::<crate::sessions::FrontendSnapshotBusy>()
                            .is_some()
                        {
                            ServiceErrorCode::SessionBusy
                        } else {
                            ServiceErrorCode::SessionUnavailable
                        }
                    })?;
                serde_json::to_value(page).map_err(|_| ServiceErrorCode::InternalError)
            }
            "session.close" => {
                empty_params(&request.payload)?;
                let id = request_id(request)?;
                if active {
                    let attachment = self
                        .attachments
                        .get_mut(id)
                        .ok_or(ServiceErrorCode::SessionUnavailable)?;
                    attachment.closing = true;
                    Ok(json!({"session_id": id, "status": "closing"}))
                } else {
                    // Idempotent for any valid ID, including a previously removed local session.
                    self.attachments.remove(id);
                    Ok(json!({"session_id": id, "status": "closed"}))
                }
            }
            _ => Err(ServiceErrorCode::UnsupportedOperation),
        }
    }

    pub(super) fn claim(&mut self, id: &str) -> Result<(), ServiceErrorCode> {
        let already_attached = self.attachments.contains_key(id);
        let session = self.attach(Some(id), AttachmentKind::DaemonClaim)?;
        if !already_attached {
            let replay = session.read_events_tolerant_bounded(
                crate::context::REPLAY_JSONL_MAX_LINES,
                crate::context::REPLAY_JSONL_MAX_BYTES,
            );
            if !matches!(replay, Ok(ref replay) if replay.diagnostics.is_empty()) {
                self.attachments.remove(id);
                return Err(ServiceErrorCode::SessionUnavailable);
            }
        }
        Ok(())
    }

    pub(super) fn detach(&mut self, id: &str, active: bool) {
        if active {
            if let Some(attachment) = self.attachments.get_mut(id) {
                attachment.explicit = false;
            }
        } else {
            self.attachments.remove(id);
        }
    }
    pub(super) fn finish_turn(&mut self, id: &str) {
        if self
            .attachments
            .get(id)
            .is_some_and(|attachment| !attachment.explicit || attachment.closing)
        {
            self.attachments.remove(id);
        }
    }

    pub(super) fn clear(&mut self) {
        self.attachments.clear();
    }
}

fn validate_id(id: &str) -> Result<(), ServiceErrorCode> {
    if id.len() > super::protocol::MAX_ID_BYTES
        || crate::sessions::validate_session_id(id.to_owned()).is_err()
    {
        return Err(ServiceErrorCode::InvalidSessionId);
    }
    Ok(())
}
fn request_id(request: &ServiceRequest) -> Result<&str, ServiceErrorCode> {
    let id = request
        .session_id
        .as_deref()
        .ok_or(ServiceErrorCode::InvalidPayload)?;
    validate_id(id)?;
    Ok(id)
}
fn empty_params(value: &Value) -> Result<(), ServiceErrorCode> {
    serde_json::from_value::<super::protocol::EmptyParams>(value.clone())
        .map(|_| ())
        .map_err(|_| ServiceErrorCode::InvalidPayload)
}
fn page_params(value: &Value) -> Result<PageParams, ServiceErrorCode> {
    let params: PageParams =
        serde_json::from_value(value.clone()).map_err(|_| ServiceErrorCode::InvalidPayload)?;
    if !(1..=MAX_PAGE_ITEMS).contains(&params.limit)
        || params.after.as_ref().is_some_and(|after| after.len() > 128)
    {
        return Err(ServiceErrorCode::LimitExceeded);
    }
    Ok(params)
}