Skip to main content

embacle_mcp/tools/
model.rs

1// ABOUTME: MCP tools for getting and setting the active model for the current provider
2// ABOUTME: Returns available models from the runner and accepts model override strings
3//
4// SPDX-License-Identifier: Apache-2.0
5// Copyright (c) 2026 dravr.ai
6
7use async_trait::async_trait;
8use serde_json::{json, Value};
9
10use dravr_tronc::mcp::protocol::{CallToolResult, ToolDefinition};
11use dravr_tronc::McpTool;
12
13use crate::state::SharedState;
14
15/// Returns the current model, default model, and available models for the active provider
16pub struct GetModel;
17
18#[async_trait]
19impl McpTool<crate::state::ServerState> for GetModel {
20    fn definition(&self) -> ToolDefinition {
21        ToolDefinition {
22            name: "get_model".to_owned(),
23            description: "Get the current model and list available models for the active provider"
24                .to_owned(),
25            input_schema: json!({
26                "type": "object",
27                "properties": {}
28            }),
29        }
30    }
31
32    async fn execute(&self, state: &SharedState, _arguments: Value) -> CallToolResult {
33        let state_guard = state.read().await;
34        let provider = state_guard.active_provider();
35        let current_model = state_guard.active_model().map(ToOwned::to_owned);
36        let runner_result = state_guard.get_runner(provider).await;
37        drop(state_guard);
38
39        let (default_model, available_models) = match runner_result {
40            Ok(runner) => (
41                runner.default_model().to_owned(),
42                runner.available_models().to_vec(),
43            ),
44            Err(e) => {
45                return CallToolResult::text(
46                    json!({
47                        "provider": provider.to_string(),
48                        "current_model": current_model,
49                        "error": format!("Could not load runner: {e}")
50                    })
51                    .to_string(),
52                );
53            }
54        };
55
56        CallToolResult::text(
57            json!({
58                "provider": provider.to_string(),
59                "current_model": current_model,
60                "default_model": default_model,
61                "available_models": available_models
62            })
63            .to_string(),
64        )
65    }
66}
67
68/// Sets the model for subsequent prompt dispatch requests
69pub struct SetModel;
70
71#[async_trait]
72impl McpTool<crate::state::ServerState> for SetModel {
73    fn definition(&self) -> ToolDefinition {
74        ToolDefinition {
75            name: "set_model".to_owned(),
76            description: "Set the model for the active provider (pass null to reset to default)"
77                .to_owned(),
78            input_schema: json!({
79                "type": "object",
80                "properties": {
81                    "model": {
82                        "type": "string",
83                        "description": "Model identifier (e.g. claude-opus-4-20250514, gpt-4o). Pass null to reset."
84                    }
85                },
86                "required": ["model"]
87            }),
88        }
89    }
90
91    async fn execute(&self, state: &SharedState, arguments: Value) -> CallToolResult {
92        let model = arguments
93            .get("model")
94            .and_then(Value::as_str)
95            .map(ToOwned::to_owned);
96
97        let mut state_guard = state.write().await;
98        state_guard.set_active_model(model.clone());
99        let provider = state_guard.active_provider();
100        drop(state_guard);
101
102        CallToolResult::text(
103            json!({
104                "provider": provider.to_string(),
105                "current_model": model,
106                "status": "updated"
107            })
108            .to_string(),
109        )
110    }
111}