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

apollo/
mcp_server.rs
