apollo/
mcp_server.rs

1//! MCP server mode — exposes apollo as an MCP server over stdio or HTTP.
2//! Other AI clients can connect to prompt apollo or use its tools.
3
4use axum::{
5    extract::State,
6    http::StatusCode,
7    response::IntoResponse,
8    routing::{get, post},
9    Json, Router,
10};
11use serde::{Deserialize, Serialize};
12use serde_json::Value;
13use std::sync::Arc;
14use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
15use tokio::sync::mpsc;
16
17use crate::agent::loop_runner::AgentRunner;
18use crate::channels::traits::{Channel, IncomingMessage, OutgoingMessage};
19use crate::tools::traits::{Tool, ToolSpec};
20
21#[derive(Debug, Deserialize)]
22struct JsonRpcRequest {
23    #[allow(dead_code)]
24    jsonrpc: String,
25    id: Option<Value>,
26    method: String,
27    #[serde(default)]
28    params: Option<Value>,
29}
30
31#[derive(Debug, Serialize)]
32struct JsonRpcResponse {
33    jsonrpc: String,
34    id: Option<Value>,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    result: Option<Value>,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    error: Option<JsonRpcErrorObj>,
39}
40
41#[derive(Debug, Serialize)]
42struct JsonRpcErrorObj {
43    code: i32,
44    message: String,
45    #[serde(skip_serializing_if = "Option::is_none")]
46    data: Option<Value>,
47}
48
49impl JsonRpcResponse {
50    fn success(id: Option<Value>, result: Value) -> Self {
51        Self {
52            jsonrpc: "2.0".to_string(),
53            id,
54            result: Some(result),
55            error: None,
56        }
57    }
58
59    fn error(id: Option<Value>, code: i32, message: impl Into<String>) -> Self {
60        Self {
61            jsonrpc: "2.0".to_string(),
62            id,
63            result: None,
64            error: Some(JsonRpcErrorObj {
65                code,
66                message: message.into(),
67                data: None,
68            }),
69        }
70    }
71}
72
73fn tool_to_mcp_schema(spec: &ToolSpec) -> Value {
74    serde_json::json!({
75        "name": spec.name,
76        "description": spec.description,
77        "inputSchema": spec.parameters
78    })
79}
80
81/// A no-op Channel used for HTTP chat requests when no Telegram token is provided.
82struct HttpChannel;
83
84#[async_trait::async_trait]
85impl Channel for HttpChannel {
86    fn name(&self) -> &str {
87        "http"
88    }
89    async fn start(&mut self) -> anyhow::Result<mpsc::Receiver<IncomingMessage>> {
90        let (_tx, rx) = mpsc::channel(1);
91        Ok(rx)
92    }
93    async fn send(&self, _msg: OutgoingMessage) -> anyhow::Result<Option<String>> {
94        Ok(None)
95    }
96    async fn stop(&mut self) -> anyhow::Result<()> {
97        Ok(())
98    }
99}
100
101/// Live-update channel that posts directly to Telegram during agent execution.
102/// Sends ⏳ on first tool call, edits with tool progress, finalizes with the response.
103struct TelegramHttpChannel {
104    token: String,
105    client: reqwest::Client,
106    /// Rate-limit progress edits: track last edit time per chat
107    last_edit: std::sync::Mutex<Option<std::time::Instant>>,
108}
109
110impl TelegramHttpChannel {
111    fn new(token: String) -> Self {
112        Self {
113            token,
114            client: crate::http::shared(),
115            last_edit: std::sync::Mutex::new(None),
116        }
117    }
118
119    fn api_url(&self, method: &str) -> String {
120        format!("https://api.telegram.org/bot{}/{}", self.token, method)
121    }
122
123    async fn tg_post(
124        &self,
125        method: &str,
126        body: serde_json::Value,
127    ) -> anyhow::Result<serde_json::Value> {
128        let resp = self
129            .client
130            .post(self.api_url(method))
131            .json(&body)
132            .send()
133            .await?;
134        Ok(resp.json().await?)
135    }
136}
137
138#[async_trait::async_trait]
139impl Channel for TelegramHttpChannel {
140    fn name(&self) -> &str {
141        "telegram-http"
142    }
143
144    async fn start(&mut self) -> anyhow::Result<mpsc::Receiver<IncomingMessage>> {
145        let (_tx, rx) = mpsc::channel(1);
146        Ok(rx)
147    }
148
149    async fn send(&self, msg: OutgoingMessage) -> anyhow::Result<Option<String>> {
150        let body = serde_json::json!({
151            "chat_id": msg.chat_id,
152            "text": msg.text,
153        });
154        let resp = self.tg_post("sendMessage", body).await?;
155        Ok(resp["result"]["message_id"]
156            .as_i64()
157            .map(|id| id.to_string()))
158    }
159
160    fn supports_draft_updates(&self) -> bool {
161        true
162    }
163
164    async fn send_draft(&self, chat_id: &str, text: &str) -> anyhow::Result<Option<String>> {
165        let resp = self
166            .tg_post(
167                "sendMessage",
168                serde_json::json!({
169                    "chat_id": chat_id,
170                    "text": text,
171                }),
172            )
173            .await?;
174        Ok(resp["result"]["message_id"]
175            .as_i64()
176            .map(|id| id.to_string()))
177    }
178
179    async fn update_draft_progress(
180        &self,
181        chat_id: &str,
182        message_id: &str,
183        text: &str,
184    ) -> anyhow::Result<()> {
185        // Rate-limit to one edit per 1.5s to avoid Telegram 429s
186        {
187            let mut last = self.last_edit.lock().unwrap();
188            if let Some(t) = *last {
189                if t.elapsed().as_millis() < 1500 {
190                    return Ok(());
191                }
192            }
193            *last = Some(std::time::Instant::now());
194        }
195
196        let msg_id: i64 = message_id.parse().unwrap_or(0);
197        let _ = self
198            .tg_post(
199                "editMessageText",
200                serde_json::json!({
201                    "chat_id": chat_id,
202                    "message_id": msg_id,
203                    "text": text,
204                }),
205            )
206            .await;
207        Ok(())
208    }
209
210    async fn finalize_draft(
211        &self,
212        chat_id: &str,
213        message_id: &str,
214        text: &str,
215    ) -> anyhow::Result<()> {
216        let msg_id: i64 = message_id.parse().unwrap_or(0);
217        let html = md_to_telegram_html(text);
218
219        // Try HTML first, fall back to plain
220        let r = self
221            .tg_post(
222                "editMessageText",
223                serde_json::json!({
224                    "chat_id": chat_id,
225                    "message_id": msg_id,
226                    "text": &html,
227                    "parse_mode": "HTML",
228                }),
229            )
230            .await?;
231
232        if !r["ok"].as_bool().unwrap_or(false) {
233            let _ = self
234                .tg_post(
235                    "editMessageText",
236                    serde_json::json!({
237                        "chat_id": chat_id,
238                        "message_id": msg_id,
239                        "text": strip_html(&html),
240                    }),
241                )
242                .await;
243        }
244        Ok(())
245    }
246
247    async fn stop(&mut self) -> anyhow::Result<()> {
248        Ok(())
249    }
250}
251
252/// Convert GitHub Flavoured Markdown to Telegram HTML.
253fn md_to_telegram_html(md: &str) -> String {
254    let mut out = Vec::new();
255    let mut in_code = false;
256    let mut code_buf: Vec<&str> = Vec::new();
257
258    for line in md.lines() {
259        if line.starts_with("```") {
260            if !in_code {
261                in_code = true;
262                code_buf.clear();
263            } else {
264                let code = esc_html(&code_buf.join("\n"));
265                out.push(format!("<pre>{}</pre>", code));
266                code_buf.clear();
267                in_code = false;
268            }
269            continue;
270        }
271        if in_code {
272            code_buf.push(line);
273            continue;
274        }
275        // Table separator → skip
276        if line.starts_with('|') && line.contains("---") {
277            continue;
278        }
279        // Table row → flatten
280        if line.starts_with('|') {
281            let cells: Vec<&str> = line
282                .split('|')
283                .map(str::trim)
284                .filter(|s| !s.is_empty())
285                .collect();
286            if !cells.is_empty() {
287                out.push(inline_to_html(&cells.join("  •  ")));
288            }
289            continue;
290        }
291        // Horizontal rule → blank
292        if line
293            .chars()
294            .all(|c| c == '-' || c == '*' || c == '_' || c == ' ')
295            && line.len() >= 3
296        {
297            out.push(String::new());
298            continue;
299        }
300        // Heading
301        if let Some(rest) = line
302            .strip_prefix("### ")
303            .or_else(|| line.strip_prefix("## "))
304            .or_else(|| line.strip_prefix("# "))
305        {
306            out.push(format!("<b>{}</b>", inline_to_html(rest)));
307            continue;
308        }
309        // Bullet
310        if let Some(rest) = line
311            .strip_prefix("- ")
312            .or_else(|| line.strip_prefix("* "))
313            .or_else(|| line.strip_prefix("• "))
314        {
315            out.push(format!("• {}", inline_to_html(rest)));
316            continue;
317        }
318        out.push(inline_to_html(line));
319    }
320    if in_code && !code_buf.is_empty() {
321        out.push(format!("<pre>{}</pre>", esc_html(&code_buf.join("\n"))));
322    }
323
324    // Collapse 3+ blank lines to 2
325    let joined = out.join("\n");
326    let mut collapsed = String::with_capacity(joined.len());
327    let mut blank_count = 0;
328    for line in joined.lines() {
329        if line.trim().is_empty() {
330            blank_count += 1;
331            if blank_count <= 2 {
332                collapsed.push('\n');
333            }
334        } else {
335            blank_count = 0;
336            collapsed.push_str(line);
337            collapsed.push('\n');
338        }
339    }
340    collapsed.trim().to_string()
341}
342
343fn esc_html(s: &str) -> String {
344    s.replace('&', "&amp;")
345        .replace('<', "&lt;")
346        .replace('>', "&gt;")
347}
348
349fn inline_to_html(text: &str) -> String {
350    let mut out = String::new();
351    let bytes = text.as_bytes();
352    let mut i = 0;
353    while i < bytes.len() {
354        // Inline code
355        if bytes[i] == b'`' {
356            if let Some(j) = text[i + 1..].find('`') {
357                out.push_str(&format!(
358                    "<code>{}</code>",
359                    esc_html(&text[i + 1..i + 1 + j])
360                ));
361                i = i + 1 + j + 1;
362                continue;
363            }
364        }
365        // Bold **...**
366        if bytes.get(i) == Some(&b'*') && bytes.get(i + 1) == Some(&b'*') {
367            if let Some(j) = text[i + 2..].find("**") {
368                out.push_str(&format!("<b>{}</b>", esc_html(&text[i + 2..i + 2 + j])));
369                i = i + 2 + j + 2;
370                continue;
371            }
372        }
373        // Link [text](url)
374        if bytes[i] == b'[' {
375            if let Some(j) = text[i + 1..].find("](") {
376                let link_text = &text[i + 1..i + 1 + j];
377                let rest = &text[i + 1 + j + 2..];
378                if let Some(k) = rest.find(')') {
379                    let url = &rest[..k];
380                    out.push_str(&format!("<a href=\"{}\">{}</a>", url, esc_html(link_text)));
381                    i = i + 1 + j + 2 + k + 1;
382                    continue;
383                }
384            }
385        }
386        // Default: escape and emit
387        let c = text[i..].chars().next().unwrap_or(' ');
388        out.push_str(&esc_html(&c.to_string()));
389        i += c.len_utf8();
390    }
391    out
392}
393
394fn strip_html(html: &str) -> String {
395    let mut out = String::with_capacity(html.len());
396    let mut in_tag = false;
397    for c in html.chars() {
398        match c {
399            '<' => in_tag = true,
400            '>' => in_tag = false,
401            _ if !in_tag => out.push(c),
402            _ => {}
403        }
404    }
405    out
406}
407
408#[derive(Clone)]
409struct McpState {
410    tools: Arc<Vec<Arc<dyn Tool>>>,
411    provider: Option<Arc<dyn crate::providers::traits::Provider>>,
412    model: Option<String>,
413    runner: Option<Arc<AgentRunner>>,
414}
415
416#[derive(Debug, Deserialize)]
417struct HttpChatRequest {
418    text: String,
419    chat_id: String,
420    /// Optional Telegram bot token — enables live progress updates during execution
421    telegram_token: Option<String>,
422}
423
424#[derive(Debug, Serialize)]
425struct HttpChatResponse {
426    text: String,
427}
428
429/// Run apollo as an MCP server over stdio.
430pub async fn run_mcp_server(
431    tools: Vec<Arc<dyn Tool>>,
432    provider: Option<Arc<dyn crate::providers::traits::Provider>>,
433    model: Option<String>,
434) -> anyhow::Result<()> {
435    let stdin = BufReader::new(tokio::io::stdin());
436    let mut stdout = tokio::io::stdout();
437    let mut lines = stdin.lines();
438
439    let tools = Arc::new(tools);
440
441    tracing::info!("MCP server started on stdio");
442
443    while let Ok(Some(line)) = lines.next_line().await {
444        let line = line.trim().to_string();
445        if line.is_empty() {
446            continue;
447        }
448
449        let request: JsonRpcRequest = match serde_json::from_str(&line) {
450            Ok(r) => r,
451            Err(_) => {
452                let resp = JsonRpcResponse::error(None, -32700, "Parse error");
453                write_response(&mut stdout, &resp).await?;
454                continue;
455            }
456        };
457
458        let response = handle_request(&request, &tools, &provider, &model).await;
459        write_response(&mut stdout, &response).await?;
460    }
461
462    Ok(())
463}
464
465/// Run apollo as an MCP server over HTTP (for Cloudflare Container deployment).
466async fn health_handler() -> Json<serde_json::Value> {
467    Json(serde_json::json!({
468        "status": "ok",
469        "version": env!("CARGO_PKG_VERSION"),
470        "service": "apollo",
471        "timestamp": chrono::Utc::now().to_rfc3339(),
472    }))
473}
474
475pub async fn run_mcp_server_http(
476    tools: Vec<Arc<dyn Tool>>,
477    provider: Option<Arc<dyn crate::providers::traits::Provider>>,
478    model: Option<String>,
479    runner: Option<Arc<AgentRunner>>,
480    port: u16,
481) -> anyhow::Result<()> {
482    let state = McpState {
483        tools: Arc::new(tools),
484        provider,
485        model,
486        runner,
487    };
488
489    let app = Router::new()
490        .route("/health", get(health_handler))
491        .route("/chat", post(handle_http_chat))
492        .route("/mcp", post(handle_http_mcp))
493        .route("/", post(handle_http_mcp))
494        .with_state(state);
495
496    let addr = std::net::SocketAddr::from(([0, 0, 0, 0], port));
497    tracing::info!("MCP HTTP server listening on {}", addr);
498    eprintln!("MCP HTTP server listening on {}", addr);
499
500    let listener = tokio::net::TcpListener::bind(addr).await?;
501    axum::serve(listener, app).await?;
502
503    Ok(())
504}
505
506async fn handle_http_chat(
507    State(state): State<McpState>,
508    Json(req): Json<HttpChatRequest>,
509) -> (StatusCode, Json<HttpChatResponse>) {
510    let Some(runner) = state.runner else {
511        return (
512            StatusCode::SERVICE_UNAVAILABLE,
513            Json(HttpChatResponse {
514                text: "Agent not initialized".to_string(),
515            }),
516        );
517    };
518
519    let msg = IncomingMessage {
520        id: uuid::Uuid::new_v4().to_string(),
521        sender_id: req.chat_id.clone(),
522        sender_name: None,
523        chat_id: req.chat_id.clone(),
524        text: req.text,
525        is_group: false,
526        reply_to: None,
527        timestamp: chrono::Utc::now(),
528    };
529
530    // Use TelegramHttpChannel for live progress if token provided, else plain HTTP channel
531    let result = if let Some(token) = req.telegram_token {
532        let channel = TelegramHttpChannel::new(token);
533        runner.handle_message(&msg, &channel).await
534    } else {
535        runner.handle_message(&msg, &HttpChannel).await
536    };
537
538    match result {
539        Ok(text) => (StatusCode::OK, Json(HttpChatResponse { text })),
540        Err(e) => (
541            StatusCode::INTERNAL_SERVER_ERROR,
542            Json(HttpChatResponse {
543                text: format!("Error: {e}"),
544            }),
545        ),
546    }
547}
548
549async fn handle_http_mcp(
550    State(state): State<McpState>,
551    Json(request): Json<JsonRpcRequest>,
552) -> impl IntoResponse {
553    let response = handle_request(&request, &state.tools, &state.provider, &state.model).await;
554    Json(response)
555}
556
557async fn write_response(
558    stdout: &mut tokio::io::Stdout,
559    response: &JsonRpcResponse,
560) -> anyhow::Result<()> {
561    let json = serde_json::to_string(response)?;
562    stdout.write_all(json.as_bytes()).await?;
563    stdout.write_all(b"\n").await?;
564    stdout.flush().await?;
565    Ok(())
566}
567
568async fn handle_request(
569    req: &JsonRpcRequest,
570    tools: &[Arc<dyn Tool>],
571    provider: &Option<Arc<dyn crate::providers::traits::Provider>>,
572    model: &Option<String>,
573) -> JsonRpcResponse {
574    match req.method.as_str() {
575        "initialize" => handle_initialize(req.id.clone()),
576        "notifications/initialized" => JsonRpcResponse::success(req.id.clone(), Value::Null),
577        "tools/list" => handle_tools_list(req.id.clone(), tools, provider.is_some()),
578        "tools/call" => {
579            handle_tools_call(req.id.clone(), req.params.as_ref(), tools, provider, model).await
580        }
581        "shutdown" => {
582            tracing::info!("MCP server shutting down");
583            JsonRpcResponse::success(req.id.clone(), Value::Null)
584        }
585        _ => JsonRpcResponse::error(req.id.clone(), -32601, "Method not found"),
586    }
587}
588
589fn handle_initialize(id: Option<Value>) -> JsonRpcResponse {
590    JsonRpcResponse::success(
591        id,
592        serde_json::json!({
593            "protocolVersion": "2024-11-05",
594            "capabilities": {
595                "tools": {}
596            },
597            "serverInfo": {
598                "name": "apollo",
599                "version": env!("CARGO_PKG_VERSION")
600            }
601        }),
602    )
603}
604
605fn handle_tools_list(
606    id: Option<Value>,
607    tools: &[Arc<dyn Tool>],
608    has_provider: bool,
609) -> JsonRpcResponse {
610    let mut mcp_tools: Vec<Value> = tools
611        .iter()
612        .map(|t| tool_to_mcp_schema(&t.spec()))
613        .collect();
614
615    // Add "ask" tool if we have a provider (allows other AIs to prompt apollo)
616    if has_provider {
617        mcp_tools.push(serde_json::json!({
618            "name": "ask",
619            "description": "Send a message to the apollo AI agent and get a response. Use this to prompt apollo for complex tasks like coding, research, or analysis.",
620            "inputSchema": {
621                "type": "object",
622                "properties": {
623                    "message": {
624                        "type": "string",
625                        "description": "The message/prompt to send to apollo"
626                    }
627                },
628                "required": ["message"]
629            }
630        }));
631    }
632
633    JsonRpcResponse::success(id, serde_json::json!({ "tools": mcp_tools }))
634}
635
636async fn handle_tools_call(
637    id: Option<Value>,
638    params: Option<&Value>,
639    tools: &[Arc<dyn Tool>],
640    provider: &Option<Arc<dyn crate::providers::traits::Provider>>,
641    model: &Option<String>,
642) -> JsonRpcResponse {
643    let params = match params {
644        Some(p) => p,
645        None => {
646            return JsonRpcResponse::error(id, -32602, "Missing params");
647        }
648    };
649
650    let name = match params.get("name").and_then(|n| n.as_str()) {
651        Some(n) => n,
652        None => {
653            return JsonRpcResponse::error(id, -32602, "Missing tool name");
654        }
655    };
656
657    let arguments = params
658        .get("arguments")
659        .cloned()
660        .unwrap_or(serde_json::json!({}));
661
662    // Handle the "ask" meta-tool
663    if name == "ask" {
664        if let Some(provider) = provider {
665            let message = arguments
666                .get("message")
667                .and_then(|m| m.as_str())
668                .unwrap_or("");
669            let model_name = model.as_deref().unwrap_or("claude-sonnet-4-5");
670
671            match provider.simple_chat(message, model_name).await {
672                Ok(response) => {
673                    return JsonRpcResponse::success(
674                        id,
675                        serde_json::json!({
676                            "content": [{"type": "text", "text": response}]
677                        }),
678                    );
679                }
680                Err(e) => {
681                    return JsonRpcResponse::success(
682                        id,
683                        serde_json::json!({
684                            "content": [{"type": "text", "text": format!("Error: {}", e)}],
685                            "isError": true
686                        }),
687                    );
688                }
689            }
690        } else {
691            return JsonRpcResponse::error(id, -32602, "No provider configured for ask tool");
692        }
693    }
694
695    // Find and execute the matching tool
696    let tool = match tools.iter().find(|t| t.name() == name) {
697        Some(t) => t,
698        None => {
699            return JsonRpcResponse::error(id, -32602, format!("Unknown tool: {}", name));
700        }
701    };
702
703    let args_str = serde_json::to_string(&arguments).unwrap_or_default();
704    match tool.execute(&args_str).await {
705        Ok(result) => JsonRpcResponse::success(
706            id,
707            serde_json::json!({
708                "content": [{"type": "text", "text": result.output}],
709                "isError": result.is_error
710            }),
711        ),
712        Err(e) => JsonRpcResponse::success(
713            id,
714            serde_json::json!({
715                "content": [{"type": "text", "text": format!("Tool error: {}", e)}],
716                "isError": true
717            }),
718        ),
719    }
720}
apollo/mcp_server.rs

apollo/
mcp_server.rs
