lean_ctx_server/tools/registered/

ctx_work_graph.rs

// SPDX-License-Identifier: Apache-2.0
use rmcp::ErrorData;
use rmcp::model::Tool;
use serde_json::{Map, Value, json};

use crate::server::tool_trait::{McpTool, ToolContext, ToolOutput, get_str};
use lean_ctx::tool_defs::tool_def;

pub struct CtxWorkGraphTool;

impl McpTool for CtxWorkGraphTool {
    fn name(&self) -> &'static str {
        "ctx_work_graph"
    }

    fn tool_def(&self) -> Tool {
        tool_def(
            "ctx_work_graph",
            "Free persistent project-scoped bounded local Work Graph for multi-agent coordination. The host launches agents; start binds a host-launched child to its node (verified context capsule, leases, launch env with the child's delivery profile), receipt or cancel ends it and a stop revokes its delivery. Actions: create|delegate|claim|start|consume|complete|receipt|accept|fuse|attribution|value_report|cancel|get|observe|list. Observe returns compact state without context payloads or execution fences. Node path claims share one lease namespace with ctx_agent path leases: an agent holding src blocks starting a node that claims src/x.",
            json!({
                "type": "object",
                "properties": {
                    "action": {"type":"string","enum":["create","delegate","claim","start","consume","complete","receipt","accept","fuse","attribution","value_report","cancel","get","observe","list"]},
                    "graph_id": {"type":"string","minLength":1,"maxLength":128},
                    "if_revision": {"type":"string","pattern":"^[0-9a-f]{64}$","description":"For observe: return only unchanged=true and revision if state matches. For delegate|claim|start|consume|complete|receipt|accept|cancel: the write_revision from observe; the write fails with a revision conflict and changes nothing if the graph changed since. Refused on other actions. Authorization is still enforced."},
                    "node_id": {"type":"string","minLength":1,"maxLength":128},
                    "parent_node_id": {"type":"string","minLength":1,"maxLength":128},
                    "to_agent": {"type":"string","minLength":1,"maxLength":128},
                    "capsule_ref": {"type":"string","minLength":1,"maxLength":4096},
                    "outcome_ref": {"type":"string","minLength":1,"maxLength":4096},
                    "tokens": {"type":"integer","minimum":0},
                    "cost_micros": {"type":"integer","minimum":0}
                    ,"max_concurrency": {"type":"integer","minimum":1,"maximum":16}
                    ,"receipt_id": {"type":"string","minLength":1,"maxLength":128}
                    ,"execution_fence": {"type":"string","pattern":"^fence:[0-9a-f]{64}$"}
                    ,"outcome": {"type":"string","enum":["accepted","rejected","partial"]}
                    ,"claims": {"type":"array","maxItems":64,"items":{"type":"object","properties":{"key":{"type":"string","minLength":1,"maxLength":128},"value":{"type":"string","maxLength":4096}},"required":["key","value"],"additionalProperties":false}}
                    ,"accepted_nodes": {"type":"array","minItems":1,"maxItems":256,"items":{"type":"string","minLength":1,"maxLength":128}}
                    ,"stop_reason": {"type":"string","enum":["manual_stop","stale","redundant","policy_denied","lease_lost","duplicate","low_value","execution_failed"]}
                    ,"timeout_ms": {"type":"integer","minimum":1,"maximum":1800000,"description":"Lease window for claim/start in ms (start: at most 30 min, default 5 min)."}
                    ,"path_claims": {"type":"array","maxItems":64,"items":{"type":"string","pattern":"^path:[^/\\\\](?:[^\\r\\n]*[^/\\\\])?$","maxLength":517}}
                    ,"task_ref": {"type":"string","minLength":1,"maxLength":128}
                    ,"policy_ref": {"type":"string","minLength":1,"maxLength":128}
                    ,"expected_outcome_ref": {"type":"string","minLength":1,"maxLength":128}
                },
                "required": ["action"],
                "additionalProperties": false
            }),
        )
    }

    fn handle(
        &self,
        args: &Map<String, Value>,
        ctx: &ToolContext,
    ) -> Result<ToolOutput, ErrorData> {
        let action = get_str(args, "action")
            .ok_or_else(|| ErrorData::invalid_params("action is required", None))?;
        // Free local reference behavior has no billing dependency. Registration,
        // project scope, ownership, bounded execution and policy still apply.
        let graph_id = get_str(args, "graph_id").unwrap_or_default();
        let agent_id = ctx
            .agent_id
            .as_ref()
            .map(|agent| agent.blocking_read().clone())
            .unwrap_or_default();
        let changed = matches!(
            action.as_str(),
            "create"
                | "delegate"
                | "claim"
                | "start"
                | "consume"
                | "complete"
                | "receipt"
                | "accept"
                | "cancel"
        );
        let result = crate::tools::ctx_work_graph::handle(crate::tools::ctx_work_graph::Request {
            action: &action,
            project_root: &ctx.project_root,
            agent_id: agent_id.as_deref().unwrap_or_default(),
            graph_id: &graph_id,
            if_revision: get_str(args, "if_revision").as_deref(),
            node_id: get_str(args, "node_id").as_deref(),
            parent_node_id: get_str(args, "parent_node_id").as_deref(),
            to_agent: get_str(args, "to_agent").as_deref(),
            capsule_ref: get_str(args, "capsule_ref").as_deref(),
            outcome_ref: get_str(args, "outcome_ref").as_deref(),
            tokens: args.get("tokens").and_then(Value::as_u64),
            cost_micros: args.get("cost_micros").and_then(Value::as_u64),
            max_concurrency: args
                .get("max_concurrency")
                .and_then(Value::as_u64)
                .and_then(|value| usize::try_from(value).ok()),
            receipt_id: get_str(args, "receipt_id").as_deref(),
            execution_fence: get_str(args, "execution_fence").as_deref(),
            outcome: get_str(args, "outcome").as_deref(),
            claims: args.get("claims"),
            accepted_nodes: args.get("accepted_nodes"),
            stop_reason: get_str(args, "stop_reason").as_deref(),
            timeout_ms: args.get("timeout_ms").and_then(Value::as_u64),
            path_claims: args.get("path_claims"),
            task_ref: get_str(args, "task_ref").as_deref(),
            policy_ref: get_str(args, "policy_ref").as_deref(),
            expected_outcome_ref: get_str(args, "expected_outcome_ref").as_deref(),
        })
        .map_err(|error| ErrorData::invalid_params(error, None))?;
        Ok(ToolOutput {
            text: result,
            original_tokens: 0,
            saved_tokens: 0,
            mode: Some(action),
            path: None,
            changed,
            shell_outcome: None,
            content_blocks: None,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn work_graph_is_free_and_keeps_registration_isolation_and_policy() {
        let isolated = lean_ctx::core::data_dir::isolated_data_dir();
        let args = Map::from_iter([("action".into(), Value::String("list".into()))]);
        let context = ToolContext {
            project_root: isolated.path().to_string_lossy().into_owned(),
            ..ToolContext::default()
        };

        let Err(error) = CtxWorkGraphTool.handle(&args, &context) else {
            panic!("the dispatcher must enforce agent registration");
        };
        assert_eq!(
            error.message,
            "agent must be registered first via ctx_agent"
        );

        let mut registry = lean_ctx::core::agents::AgentRegistry::new();
        let agent = registry
            .register(
                "codex",
                Some("signed graph test"),
                &context.project_root,
                None,
            )
            .unwrap();
        registry.save().unwrap();
        let context = ToolContext {
            agent_id: Some(std::sync::Arc::new(tokio::sync::RwLock::new(Some(agent)))),
            ..context
        };
        let create = json!({
            "action": "create", "graph_id": "paid-graph", "node_id": "root",
            "capsule_ref": "capsule:root", "task_ref": "task:root",
            "policy_ref": "policy:root", "expected_outcome_ref": "outcome:root",
            "tokens": 100, "cost_micros": 50
        });
        let created = CtxWorkGraphTool
            .handle(create.as_object().unwrap(), &context)
            .unwrap();
        assert!(created.changed);
        let listed = CtxWorkGraphTool.handle(&args, &context).unwrap();
        assert!(!listed.changed);
        let listed: Value = serde_json::from_str(&listed.text).unwrap();
        assert_eq!(listed["graph_ids"], json!(["paid-graph"]));
        let get = json!({"action": "get", "graph_id": "paid-graph"});
        let before = CtxWorkGraphTool
            .handle(get.as_object().unwrap(), &context)
            .unwrap();
        let observe = json!({"action": "observe", "graph_id": "paid-graph"});
        let observed = CtxWorkGraphTool
            .handle(observe.as_object().unwrap(), &context)
            .unwrap();
        assert!(!observed.changed);
        let observation: Value = serde_json::from_str(&observed.text).unwrap();
        assert!(observation.get("nodes").is_some());
        assert!(!observed.text.contains("capsule:root"));
        let conditional = json!({
            "action": "observe", "graph_id": "paid-graph",
            "if_revision": observation["revision"]
        });
        let unchanged = CtxWorkGraphTool
            .handle(conditional.as_object().unwrap(), &context)
            .unwrap();
        assert!(!unchanged.changed);
        let unchanged: Value = serde_json::from_str(&unchanged.text).unwrap();
        assert_eq!(unchanged["unchanged"], true);
        assert!(unchanged.get("nodes").is_none());
        let wrong_project = ToolContext {
            project_root: isolated
                .path()
                .join("other-project")
                .to_string_lossy()
                .into_owned(),
            agent_id: context.agent_id.clone(),
            ..ToolContext::default()
        };
        assert!(
            CtxWorkGraphTool
                .handle(conditional.as_object().unwrap(), &wrong_project)
                .is_err()
        );
        assert_eq!(
            before.text,
            CtxWorkGraphTool
                .handle(get.as_object().unwrap(), &context)
                .unwrap()
                .text
        );
        let policy_dir = std::path::Path::new(&context.project_root).join(".lean-ctx");
        std::fs::create_dir_all(&policy_dir).unwrap();
        std::fs::write(policy_dir.join("policy.toml"),
            "name = \"deny-execution\"\nversion = \"1.0.0\"\ndescription = \"test\"\n[context]\ndeny_tools = [\"ctx_work_graph\"]\n").unwrap();
        let start = json!({
            "action": "start", "graph_id": "paid-graph", "node_id": "root"
        });
        let error = CtxWorkGraphTool
            .handle(start.as_object().unwrap(), &context)
            .err()
            .expect("project policy must deny execution");
        assert_eq!(error.message, "project policy denies Work Graph execution");
        let after = CtxWorkGraphTool
            .handle(get.as_object().unwrap(), &context)
            .unwrap();
        assert_eq!(
            before.text, after.text,
            "denial must not claim or mutate the graph"
        );
        std::fs::remove_file(policy_dir.join("policy.toml")).unwrap();
        let mut free_create = create.clone();
        free_create["graph_id"] = json!("free-graph");
        assert!(
            CtxWorkGraphTool
                .handle(free_create.as_object().unwrap(), &context)
                .expect("registered agents can create local graphs without a paid account")
                .changed
        );
        assert!(
            CtxWorkGraphTool
                .handle(conditional.as_object().unwrap(), &context)
                .is_ok()
        );
        let cancel = json!({
            "action": "cancel", "graph_id": "paid-graph", "node_id": "root",
            "stop_reason": "manual_stop"
        });
        let unregistered = ToolContext {
            project_root: context.project_root.clone(),
            ..ToolContext::default()
        };
        let error = CtxWorkGraphTool
            .handle(cancel.as_object().unwrap(), &unregistered)
            .err()
            .expect("cancel still requires registration");
        assert_eq!(
            error.message,
            "agent must be registered first via ctx_agent"
        );
        let other_project = ToolContext {
            project_root: isolated
                .path()
                .join("other-project")
                .to_string_lossy()
                .into_owned(),
            agent_id: context.agent_id.clone(),
            ..ToolContext::default()
        };
        let error = CtxWorkGraphTool
            .handle(cancel.as_object().unwrap(), &other_project)
            .err()
            .expect("cancel must preserve project isolation");
        assert!(
            error
                .message
                .contains("agent is not active in this project")
        );
        assert!(
            CtxWorkGraphTool
                .handle(cancel.as_object().unwrap(), &context)
                .unwrap()
                .changed
        );
    }

    #[test]
    fn forged_legacy_plan_files_never_bypass_agent_registration() {
        let isolated = lean_ctx::core::data_dir::isolated_data_dir();
        let cloud = isolated.path().join("cloud");
        std::fs::create_dir_all(&cloud).unwrap();
        std::fs::write(cloud.join("plan.txt"), "enterprise").unwrap();
        std::fs::write(
            cloud.join("plan.json"),
            br#"{"plan":"enterprise","verified_at":9007199254740991}"#,
        )
        .unwrap();
        let args = Map::from_iter([("action".into(), Value::String("list".into()))]);
        let context = ToolContext {
            project_root: isolated.path().to_string_lossy().into_owned(),
            ..ToolContext::default()
        };

        let Err(error) = CtxWorkGraphTool.handle(&args, &context) else {
            panic!("forged plan files must not bypass agent registration");
        };

        assert_eq!(
            error.message,
            "agent must be registered first via ctx_agent"
        );
    }
}