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::schema::{Tool, ToolResponse};
11use dravr_tronc::{McpTool, ToolContext};
12
13use crate::state::{ServerState, 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<ServerState> for GetModel {
20    fn definition(&self) -> Tool {
21        Tool {
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            annotations: None,
30        }
31    }
32
33    async fn execute(
34        &self,
35        state: &SharedState,
36        _ctx: &ToolContext,
37        _arguments: Value,
38    ) -> ToolResponse {
39        let provider = state.active_provider().await;
40        let current_model = state.active_model().await;
41        let runner_result = state.get_runner(provider).await;
42
43        let (default_model, available_models) = match runner_result {
44            Ok(runner) => (
45                runner.default_model().to_owned(),
46                runner.available_models().to_vec(),
47            ),
48            Err(e) => {
49                return ToolResponse::text(
50                    json!({
51                        "provider": provider.to_string(),
52                        "current_model": current_model,
53                        "error": format!("Could not load runner: {e}")
54                    })
55                    .to_string(),
56                );
57            }
58        };
59
60        ToolResponse::text(
61            json!({
62                "provider": provider.to_string(),
63                "current_model": current_model,
64                "default_model": default_model,
65                "available_models": available_models
66            })
67            .to_string(),
68        )
69    }
70}
71
72/// Sets the model for subsequent prompt dispatch requests
73pub struct SetModel;
74
75#[async_trait]
76impl McpTool<ServerState> for SetModel {
77    fn definition(&self) -> Tool {
78        Tool {
79            name: "set_model".to_owned(),
80            description: "Set the model for the active provider (pass null to reset to default)"
81                .to_owned(),
82            input_schema: json!({
83                "type": "object",
84                "properties": {
85                    "model": {
86                        "type": "string",
87                        "description": "Model identifier (e.g. claude-opus-4-20250514, gpt-4o). Pass null to reset."
88                    }
89                },
90                "required": ["model"]
91            }),
92            annotations: None,
93        }
94    }
95
96    async fn execute(
97        &self,
98        state: &SharedState,
99        _ctx: &ToolContext,
100        arguments: Value,
101    ) -> ToolResponse {
102        let model = arguments
103            .get("model")
104            .and_then(Value::as_str)
105            .map(ToOwned::to_owned);
106
107        state.set_active_model(model.clone()).await;
108        let provider = state.active_provider().await;
109
110        ToolResponse::text(
111            json!({
112                "provider": provider.to_string(),
113                "current_model": model,
114                "status": "updated"
115            })
116            .to_string(),
117        )
118    }
119}