1pub mod actions;
6
7pub use actions::McpClientProtocol;
8
9use anyhow::{Context, Result};
10use std::net::SocketAddr;
11use std::sync::Arc;
12use tokio::sync::mpsc;
13use tracing::{debug, error, info};
14
15use crate::client::llm_budget::call_llm_for_client;
16use crate::client::mcp::actions::{MCP_CLIENT_CONNECTED_EVENT, MCP_CLIENT_RESPONSE_RECEIVED_EVENT};
17use crate::llm::actions::client_trait::{Client, ClientActionResult};
18use crate::llm::actions::protocol_trait::Protocol;
19use crate::llm::ollama_client::OllamaClient;
20use crate::logging::emit::Log;
21use crate::protocol::Event;
22use crate::state::app_state::AppState;
23use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
24use crate::state::{AccessLogOwner, ClientId, ClientStatus};
25use serde_json::{json, Value};
26
27#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
29struct JsonRpcRequest {
30 jsonrpc: String,
31 method: String,
32 #[serde(skip_serializing_if = "Option::is_none")]
33 params: Option<Value>,
34 id: i64,
35}
36
37#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
39struct JsonRpcResponse {
40 jsonrpc: String,
41 #[serde(skip_serializing_if = "Option::is_none")]
42 result: Option<Value>,
43 #[serde(skip_serializing_if = "Option::is_none")]
44 error: Option<JsonRpcError>,
45 id: Option<i64>,
46}
47
48#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
50struct JsonRpcNotification {
51 jsonrpc: String,
52 method: String,
53 #[serde(skip_serializing_if = "Option::is_none")]
54 params: Option<Value>,
55}
56
57#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
59struct JsonRpcError {
60 code: i32,
61 message: String,
62 #[serde(skip_serializing_if = "Option::is_none")]
63 data: Option<Value>,
64}
65
66const DEFAULT_CLIENT_NAME: &str = "netget-mcp-client";
68
69const DEFAULT_CLIENT_VERSION: &str = "0.1.0";
71
72pub struct McpClient;
74
75impl McpClient {
76 pub async fn connect_with_llm_actions(
78 remote_addr: String,
79 llm_client: OllamaClient,
80 app_state: Arc<AppState>,
81 status_tx: mpsc::UnboundedSender<String>,
82 client_id: ClientId,
83 startup_params: Option<crate::protocol::StartupParams>,
84 ) -> Result<SocketAddr> {
85 info!("MCP client {} connecting to {}", client_id, remote_addr);
86
87 let (client_name, client_version) = match &startup_params {
92 Some(params) => (
93 params.get_optional_string("client_name")?,
94 params.get_optional_string("client_version")?,
95 ),
96 None => (None, None),
97 };
98 let client_name = client_name.unwrap_or_else(|| DEFAULT_CLIENT_NAME.to_string());
99 let client_version = client_version.unwrap_or_else(|| DEFAULT_CLIENT_VERSION.to_string());
100
101 let http_client = reqwest::Client::builder()
103 .timeout(std::time::Duration::from_secs(30))
104 .build()
105 .context("Failed to build HTTP client")?;
106
107 let base_url = if remote_addr.starts_with("http://") || remote_addr.starts_with("https://")
109 {
110 remote_addr.clone()
111 } else {
112 format!("http://{}", remote_addr)
113 };
114
115 app_state
116 .with_client_mut(client_id, |client| {
117 client.set_protocol_field("base_url".to_string(), serde_json::json!(base_url));
118 client.set_protocol_field("request_id".to_string(), serde_json::json!(1));
119 client.set_protocol_field("initialized".to_string(), serde_json::json!(false));
120 })
121 .await;
122
123 let init_response = Self::send_initialize_request(
125 &http_client,
126 &base_url,
127 client_id,
128 &app_state,
129 &status_tx,
130 &client_name,
131 &client_version,
132 )
133 .await?;
134
135 let server_info = init_response
137 .get("serverInfo")
138 .and_then(|v| v.as_object())
139 .context("Missing serverInfo in initialize response")?;
140
141 let server_name = server_info
142 .get("name")
143 .and_then(|v| v.as_str())
144 .unwrap_or("unknown");
145
146 let server_version = server_info
147 .get("version")
148 .and_then(|v| v.as_str())
149 .unwrap_or("unknown");
150
151 let capabilities = init_response
152 .get("capabilities")
153 .cloned()
154 .unwrap_or(json!({}));
155
156 app_state
158 .with_client_mut(client_id, |client| {
159 client.set_protocol_field("server_info".to_string(), json!(server_info));
160 client.set_protocol_field("capabilities".to_string(), capabilities.clone());
161 })
162 .await;
163
164 Self::send_initialized_notification(&http_client, &base_url, &status_tx).await?;
166
167 app_state
169 .with_client_mut(client_id, |client| {
170 client.set_protocol_field("initialized".to_string(), serde_json::json!(true));
171 })
172 .await;
173
174 app_state
176 .update_client_status(client_id, ClientStatus::Connected)
177 .await;
178 Log::new(Some(&status_tx)).info(format!(
179 "MCP client {} initialized with server {}",
180 client_id, server_name
181 ));
182 let _ = status_tx.send("__UPDATE_UI__".to_string());
183
184 info!("MCP client {} initialization complete", client_id);
185
186 let event = Event::new(
188 &MCP_CLIENT_CONNECTED_EVENT,
189 json!({
190 "server_name": server_name,
191 "server_version": server_version,
192 "capabilities": capabilities,
193 }),
194 );
195
196 let command_rx =
203 crate::client::command_support::register_command_channel(&app_state, client_id).await;
204 let cmd_task = tokio::spawn(Self::command_loop(
205 command_rx,
206 client_id,
207 http_client.clone(),
208 app_state.clone(),
209 llm_client.clone(),
210 status_tx.clone(),
211 ));
212 app_state.register_client_task(client_id, cmd_task).await;
213
214 let http_client_clone = http_client.clone();
216 let task_registrar = app_state.clone();
219 let task_handle = tokio::spawn(async move {
220 let protocol = Arc::new(McpClientProtocol::new());
221
222 let instruction = app_state
224 .get_instruction_for_client(client_id)
225 .await
226 .unwrap_or_default();
227 let memory = app_state
228 .get_memory_for_client(client_id)
229 .await
230 .unwrap_or_default();
231
232 match call_llm_for_client(
234 &llm_client,
235 &app_state,
236 client_id.to_string(),
237 &instruction,
238 &memory,
239 Some(&event),
240 protocol.as_ref(),
241 &status_tx,
242 )
243 .await
244 {
245 Ok(result) => {
246 if let Some(mem) = result.memory_updates {
248 app_state.set_memory_for_client(client_id, mem).await;
249 }
250
251 if let Err(e) = Self::execute_llm_actions(
253 client_id,
254 result.actions,
255 &http_client_clone,
256 &llm_client,
257 &app_state,
258 &status_tx,
259 protocol.clone(),
260 )
261 .await
262 {
263 Log::new(Some(&status_tx))
264 .error(format!("Failed to execute LLM actions: {}", e));
265 }
266 }
267 Err(e) => {
268 Log::new(Some(&status_tx)).error(format!("Failed to call LLM: {}", e));
269 }
270 }
271 });
272 task_registrar
273 .register_client_task(client_id, task_handle)
274 .await;
275
276 Ok("0.0.0.0:0".parse().unwrap())
278 }
279
280 async fn send_initialize_request(
286 http_client: &reqwest::Client,
287 base_url: &str,
288 client_id: ClientId,
289 app_state: &Arc<AppState>,
290 status_tx: &mpsc::UnboundedSender<String>,
291 client_name: &str,
292 client_version: &str,
293 ) -> Result<Value> {
294 let request_id: i64 = app_state
295 .with_client_mut(client_id, |client| {
296 let id = client
297 .get_protocol_field("request_id")
298 .and_then(|v| v.as_i64())
299 .unwrap_or(1);
300 client.set_protocol_field("request_id".to_string(), json!(id + 1));
301 id
302 })
303 .await
304 .context("Client not found")?;
305
306 let request = JsonRpcRequest {
307 jsonrpc: "2.0".to_string(),
308 method: "initialize".to_string(),
309 params: Some(json!({
310 "protocolVersion": "2024-11-05",
311 "capabilities": {
312 "roots": {
313 "listChanged": false
314 }
315 },
316 "clientInfo": {
317 "name": client_name,
318 "version": client_version
319 }
320 })),
321 id: request_id,
322 };
323
324 debug!("Sending initialize request: {:?}", request);
325 Log::new(Some(status_tx)).info("Sending MCP initialize request".to_string());
326
327 let response = http_client
328 .post(base_url)
329 .json(&request)
330 .send()
331 .await
332 .context("Failed to send initialize request")?;
333
334 let status = response.status();
335 let response_text = crate::client::http_fetch::read_response_text(
336 response,
337 crate::client::http_fetch::transport::MAX_RESPONSE_BODY_BYTES,
338 )
339 .await?;
340
341 debug!(
342 "Initialize response status: {}, body: {}",
343 status, response_text
344 );
345
346 if !status.is_success() {
347 return Err(anyhow::anyhow!("HTTP error {}: {}", status, response_text));
348 }
349
350 let json_response: JsonRpcResponse =
351 serde_json::from_str(&response_text).context("Failed to parse initialize response")?;
352
353 if let Some(error) = json_response.error {
354 return Err(anyhow::anyhow!(
355 "JSON-RPC error {}: {}",
356 error.code,
357 error.message
358 ));
359 }
360
361 json_response
362 .result
363 .context("Missing result in initialize response")
364 }
365
366 async fn send_initialized_notification(
377 http_client: &reqwest::Client,
378 base_url: &str,
379 status_tx: &mpsc::UnboundedSender<String>,
380 ) -> Result<()> {
381 let notification = JsonRpcNotification {
382 jsonrpc: "2.0".to_string(),
383 method: "notifications/initialized".to_string(),
384 params: Some(json!({})),
385 };
386
387 debug!("Sending initialized notification: {:?}", notification);
388 Log::new(Some(status_tx)).info("Sending MCP initialized notification".to_string());
389
390 http_client
391 .post(base_url)
392 .json(¬ification)
393 .send()
394 .await
395 .context("Failed to send initialized notification")?;
396
397 Ok(())
398 }
399
400 fn execute_llm_actions<'a>(
402 client_id: ClientId,
403 actions: Vec<Value>,
404 http_client: &'a reqwest::Client,
405 llm_client: &'a OllamaClient,
406 app_state: &'a Arc<AppState>,
407 status_tx: &'a mpsc::UnboundedSender<String>,
408 protocol: Arc<McpClientProtocol>,
409 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + 'a>> {
410 Box::pin(async move {
411 for action in actions {
412 debug!("Executing action: {:?}", action);
413
414 let action_result = protocol.as_ref().execute_action(action.clone())?;
416
417 match Self::apply_action(
418 client_id,
419 action_result,
420 Notify::Inline,
421 http_client,
422 llm_client,
423 app_state,
424 status_tx,
425 protocol.clone(),
426 )
427 .await
428 {
429 Ok(Applied::Disconnect) => break,
430 Ok(Applied::Executed(_)) => {}
431 Err(e) => {
432 Log::new(Some(status_tx))
433 .error(format!("Failed to execute MCP action: {}", e));
434 }
435 }
436 }
437
438 Ok(())
439 })
440 }
441
442 #[allow(clippy::too_many_arguments)]
449 fn apply_action<'a>(
450 client_id: ClientId,
451 result: ClientActionResult,
452 notify: Notify,
453 http_client: &'a reqwest::Client,
454 llm_client: &'a OllamaClient,
455 app_state: &'a Arc<AppState>,
456 status_tx: &'a mpsc::UnboundedSender<String>,
457 protocol: Arc<McpClientProtocol>,
458 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Applied>> + Send + 'a>> {
459 Box::pin(async move {
460 match result {
461 ClientActionResult::Disconnect => {
462 app_state
463 .update_client_status(client_id, ClientStatus::Disconnected)
464 .await;
465 app_state.remove_client_handle(client_id).await;
468 Log::new(Some(status_tx))
469 .info(format!("MCP client {} disconnected", client_id));
470 Ok(Applied::Disconnect)
471 }
472 ClientActionResult::Custom { name, data } => {
473 let response = Self::execute_mcp_action(
474 client_id,
475 &name,
476 &data,
477 http_client,
478 app_state,
479 status_tx,
480 )
481 .await?;
482
483 let event_data = json!({
484 "method": name,
485 "result": response,
486 });
487
488 match notify {
489 Notify::Inline => {
490 Self::notify_response(
491 client_id,
492 event_data,
493 http_client,
494 llm_client,
495 app_state,
496 status_tx,
497 protocol.clone(),
498 )
499 .await;
500 }
501 Notify::Deferred => {
502 let http_clone = http_client.clone();
509 let llm_clone = llm_client.clone();
510 let state_clone = app_state.clone();
511 let status_clone = status_tx.clone();
512 let protocol_clone = protocol.clone();
513 let notify_handle = tokio::spawn(async move {
514 McpClient::notify_response(
515 client_id,
516 event_data,
517 &http_clone,
518 &llm_clone,
519 &state_clone,
520 &status_clone,
521 protocol_clone,
522 )
523 .await;
524 });
525 app_state
529 .register_client_task(client_id, notify_handle)
530 .await;
531 }
532 }
533
534 Ok(Applied::Executed(format!("{name} completed")))
535 }
536 other => {
537 debug!("Ignoring non-custom action result: {:?}", other);
538 Ok(Applied::Executed(
539 "action produced no MCP request".to_string(),
540 ))
541 }
542 }
543 })
544 }
545
546 #[allow(clippy::too_many_arguments)]
549 async fn notify_response(
550 client_id: ClientId,
551 event_data: Value,
552 http_client: &reqwest::Client,
553 llm_client: &OllamaClient,
554 app_state: &Arc<AppState>,
555 status_tx: &mpsc::UnboundedSender<String>,
556 protocol: Arc<McpClientProtocol>,
557 ) {
558 let event = Event::new(&MCP_CLIENT_RESPONSE_RECEIVED_EVENT, event_data);
559
560 let instruction = app_state
561 .get_instruction_for_client(client_id)
562 .await
563 .unwrap_or_default();
564 let memory = app_state
565 .get_memory_for_client(client_id)
566 .await
567 .unwrap_or_default();
568
569 match call_llm_for_client(
570 llm_client,
571 app_state,
572 client_id.to_string(),
573 &instruction,
574 &memory,
575 Some(&event),
576 protocol.as_ref(),
577 status_tx,
578 )
579 .await
580 {
581 Ok(result) => {
582 if let Some(mem) = result.memory_updates {
583 app_state.set_memory_for_client(client_id, mem).await;
584 }
585 if let Err(e) = Self::execute_llm_actions(
586 client_id,
587 result.actions,
588 http_client,
589 llm_client,
590 app_state,
591 status_tx,
592 protocol.clone(),
593 )
594 .await
595 {
596 error!("Failed to execute nested actions: {}", e);
597 }
598 }
599 Err(e) => {
600 error!("Failed to call LLM: {}", e);
601 }
602 }
603 }
604
605 async fn command_loop(
614 mut command_rx: tokio::sync::mpsc::Receiver<ClientCommand>,
615 client_id: ClientId,
616 http_client: reqwest::Client,
617 app_state: Arc<AppState>,
618 llm_client: OllamaClient,
619 status_tx: mpsc::UnboundedSender<String>,
620 ) {
621 let protocol = Arc::new(McpClientProtocol::new());
622
623 while let Some(command) = command_rx.recv().await {
624 let action = command.action.clone();
625 let outcome = match protocol.as_ref().execute_action(action.clone()) {
626 Err(e) => Ok(ClientSendOutcome::Rejected {
627 error: e.to_string(),
628 }),
629 Ok(result) => Self::apply_action(
635 client_id,
636 result,
637 Notify::Deferred,
638 &http_client,
639 &llm_client,
640 &app_state,
641 &status_tx,
642 protocol.clone(),
643 )
644 .await
645 .map(|applied| match applied {
646 Applied::Disconnect => ClientSendOutcome::Disconnected,
647 Applied::Executed(detail) => ClientSendOutcome::Executed { detail },
648 }),
649 };
650
651 let outcome_json = match &outcome {
652 Ok(outcome) => serde_json::to_value(outcome).unwrap_or(serde_json::Value::Null),
653 Err(e) => serde_json::json!({"error": e.to_string()}),
654 };
655 app_state
656 .record_access_log(
657 AccessLogOwner::Client(client_id.as_u32()),
658 protocol.protocol_name(),
659 None,
660 "injected_action",
661 action,
662 vec![outcome_json],
663 )
664 .await;
665
666 let disconnect = matches!(outcome, Ok(ClientSendOutcome::Disconnected));
667 if let Err(e) = &outcome {
668 error!("MCP client {} injected action failed: {}", client_id, e);
669 let _ = status_tx.send(format!(
670 "[WARN] Client {} injected action failed: {}",
671 client_id, e
672 ));
673 }
674 let _ = status_tx.send("__UPDATE_UI__".to_string());
675 crate::client::command_support::reply(command, outcome);
676
677 if disconnect {
678 break;
679 }
680 }
681
682 info!("MCP client {} command loop stopped", client_id);
683 app_state.remove_client_handle(client_id).await;
684 let _ = status_tx.send("__UPDATE_UI__".to_string());
685 }
686
687 async fn execute_mcp_action(
689 client_id: ClientId,
690 action_name: &str,
691 data: &Value,
692 http_client: &reqwest::Client,
693 app_state: &Arc<AppState>,
694 status_tx: &mpsc::UnboundedSender<String>,
695 ) -> Result<Value> {
696 let (base_url, request_id) = app_state
705 .with_client_mut(client_id, |client| {
706 let url = client
707 .get_protocol_field("base_url")
708 .and_then(|v| v.as_str())
709 .map(|s| s.to_string());
710
711 let id = client
712 .get_protocol_field("request_id")
713 .and_then(|v| v.as_i64());
714
715 if let Some(id) = id {
716 client.set_protocol_field("request_id".to_string(), json!(id + 1));
717 }
718
719 (url, id)
720 })
721 .await
722 .context("Client not found")?;
723
724 let base_url = base_url.context(
725 "MCP client has no base_url; it was never initialized or its state was reset",
726 )?;
727 let request_id = request_id.context(
728 "MCP client has no request_id; it was never initialized or its state was reset",
729 )?;
730
731 let (method, params) = match action_name {
732 "mcp_list_resources" => ("resources/list".to_string(), None),
733 "mcp_read_resource" => {
734 let uri = data
735 .get("uri")
736 .and_then(|v| v.as_str())
737 .context("Missing uri in read_resource")?;
738 ("resources/read".to_string(), Some(json!({"uri": uri})))
739 }
740 "mcp_list_tools" => ("tools/list".to_string(), None),
741 "mcp_call_tool" => {
742 let name = data
743 .get("name")
744 .and_then(|v| v.as_str())
745 .context("Missing name in call_tool")?;
746 let arguments = data.get("arguments").cloned().unwrap_or(json!({}));
747 (
748 "tools/call".to_string(),
749 Some(json!({
750 "name": name,
751 "arguments": arguments
752 })),
753 )
754 }
755 "mcp_list_prompts" => ("prompts/list".to_string(), None),
756 "mcp_get_prompt" => {
757 let name = data
758 .get("name")
759 .and_then(|v| v.as_str())
760 .context("Missing name in get_prompt")?;
761 let mut params = serde_json::Map::new();
766 params.insert("name".to_string(), json!(name));
767 if let Some(arguments) = data.get("arguments") {
768 if !arguments.is_null() {
769 params.insert("arguments".to_string(), arguments.clone());
770 }
771 }
772 ("prompts/get".to_string(), Some(Value::Object(params)))
773 }
774 _ => return Err(anyhow::anyhow!("Unknown MCP action: {}", action_name)),
775 };
776
777 let request = JsonRpcRequest {
778 jsonrpc: "2.0".to_string(),
779 method: method.clone(),
780 params,
781 id: request_id,
782 };
783
784 debug!("Sending MCP request: {:?}", request);
785 Log::new(Some(status_tx)).info(format!("Sending MCP request: {}", method));
786
787 let response = http_client
788 .post(&base_url)
789 .json(&request)
790 .send()
791 .await
792 .context("Failed to send MCP request")?;
793
794 let status = response.status();
795 let response_text = crate::client::http_fetch::read_response_text(
796 response,
797 crate::client::http_fetch::transport::MAX_RESPONSE_BODY_BYTES,
798 )
799 .await?;
800
801 debug!("MCP response status: {}, body: {}", status, response_text);
802
803 if !status.is_success() {
804 return Err(anyhow::anyhow!("HTTP error {}: {}", status, response_text));
805 }
806
807 let json_response: JsonRpcResponse =
808 serde_json::from_str(&response_text).context("Failed to parse MCP response")?;
809
810 if let Some(error) = json_response.error {
811 return Err(anyhow::anyhow!(
812 "JSON-RPC error {}: {}",
813 error.code,
814 error.message
815 ));
816 }
817
818 json_response
819 .result
820 .context("Missing result in MCP response")
821 }
822}
823
824#[derive(Clone, Copy)]
830enum Notify {
831 Inline,
833 Deferred,
836}
837
838enum Applied {
842 Executed(String),
844 Disconnect,
846}