cookie_agent_server/service/

routes.rs

use async_trait::async_trait;
use cookie_agent_engine::{EngineError, session::SessionError};
use cookie_agent_protocol::{
    AgentUsageParams, AgentUsageResult, ApprovalListParams, ApprovalListResult,
    ApprovalRespondErrorCode, ApprovalRespondParams, ApprovalRespondResult, EventsSubscribeParams,
    EventsSubscribeResult, GlobalUsageParams, GlobalUsageResult, McpAuthBeginParams,
    McpAuthBeginResult, McpAuthCancelParams, McpAuthCancelResult, McpServerAddParams,
    McpServerEditParams, McpServerListParams, McpServerListResult, McpServerMutationResult,
    McpServerNameParams, McpServerPersistParams, McpServerSetEnabledParams, ProviderConnectParams,
    ProviderConnectResult, ProviderDisconnectParams, ProviderDisconnectResult, RunCancelParams,
    RunCancelResult, RunRecallSteerParams, RunRecallSteerResult, RunStartParams, RunStartResult,
    RunSteerParams, RunSteerResult, RunToolStdinParams, RunToolStdinResult,
    RuntimeSnapshotGetParams, RuntimeSnapshotResult, ServerContext, ServerFault, ServerProtocol,
    SessionChildrenParams, SessionChildrenResult, SessionCompactParams, SessionCompactResult,
    SessionCreateParams, SessionCreateResult, SessionForkParams, SessionForkResult,
    SessionGetParams, SessionGetResult, SessionGoalGetParams, SessionGoalGetResult,
    SessionGoalLifecycleParams, SessionGoalLifecycleResult, SessionGoalSetParams,
    SessionGoalSetResult, SessionListParams, SessionListResult, SessionPermissionClearParams,
    SessionPermissionGetParams, SessionPermissionGetResult, SessionPermissionMutationResult,
    SessionPermissionSetParams, SessionProducersParams, SessionProducersResult,
    SessionRenameErrorCode, SessionRenameParams, SessionRenameResult, SessionResumeParams,
    SessionResumeResult, SessionRevertParams, SessionRevertResult, SessionSetPermissionModeParams,
    SessionSetPermissionModeResult, SessionTreeParams, SessionTreeResult, SessionTreeUsageResult,
    SessionUsageParams, SessionUsageResult, SkillsGetParams, SkillsGetResult, SkillsListParams,
    SkillsListResult,
};

use super::Server;
use crate::rpc::{RpcFault, engine_fault, run_start_fault};

type Result<T> = std::result::Result<T, ServerFault>;

fn rpc_origin() -> cookie_agent_protocol::EventOrigin {
    cookie_agent_protocol::EventOrigin::new("client:rpc").expect("static event origin is valid")
}

#[async_trait]
impl ServerProtocol for Server {
    async fn connected(&self, context: ServerContext) {
        self.start_runtime_notifications(context.clone());
        self.start_engine_event_notifications(context);
    }

    async fn create_session(&self, params: SessionCreateParams) -> Result<SessionCreateResult> {
        self.engine
            .create_session(params.selection)
            .map(|session| SessionCreateResult { session })
            .map_err(protocol_fault)
    }

    async fn list_sessions(&self, params: SessionListParams) -> Result<SessionListResult> {
        let sessions = self
            .engine
            .list_sessions()
            .into_iter()
            .filter(|session| {
                params
                    .cwd_identity
                    .as_ref()
                    .is_none_or(|cwd| &session.cwd_identity == cwd)
            })
            .collect();
        Ok(SessionListResult { sessions })
    }

    async fn get_session(&self, params: SessionGetParams) -> Result<SessionGetResult> {
        self.engine
            .get_session(params.session_id)
            .map(|session| SessionGetResult { session })
            .map_err(protocol_fault)
    }

    async fn get_session_goal(&self, params: SessionGoalGetParams) -> Result<SessionGoalGetResult> {
        self.engine
            .get_session_goal(params)
            .await
            .map_err(protocol_fault)
    }

    async fn set_session_goal(&self, params: SessionGoalSetParams) -> Result<SessionGoalSetResult> {
        self.engine
            .set_session_goal(params, rpc_origin())
            .await
            .map_err(protocol_fault)
    }

    async fn change_session_goal_lifecycle(
        &self,
        params: SessionGoalLifecycleParams,
    ) -> Result<SessionGoalLifecycleResult> {
        self.engine
            .change_session_goal_lifecycle(params, rpc_origin())
            .await
            .map_err(protocol_fault)
    }

    async fn session_producers(
        &self,
        params: SessionProducersParams,
    ) -> Result<SessionProducersResult> {
        self.engine
            .session_producers(params)
            .await
            .map_err(protocol_fault)
    }

    async fn session_usage(&self, params: SessionUsageParams) -> Result<SessionUsageResult> {
        self.engine
            .session_usage(params.session_id)
            .map_err(protocol_fault)
    }

    async fn session_tree_usage(
        &self,
        params: SessionUsageParams,
    ) -> Result<SessionTreeUsageResult> {
        self.engine
            .session_tree_usage(params.session_id)
            .map_err(|error| RpcFault::session_tree_usage(error).into())
    }

    async fn agent_usage(&self, params: AgentUsageParams) -> Result<AgentUsageResult> {
        Ok(self.engine.agent_usage(params.agent_id))
    }

    async fn global_usage(&self, _: GlobalUsageParams) -> Result<GlobalUsageResult> {
        Ok(self.engine.global_usage())
    }

    async fn session_children(
        &self,
        params: SessionChildrenParams,
    ) -> Result<SessionChildrenResult> {
        Ok(SessionChildrenResult {
            children: self
                .engine
                .children(params.session_id)
                .map_err(protocol_fault)?,
        })
    }

    async fn session_tree(&self, params: SessionTreeParams) -> Result<SessionTreeResult> {
        self.engine
            .tree(params.session_id)
            .map(|tree| SessionTreeResult { tree })
            .map_err(protocol_fault)
    }

    async fn resume_session(&self, params: SessionResumeParams) -> Result<SessionResumeResult> {
        self.engine
            .resume(params.session_id)
            .await
            .map(|session| SessionResumeResult { session })
            .map_err(protocol_fault)
    }

    async fn rename_session(&self, params: SessionRenameParams) -> Result<SessionRenameResult> {
        self.engine
            .rename_session(params.clone(), rpc_origin())
            .await
            .map_err(|error| rename_fault(&params, error).into())
    }

    async fn set_permission_mode(
        &self,
        params: SessionSetPermissionModeParams,
    ) -> Result<SessionSetPermissionModeResult> {
        self.engine
            .set_permission_mode(params.session_id, params.mode)
            .map_err(protocol_fault)?;
        Ok(SessionSetPermissionModeResult {})
    }

    async fn get_session_permissions(
        &self,
        params: SessionPermissionGetParams,
    ) -> Result<SessionPermissionGetResult> {
        self.engine
            .get_session_permissions(params.session_id)
            .map_err(protocol_fault)
    }

    async fn set_session_permission(
        &self,
        params: SessionPermissionSetParams,
    ) -> Result<SessionPermissionMutationResult> {
        self.engine
            .set_session_permission(
                params.session_id,
                params.action,
                params.resource,
                params.effect,
                rpc_origin(),
            )
            .await
            .map_err(protocol_fault)
    }

    async fn clear_session_permission(
        &self,
        params: SessionPermissionClearParams,
    ) -> Result<SessionPermissionMutationResult> {
        self.engine
            .clear_session_permission(
                params.session_id,
                params.action,
                &params.resource,
                rpc_origin(),
            )
            .await
            .map_err(protocol_fault)
    }

    async fn list_skills(&self, params: SkillsListParams) -> Result<SkillsListResult> {
        self.engine
            .list_skills(params.session_id)
            .map_err(protocol_fault)
    }

    async fn get_skill(&self, params: SkillsGetParams) -> Result<SkillsGetResult> {
        self.engine
            .get_skill(params.session_id, &params.name, &params.args)
            .map_err(protocol_fault)
    }

    async fn compact_session(&self, params: SessionCompactParams) -> Result<SessionCompactResult> {
        self.engine
            .compact_session_result(
                params.session_id,
                params.focus.as_ref().map(|focus| focus.as_str()),
                rpc_origin(),
            )
            .await
            .map_err(protocol_fault)
    }

    async fn revert_session(&self, params: SessionRevertParams) -> Result<SessionRevertResult> {
        self.engine
            .revert_session(params.session_id, params.through_seq, rpc_origin())
            .await
            .map_err(protocol_fault)
    }

    async fn fork_session(&self, params: SessionForkParams) -> Result<SessionForkResult> {
        self.engine
            .fork_session(params.session_id, params.through_seq, rpc_origin())
            .await
            .map_err(protocol_fault)
    }

    async fn start_run(&self, params: RunStartParams) -> Result<RunStartResult> {
        match self.engine.start_run(params.clone(), rpc_origin()).await {
            Ok(result) => Ok(result),
            Err(EngineError::RunIdempotencyConflict) => {
                Err(RpcFault::run_start_conflict(&params).into())
            }
            Err(error) => Err(run_start_fault(error).into()),
        }
    }

    async fn steer_run(&self, params: RunSteerParams) -> Result<RunSteerResult> {
        self.engine
            .steer(params.run_id, params.input, rpc_origin())
            .await
            .map_err(protocol_fault)
    }

    async fn recall_steer(&self, params: RunRecallSteerParams) -> Result<RunRecallSteerResult> {
        self.engine
            .recall_steer(params.run_id)
            .await
            .map_err(protocol_fault)
    }

    async fn cancel_run(&self, params: RunCancelParams) -> Result<RunCancelResult> {
        self.engine
            .cancel_run(params.run_id)
            .await
            .map_err(protocol_fault)
    }

    async fn tool_stdin(&self, params: RunToolStdinParams) -> Result<RunToolStdinResult> {
        self.engine.tool_stdin(params).await.map_err(protocol_fault)
    }

    async fn subscribe_events(
        &self,
        params: EventsSubscribeParams,
        context: &ServerContext,
    ) -> Result<EventsSubscribeResult> {
        let (result, receiver) = match self
            .engine
            .subscribe(params.session_id, params.cursor)
            .await
        {
            Ok(subscription) => subscription,
            Err(EngineError::SessionOwnedByAnotherProcess(_)) => {
                // Protocol 17 has no separate snapshot RPC. Return only the replay
                // half of this response and deliberately register no live tail.
                return self
                    .engine
                    .snapshot_events(params.session_id, params.cursor)
                    .map_err(protocol_fault);
            }
            Err(error) => return Err(protocol_fault(error)),
        };
        context.register_session_subscription(params.session_id);
        self.start_event_tail(receiver, context.clone());
        for event in &result.events {
            if let cookie_agent_protocol::EventPayload::ToolCallStarted { start } = &event.payload {
                self.start_output_tail(start.tool_call_id, context.clone());
            }
        }
        Ok(result)
    }

    async fn respond_approval(
        &self,
        params: ApprovalRespondParams,
    ) -> Result<ApprovalRespondResult> {
        self.engine
            .approval_respond(params.clone(), rpc_origin())
            .await
            .map_err(|error| approval_fault(&params, error).into())
    }

    async fn list_approvals(&self, params: ApprovalListParams) -> Result<ApprovalListResult> {
        Ok(self
            .engine
            .list_approvals(params.root_session_id, params.status))
    }

    async fn begin_mcp_auth(&self, params: McpAuthBeginParams) -> Result<McpAuthBeginResult> {
        let authorization_url = self
            .engine
            .begin_mcp_auth(params.server.clone())
            .await
            .map_err(protocol_fault)?;
        Ok(McpAuthBeginResult {
            server: params.server,
            authorization_url,
        })
    }

    async fn cancel_mcp_auth(&self, params: McpAuthCancelParams) -> Result<McpAuthCancelResult> {
        self.engine
            .cancel_mcp_auth(params.server.clone())
            .await
            .map_err(protocol_fault)?;
        Ok(McpAuthCancelResult {
            server: params.server,
        })
    }

    async fn list_mcp_servers(&self, _: McpServerListParams) -> Result<McpServerListResult> {
        Ok(McpServerListResult {
            servers: self.engine.list_mcp_servers(),
        })
    }

    async fn add_mcp_server(&self, params: McpServerAddParams) -> Result<McpServerMutationResult> {
        self.engine
            .add_mcp_server(params.name, params.definition)
            .await
            .map(|server| McpServerMutationResult { server })
            .map_err(protocol_fault)
    }

    async fn edit_mcp_server(
        &self,
        params: McpServerEditParams,
    ) -> Result<McpServerMutationResult> {
        self.engine
            .edit_mcp_server(params.name, params.definition)
            .await
            .map(|server| McpServerMutationResult { server })
            .map_err(protocol_fault)
    }

    async fn remove_mcp_server(
        &self,
        params: McpServerNameParams,
    ) -> Result<McpServerMutationResult> {
        self.engine
            .remove_mcp_server(params.name)
            .await
            .map(|server| McpServerMutationResult { server })
            .map_err(protocol_fault)
    }

    async fn set_mcp_server_enabled(
        &self,
        params: McpServerSetEnabledParams,
    ) -> Result<McpServerMutationResult> {
        self.engine
            .set_mcp_server_enabled(params.name, params.enabled)
            .await
            .map(|server| McpServerMutationResult { server })
            .map_err(protocol_fault)
    }

    async fn reconnect_mcp_server(
        &self,
        params: McpServerNameParams,
    ) -> Result<McpServerMutationResult> {
        self.engine
            .reconnect_mcp_server(params.name)
            .await
            .map(|server| McpServerMutationResult { server })
            .map_err(protocol_fault)
    }

    async fn persist_mcp_server(
        &self,
        params: McpServerPersistParams,
    ) -> Result<McpServerMutationResult> {
        self.engine
            .persist_mcp_server(params.name, params.target)
            .await
            .map(|server| McpServerMutationResult { server })
            .map_err(protocol_fault)
    }

    async fn runtime_snapshot(&self, _: RuntimeSnapshotGetParams) -> Result<RuntimeSnapshotResult> {
        self.engine.runtime_snapshot().map_err(protocol_fault)
    }

    async fn connect_provider(
        &self,
        params: ProviderConnectParams,
    ) -> Result<ProviderConnectResult> {
        // Provider connect takes the cross-process store lock; keep the
        // synchronous disk work off the async worker.
        let server = self.clone();
        tokio::task::spawn_blocking(move || Server::connect_provider(&server, params))
            .await
            .map_err(|_| ServerFault::from(RpcFault::engine()))?
            .map_err(Into::into)
    }

    async fn disconnect_provider(
        &self,
        params: ProviderDisconnectParams,
    ) -> Result<ProviderDisconnectResult> {
        // Provider disconnect takes the cross-process store lock; keep the
        // synchronous disk work off the async worker.
        let server = self.clone();
        tokio::task::spawn_blocking(move || Server::disconnect_provider(&server, params))
            .await
            .map_err(|_| ServerFault::from(RpcFault::engine()))?
            .map_err(Into::into)
    }
}

fn protocol_fault(error: EngineError) -> ServerFault {
    engine_fault(error).into()
}

fn rename_fault(request: &SessionRenameParams, error: EngineError) -> RpcFault {
    let code = match error {
        EngineError::RenameConflict => SessionRenameErrorCode::IdempotencyConflict,
        EngineError::Session(SessionError::Missing(_)) | EngineError::MissingActor(_) => {
            SessionRenameErrorCode::SessionNotFound
        }
        _ => return engine_fault(error),
    };
    RpcFault::rename(request, code)
}

fn approval_fault(request: &ApprovalRespondParams, error: EngineError) -> RpcFault {
    match error {
        EngineError::ApprovalResponse(failure) => RpcFault::approval_failure(request, &failure),
        EngineError::ApprovalNotPending { .. } => RpcFault::approval(
            request,
            ApprovalRespondErrorCode::ApprovalNotPending,
            None,
            None,
        ),
        EngineError::Session(SessionError::Missing(_)) | EngineError::MissingActor(_) => {
            RpcFault::approval(
                request,
                ApprovalRespondErrorCode::ApprovalNotFound,
                None,
                None,
            )
        }
        _ => engine_fault(error),
    }
}