1use std::collections::BTreeMap;
2
3use rmcp::handler::server::tool::ToolRouter;
4use rmcp::handler::server::wrapper::Parameters;
5use rmcp::model::*;
6use rmcp::schemars;
7use rmcp::{tool, tool_handler, tool_router, ErrorData as McpError, ServerHandler, ServiceExt};
8use serde::Deserialize;
9use tokio::task::JoinSet;
10
11use crate::fetch::FetchOptions;
12use crate::model::SearchResult;
13use crate::query::{FlightLeg, Passengers, QueryParams, Seat, SearchQuery, TripType};
14
15#[derive(Debug, Deserialize, schemars::JsonSchema)]
16struct SearchArgs {
17 #[schemars(
18 description = "Departure airport IATA code, exactly 3 uppercase letters. Example: HEL, JFK, LAX"
19 )]
20 from: String,
21 #[schemars(
22 description = "Arrival airport IATA code(s). Comma-separate for multi-destination. Examples: BCN or BCN,ATH,AYT"
23 )]
24 to: String,
25 #[schemars(description = "Departure date in YYYY-MM-DD format. Example: 2026-03-01")]
26 date: String,
27 #[schemars(
28 description = "Return date in YYYY-MM-DD for round-trip. Auto-sets trip type to round-trip"
29 )]
30 return_date: Option<String>,
31 #[schemars(
32 description = "One of: economy, premium-economy, business, first. Default: economy"
33 )]
34 seat: Option<String>,
35 #[schemars(description = "Maximum stops. 0 = nonstop only. Omit for any number of stops")]
36 max_stops: Option<u32>,
37 #[schemars(description = "Filter airlines by IATA code, comma-separated. Example: AY,IB")]
38 airlines: Option<String>,
39 #[schemars(description = "Adult passengers (12+). Default: 1")]
40 adults: Option<u32>,
41 #[schemars(description = "Child passengers (2-11). Default: 0")]
42 children: Option<u32>,
43 #[schemars(description = "Infants with own seat (under 2). Default: 0")]
44 infants_in_seat: Option<u32>,
45 #[schemars(description = "Infants on adult's lap (under 2). Default: 0")]
46 infants_on_lap: Option<u32>,
47 #[schemars(description = "Currency code. Examples: USD, EUR, JPY. Default: USD")]
48 currency: Option<String>,
49 #[schemars(description = "Return only N cheapest results")]
50 top: Option<usize>,
51}
52
53#[derive(Debug, Deserialize, schemars::JsonSchema)]
54struct GetUrlArgs {
55 #[schemars(
56 description = "Departure airport IATA code, exactly 3 uppercase letters. Example: HEL, JFK, LAX"
57 )]
58 from: String,
59 #[schemars(
60 description = "Arrival airport IATA code(s). Comma-separate for multi-destination. Examples: BCN or BCN,ATH,AYT"
61 )]
62 to: String,
63 #[schemars(description = "Departure date in YYYY-MM-DD format. Example: 2026-03-01")]
64 date: String,
65 #[schemars(
66 description = "Return date in YYYY-MM-DD for round-trip. Auto-sets trip type to round-trip"
67 )]
68 return_date: Option<String>,
69 #[schemars(
70 description = "One of: economy, premium-economy, business, first. Default: economy"
71 )]
72 seat: Option<String>,
73 #[schemars(description = "Adult passengers (12+). Default: 1")]
74 adults: Option<u32>,
75 #[schemars(description = "Currency code. Examples: USD, EUR, JPY. Default: USD")]
76 currency: Option<String>,
77}
78
79#[derive(Debug, Deserialize, schemars::JsonSchema)]
80struct OpenUrlArgs {
81 #[schemars(description = "URL to open. Must start with http:// or https://")]
82 url: String,
83}
84
85fn parse_legs(
86 from: &str,
87 to: &str,
88 date: &str,
89 return_date: Option<&str>,
90 max_stops: Option<u32>,
91 airlines: Option<&str>,
92) -> (Vec<FlightLeg>, TripType) {
93 let parsed_airlines: Option<Vec<String>> = airlines
94 .map(|s| s.split(',').map(|a| a.trim().to_uppercase()).collect());
95
96 let mut legs = vec![FlightLeg {
97 date: date.to_string(),
98 from_airport: from.to_uppercase(),
99 to_airport: to.to_uppercase(),
100 max_stops,
101 airlines: parsed_airlines.clone(),
102 }];
103
104 let trip = if let Some(ret) = return_date {
105 legs.push(FlightLeg {
106 date: ret.to_string(),
107 from_airport: to.to_uppercase(),
108 to_airport: from.to_uppercase(),
109 max_stops,
110 airlines: parsed_airlines,
111 });
112 TripType::RoundTrip
113 } else {
114 TripType::OneWay
115 };
116
117 (legs, trip)
118}
119
120fn tool_error(msg: impl Into<String>) -> Result<CallToolResult, McpError> {
121 Ok(CallToolResult::error(vec![Content::text(msg.into())]))
122}
123
124fn apply_top(result: &mut SearchResult, n: usize) {
125 result
126 .flights
127 .sort_by_key(|f| f.price.unwrap_or(i64::MAX));
128 result.flights.truncate(n);
129}
130
131#[derive(Debug, Clone)]
132struct FlyrMcp {
133 tool_router: ToolRouter<Self>,
134}
135
136#[tool_router]
137impl FlyrMcp {
138 fn new() -> Self {
139 Self {
140 tool_router: Self::tool_router(),
141 }
142 }
143
144 #[tool(
145 description = "Search for flights and return results as JSON. Searches Google Flights for available flights between airports on specific dates. Returns flight options with prices, airlines, duration, stops, and schedule. Comma-separate 'to' for multi-destination comparison. To open results in browser: call flyr_get_url with the same parameters, then call open_url with the returned URL."
146 )]
147 async fn flyr_search(
148 &self,
149 Parameters(args): Parameters<SearchArgs>,
150 ) -> Result<CallToolResult, McpError> {
151 let is_multi = args.to.contains(',');
152
153 if is_multi {
154 let from = args.from.to_uppercase();
155 let date = args.date;
156
157 let seat = match args
158 .seat
159 .as_deref()
160 .map(Seat::from_str_loose)
161 .transpose()
162 {
163 Ok(s) => s.unwrap_or(Seat::Economy),
164 Err(e) => return tool_error(e.to_string()),
165 };
166
167 let passengers = Passengers {
168 adults: args.adults.unwrap_or(1),
169 children: args.children.unwrap_or(0),
170 infants_in_seat: args.infants_in_seat.unwrap_or(0),
171 infants_on_lap: args.infants_on_lap.unwrap_or(0),
172 };
173
174 let airlines: Option<Vec<String>> = args
175 .airlines
176 .as_ref()
177 .map(|s| s.split(',').map(|a| a.trim().to_uppercase()).collect());
178
179 let currency = args.currency.unwrap_or_else(|| "USD".into());
180
181 let destinations: Vec<String> = args
182 .to
183 .split(',')
184 .map(|s| s.trim().to_uppercase())
185 .filter(|s| !s.is_empty())
186 .collect();
187
188 let mut join_set = JoinSet::new();
189
190 for dest in &destinations {
191 let mut legs = vec![FlightLeg {
192 date: date.clone(),
193 from_airport: from.clone(),
194 to_airport: dest.clone(),
195 max_stops: args.max_stops,
196 airlines: airlines.clone(),
197 }];
198
199 let trip = if let Some(ref ret) = args.return_date {
200 legs.push(FlightLeg {
201 date: ret.clone(),
202 from_airport: dest.clone(),
203 to_airport: from.clone(),
204 max_stops: args.max_stops,
205 airlines: airlines.clone(),
206 });
207 TripType::RoundTrip
208 } else {
209 TripType::OneWay
210 };
211
212 let params = QueryParams {
213 legs,
214 passengers: passengers.clone(),
215 seat: seat.clone(),
216 trip,
217 language: "en".into(),
218 currency: currency.clone(),
219 };
220
221 if let Err(e) = params.validate() {
222 return tool_error(format!("{dest}: {e}"));
223 }
224
225 let dest_code = dest.clone();
226 let top = args.top;
227 join_set.spawn(async move {
228 let result =
229 crate::search(SearchQuery::Structured(params), FetchOptions::default())
230 .await;
231 (dest_code, result, top)
232 });
233 }
234
235 let mut results: BTreeMap<String, SearchResult> = BTreeMap::new();
236 while let Some(join_result) = join_set.join_next().await {
237 let (dest_code, search_result, top) = join_result.unwrap();
238 match search_result {
239 Ok(mut result) => {
240 if let Some(n) = top {
241 apply_top(&mut result, n);
242 }
243 results.insert(dest_code, result);
244 }
245 Err(e) => {
246 results.insert(dest_code.clone(), SearchResult::default());
247 eprintln!("warning: {dest_code}: {e}");
248 }
249 }
250 }
251
252 let json = serde_json::to_string_pretty(&results).unwrap();
253 Ok(CallToolResult::success(vec![Content::text(json)]))
254 } else {
255 let (legs, trip) = parse_legs(
256 &args.from,
257 &args.to,
258 &args.date,
259 args.return_date.as_deref(),
260 args.max_stops,
261 args.airlines.as_deref(),
262 );
263
264 let seat = match args
265 .seat
266 .as_deref()
267 .map(Seat::from_str_loose)
268 .transpose()
269 {
270 Ok(s) => s.unwrap_or(Seat::Economy),
271 Err(e) => return tool_error(e.to_string()),
272 };
273
274 let passengers = Passengers {
275 adults: args.adults.unwrap_or(1),
276 children: args.children.unwrap_or(0),
277 infants_in_seat: args.infants_in_seat.unwrap_or(0),
278 infants_on_lap: args.infants_on_lap.unwrap_or(0),
279 };
280
281 let currency = args.currency.unwrap_or_else(|| "USD".into());
282
283 let params = QueryParams {
284 legs,
285 passengers,
286 seat,
287 trip,
288 language: "en".into(),
289 currency,
290 };
291
292 if let Err(e) = params.validate() {
293 return tool_error(e.to_string());
294 }
295
296 match crate::search(SearchQuery::Structured(params), FetchOptions::default()).await {
297 Ok(mut result) => {
298 if let Some(n) = args.top {
299 apply_top(&mut result, n);
300 }
301 let json = serde_json::to_string_pretty(&result).unwrap();
302 Ok(CallToolResult::success(vec![Content::text(json)]))
303 }
304 Err(e) => tool_error(e.to_string()),
305 }
306 }
307 }
308
309 #[tool(
310 description = "Generate a Google Flights URL for the given search parameters. This is the ONLY way to get a valid Google Flights URL. Returns an encoded URL that can be opened in a browser with open_url. NEVER construct Google Flights URLs manually -- always use this tool."
311 )]
312 async fn flyr_get_url(
313 &self,
314 Parameters(args): Parameters<GetUrlArgs>,
315 ) -> Result<CallToolResult, McpError> {
316 let is_multi = args.to.contains(',');
317
318 if is_multi {
319 let seat = match args
320 .seat
321 .as_deref()
322 .map(Seat::from_str_loose)
323 .transpose()
324 {
325 Ok(s) => s.unwrap_or(Seat::Economy),
326 Err(e) => return tool_error(e.to_string()),
327 };
328
329 let passengers = Passengers {
330 adults: args.adults.unwrap_or(1),
331 ..Default::default()
332 };
333
334 let currency = args.currency.unwrap_or_else(|| "USD".into());
335
336 let destinations: Vec<String> = args
337 .to
338 .split(',')
339 .map(|s| s.trim().to_uppercase())
340 .filter(|s| !s.is_empty())
341 .collect();
342
343 let mut urls = Vec::new();
344 for dest in &destinations {
345 let mut legs = vec![FlightLeg {
346 date: args.date.clone(),
347 from_airport: args.from.to_uppercase(),
348 to_airport: dest.clone(),
349 max_stops: None,
350 airlines: None,
351 }];
352
353 let trip = if let Some(ref ret) = args.return_date {
354 legs.push(FlightLeg {
355 date: ret.clone(),
356 from_airport: dest.clone(),
357 to_airport: args.from.to_uppercase(),
358 max_stops: None,
359 airlines: None,
360 });
361 TripType::RoundTrip
362 } else {
363 TripType::OneWay
364 };
365
366 let params = QueryParams {
367 legs,
368 passengers: passengers.clone(),
369 seat: seat.clone(),
370 trip,
371 language: "en".into(),
372 currency: currency.clone(),
373 };
374
375 if let Err(e) = params.validate() {
376 return tool_error(format!("{dest}: {e}"));
377 }
378
379 urls.push(crate::generate_browser_url(¶ms));
380 }
381
382 Ok(CallToolResult::success(vec![Content::text(
383 urls.join("\n"),
384 )]))
385 } else {
386 let (legs, trip) = parse_legs(
387 &args.from,
388 &args.to,
389 &args.date,
390 args.return_date.as_deref(),
391 None,
392 None,
393 );
394
395 let seat = match args
396 .seat
397 .as_deref()
398 .map(Seat::from_str_loose)
399 .transpose()
400 {
401 Ok(s) => s.unwrap_or(Seat::Economy),
402 Err(e) => return tool_error(e.to_string()),
403 };
404
405 let passengers = Passengers {
406 adults: args.adults.unwrap_or(1),
407 ..Default::default()
408 };
409
410 let currency = args.currency.unwrap_or_else(|| "USD".into());
411
412 let params = QueryParams {
413 legs,
414 passengers,
415 seat,
416 trip,
417 language: "en".into(),
418 currency,
419 };
420
421 if let Err(e) = params.validate() {
422 return tool_error(e.to_string());
423 }
424
425 let url = crate::generate_browser_url(¶ms);
426 Ok(CallToolResult::success(vec![Content::text(url)]))
427 }
428 }
429
430 #[tool(description = "Open a URL in the default web browser. IMPORTANT: To open flight results, you MUST call flyr_get_url first to get the URL, then pass that URL here. NEVER construct Google Flights URLs yourself -- they require special encoding that only flyr_get_url can produce.")]
431 async fn open_url(
432 &self,
433 Parameters(args): Parameters<OpenUrlArgs>,
434 ) -> Result<CallToolResult, McpError> {
435 if !args.url.starts_with("http://") && !args.url.starts_with("https://") {
436 return tool_error("URL must start with http:// or https://");
437 }
438
439 match open::that(&args.url) {
440 Ok(()) => Ok(CallToolResult::success(vec![Content::text(format!(
441 "Opened: {}",
442 args.url
443 ))])),
444 Err(e) => tool_error(format!("failed to open browser: {e}")),
445 }
446 }
447}
448
449#[tool_handler]
450impl ServerHandler for FlyrMcp {
451 fn get_info(&self) -> ServerInfo {
452 ServerInfo {
453 protocol_version: ProtocolVersion::V_2024_11_05,
454 capabilities: ServerCapabilities::builder().enable_tools().build(),
455 server_info: Implementation {
456 name: "flyr".into(),
457 version: env!("CARGO_PKG_VERSION").into(),
458 ..Default::default()
459 },
460 instructions: Some(
461 "Flight search tool. Workflow: (1) flyr_search to find flights. (2) To open in browser: call flyr_get_url with same params to get URL, then call open_url with that URL. NEVER construct Google Flights URLs yourself -- they require special protobuf encoding.".into(),
462 ),
463 }
464 }
465}
466
467pub async fn run() {
468 let service = FlyrMcp::new()
469 .serve(rmcp::transport::stdio())
470 .await
471 .expect("failed to start MCP server");
472 service.waiting().await.expect("MCP server error");
473}
474
475#[cfg(test)]
476mod tests {
477 use super::*;
478
479 #[test]
480 fn parse_legs_from_to_date() {
481 let (legs, trip) = parse_legs("HEL", "BCN", "2026-03-01", None, None, None);
482 assert_eq!(legs.len(), 1);
483 assert_eq!(legs[0].from_airport, "HEL");
484 assert_eq!(legs[0].to_airport, "BCN");
485 assert_eq!(legs[0].date, "2026-03-01");
486 assert!(matches!(trip, TripType::OneWay));
487 }
488
489 #[test]
490 fn parse_legs_with_return_date() {
491 let (legs, trip) =
492 parse_legs("HEL", "BCN", "2026-03-01", Some("2026-03-08"), None, None);
493 assert_eq!(legs.len(), 2);
494 assert_eq!(legs[0].from_airport, "HEL");
495 assert_eq!(legs[0].to_airport, "BCN");
496 assert_eq!(legs[1].from_airport, "BCN");
497 assert_eq!(legs[1].to_airport, "HEL");
498 assert_eq!(legs[1].date, "2026-03-08");
499 assert!(matches!(trip, TripType::RoundTrip));
500 }
501
502 #[test]
503 fn parse_legs_with_airlines() {
504 let (legs, _) = parse_legs("HEL", "BCN", "2026-03-01", None, Some(1), Some("AY,IB"));
505 assert_eq!(legs[0].max_stops, Some(1));
506 assert_eq!(
507 legs[0].airlines,
508 Some(vec!["AY".to_string(), "IB".to_string()])
509 );
510 }
511}