yorishiro-server 0.12.1

use http::request::Parts;
use rmcp::ErrorData;
use rmcp::handler::server::common::Extension;
use rmcp::handler::server::wrapper::Parameters;
use rmcp::model::CallToolResult;
use rmcp::tool;
use rmcp::tool_router;
use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::Value;
use uuid::Uuid;
use yorishiro_core::repositories::entities;
use yorishiro_core::services::auth::ApiKeyScope;

use super::{YorishiroMcpServer, authorized, mcp_try, ok_json};

#[derive(Debug, Deserialize, JsonSchema)]
pub struct CreateEntityArgs {
    /// Name of the schema this entity conforms to. The tenant's current active version is used.
    pub schema_name: String,
    /// entity_type name declared in the schema.
    pub entity_type: String,
    /// Entity body (JSON object) conforming to the schema's `fields` definition.
    pub data: Value,
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct GetEntityArgs {
    pub id: Uuid,
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct UpdateEntityArgs {
    pub id: Uuid,
    /// Replacement entity body. Validated against the schema version in effect when the entity was created.
    pub data: Value,
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct DeleteEntityArgs {
    pub id: Uuid,
}

#[derive(Debug, Deserialize, JsonSchema)]
pub struct ListEntitiesArgs {
    pub entity_type: Option<String>,
    /// JSONB containment filter matched against entity data, e.g. `{"status": "active"}`.
    pub filter: Option<Value>,
    /// Maximum number of results (defaults to 50 if omitted).
    pub limit: Option<i64>,
    /// Number of records to skip (defaults to 0 if omitted).
    pub offset: Option<i64>,
}

#[tool_router(vis = "pub(crate)", router = tool_router_entities)]
impl YorishiroMcpServer {
    #[tool(description = "Create a new entity (requires write scope)")]
    pub async fn create_entity(
        &self,
        Parameters(args): Parameters<CreateEntityArgs>,
        Extension(parts): Extension<Parts>,
    ) -> Result<CallToolResult, ErrorData> {
        let mut authorized = authorized!(&self.state, &parts, ApiKeyScope::Write);

        let input = entities::CreateEntityInput {
            schema_name: args.schema_name,
            entity_type: args.entity_type,
            data: args.data,
        };

        let workspace_id = authorized.ctx.workspace_id;
        let created_by = authorized.ctx.user_id;
        let record =
            mcp_try!(entities::create(authorized.conn(), workspace_id, input, created_by).await);
        self.state
            .spawn_embedding_sync(authorized.ctx.tenant_id, workspace_id, record.clone());
        ok_json(record)
    }

    #[tool(description = "Get a single entity by ID (requires read scope)")]
    pub async fn get_entity(
        &self,
        Parameters(args): Parameters<GetEntityArgs>,
        Extension(parts): Extension<Parts>,
    ) -> Result<CallToolResult, ErrorData> {
        let mut authorized = authorized!(&self.state, &parts, ApiKeyScope::Read);

        let workspace_id = authorized.ctx.workspace_id;
        let record = mcp_try!(entities::get(authorized.conn(), workspace_id, args.id).await);
        ok_json(record)
    }

    #[tool(description = "Replace the data of an existing entity (requires write scope)")]
    pub async fn update_entity(
        &self,
        Parameters(args): Parameters<UpdateEntityArgs>,
        Extension(parts): Extension<Parts>,
    ) -> Result<CallToolResult, ErrorData> {
        let mut authorized = authorized!(&self.state, &parts, ApiKeyScope::Write);

        let workspace_id = authorized.ctx.workspace_id;
        let updated_by = authorized.ctx.user_id;
        let record = mcp_try!(
            entities::update(
                authorized.conn(),
                workspace_id,
                args.id,
                args.data,
                updated_by,
            )
            .await
        );
        self.state
            .spawn_embedding_sync(authorized.ctx.tenant_id, workspace_id, record.clone());
        ok_json(record)
    }

    #[tool(description = "Delete an entity (requires write scope)")]
    pub async fn delete_entity(
        &self,
        Parameters(args): Parameters<DeleteEntityArgs>,
        Extension(parts): Extension<Parts>,
    ) -> Result<CallToolResult, ErrorData> {
        let mut authorized = authorized!(&self.state, &parts, ApiKeyScope::Write);

        let workspace_id = authorized.ctx.workspace_id;
        mcp_try!(entities::delete(authorized.conn(), workspace_id, args.id).await);
        ok_json(serde_json::json!({ "deleted": true }))
    }

    #[tool(description = "List entities (requires read scope)")]
    pub async fn list_entities(
        &self,
        Parameters(args): Parameters<ListEntitiesArgs>,
        Extension(parts): Extension<Parts>,
    ) -> Result<CallToolResult, ErrorData> {
        let mut authorized = authorized!(&self.state, &parts, ApiKeyScope::Read);

        let default = entities::ListEntitiesQuery::default();
        let query = entities::ListEntitiesQuery {
            entity_type: args.entity_type,
            filter: args.filter,
            limit: args.limit.unwrap_or(default.limit),
            offset: args.offset.unwrap_or(default.offset),
        };

        let workspace_id = authorized.ctx.workspace_id;
        let records = mcp_try!(entities::list(authorized.conn(), workspace_id, query).await);
        ok_json(records)
    }
}