1use serde_json::{Map, Value};
2
3use crate::oauth::UpstreamOAuthProvider;
4use crate::process::guard::SpawnGuard;
5use crate::upstream::{
6 CapScope, McpRequestOutcome, McpRoundTrip, PromptDescriptor, ResourceDescriptor,
7 ToolDescriptor, UpstreamError, UpstreamHealth, UpstreamSnapshot,
8};
9
10use super::tools::matches_filter;
11use super::{SubjectPoolEntry, ToolCall, UpstreamPool, live};
12
13impl UpstreamPool {
14 pub fn install_oauth_provider(&self, provider: std::sync::Arc<dyn UpstreamOAuthProvider>) {
15 *self
16 .oauth_provider
17 .write()
18 .expect("oauth provider lock poisoned") = Some(provider);
19 self.subject_entries
20 .write()
21 .expect("subject pool lock poisoned")
22 .clear();
23 }
24
25 pub fn evict_oauth_subject(&self, upstream: &str, subject: &str) {
26 self.subject_entries
27 .write()
28 .expect("subject pool lock poisoned")
29 .remove(&(upstream.to_owned(), subject.to_owned()));
30 if let Some(provider) = self
31 .oauth_provider
32 .read()
33 .expect("oauth provider lock poisoned")
34 .as_ref()
35 {
36 provider.evict_subject(upstream, subject);
37 }
38 }
39
40 pub async fn discover_for_subject(
41 &self,
42 subject: Option<&str>,
43 ) -> Result<Vec<UpstreamSnapshot>, UpstreamError> {
44 let Some(subject) = subject else {
45 return self.discover().await;
46 };
47 let configs = self.configured_upstreams();
48 for (name, oauth_enabled) in configs {
49 let result = if oauth_enabled {
50 self.ensure_subject_connected(&name, subject).await
51 } else {
52 self.ensure_connected(&name).await
53 };
54 if let Err(error) = result {
55 tracing::warn!(upstream = %name, subject, error = %error, "subject discovery failed");
56 }
57 }
58 Ok(self.snapshots_for_subject(subject))
59 }
60
61 pub async fn exposed_tools_for_subject(
62 &self,
63 upstream: &str,
64 subject: Option<&str>,
65 ) -> Result<Vec<ToolDescriptor>, UpstreamError> {
66 let Some(subject) = subject else {
67 return self.exposed_tools(upstream);
68 };
69 let Some((snapshot, config)) = self.subject_snapshot_and_config(upstream, subject)? else {
70 return self.exposed_tools(upstream);
71 };
72 let tools: Vec<ToolDescriptor> = snapshot
73 .tools
74 .into_iter()
75 .filter(|tool| matches_filter(config.expose_tools.as_deref(), &tool.name))
76 .collect();
77 let bytes = serde_json::to_vec(&tools).map_or(usize::MAX, |bytes| bytes.len());
78 self.response_caps().enforce(CapScope::ToolsList, bytes)?;
79 Ok(tools)
80 }
81
82 pub async fn call_tool_for_subject(
83 &self,
84 call: ToolCall,
85 subject: Option<&str>,
86 ) -> Result<Value, UpstreamError> {
87 let Some(subject) = subject else {
88 return self.call_tool(call).await;
89 };
90 if !self.config_is_oauth(&call.upstream)? {
91 return self.call_tool(call).await;
92 }
93 self.ensure_subject_connected(&call.upstream, subject)
94 .await?;
95 let (peer, upstream) = self.with_subject_entry(&call.upstream, subject, |entry| {
96 ensure_subject_routable(&entry.snapshot)?;
97 if !entry
98 .snapshot
99 .tools
100 .iter()
101 .any(|candidate| candidate.name == call.tool)
102 {
103 return Err(UpstreamError::NotExposed {
104 upstream: call.upstream.clone(),
105 item: call.tool.clone(),
106 });
107 }
108 Ok((entry.live.peer(), entry.snapshot.name.clone()))
109 })?;
110 let result = live::call_live_tool(&upstream, peer, call.tool, call.params).await?;
111 let bytes = serde_json::to_vec(&result).map_or(usize::MAX, |bytes| bytes.len());
112 self.response_caps().enforce(CapScope::ToolsCall, bytes)?;
113 Ok(result)
114 }
115
116 pub async fn call_tool_once_for_subject(
117 &self,
118 call: ToolCall,
119 round_trip: McpRoundTrip,
120 subject: Option<&str>,
121 ) -> Result<McpRequestOutcome, UpstreamError> {
122 let Some(subject) = subject else {
123 return self.call_tool_once(call, round_trip).await;
124 };
125 if !self.config_is_oauth(&call.upstream)? {
126 return self.call_tool_once(call, round_trip).await;
127 }
128 self.ensure_subject_connected(&call.upstream, subject)
129 .await?;
130 let (peer, upstream) = self.with_subject_entry(&call.upstream, subject, |entry| {
131 ensure_subject_routable(&entry.snapshot)?;
132 if !entry
133 .snapshot
134 .tools
135 .iter()
136 .any(|candidate| candidate.name == call.tool)
137 {
138 return Err(UpstreamError::NotExposed {
139 upstream: call.upstream.clone(),
140 item: call.tool.clone(),
141 });
142 }
143 Ok((entry.live.peer(), entry.snapshot.name.clone()))
144 })?;
145 let outcome = if round_trip.request_meta.is_some() {
146 let config = self.config_for_subject(&call.upstream)?;
147 let provider = self.oauth_provider()?;
148 live::call_live_tool_once_scoped(
149 &config,
150 live::LiveConnectContext::oauth(self.response_caps(), subject, provider),
151 call.tool,
152 call.params,
153 round_trip,
154 )
155 .await?
156 } else {
157 live::call_live_tool_once(&upstream, peer, call.tool, call.params, round_trip).await?
158 };
159 let bytes = serde_json::to_vec(outcome.payload()).map_or(usize::MAX, |bytes| bytes.len());
160 self.response_caps().enforce(CapScope::ToolsCall, bytes)?;
161 Ok(outcome)
162 }
163
164 pub async fn list_resources_for_subject(
165 &self,
166 upstream: &str,
167 subject: Option<&str>,
168 ) -> Result<Vec<ResourceDescriptor>, UpstreamError> {
169 let Some(subject) = subject else {
170 return self.list_resources(upstream).await;
171 };
172 let Some((snapshot, config)) = self.subject_snapshot_and_config(upstream, subject)? else {
173 return self.list_resources(upstream).await;
174 };
175 if !config.proxy_resources {
176 return Ok(Vec::new());
177 }
178 let resources: Vec<ResourceDescriptor> = snapshot
179 .resources
180 .into_iter()
181 .filter(|resource| matches_filter(config.expose_resources.as_deref(), &resource.uri))
182 .collect();
183 let bytes = serde_json::to_vec(&resources).map_or(usize::MAX, |bytes| bytes.len());
184 self.response_caps()
185 .enforce(CapScope::ResourcesList, bytes)?;
186 Ok(resources)
187 }
188
189 pub async fn read_resource_for_subject(
190 &self,
191 upstream: &str,
192 uri: &str,
193 subject: Option<&str>,
194 ) -> Result<Value, UpstreamError> {
195 let Some(subject) = subject else {
196 return self.read_resource(upstream, uri).await;
197 };
198 if !self.config_is_oauth(upstream)? {
199 return self.read_resource(upstream, uri).await;
200 }
201 self.ensure_subject_connected(upstream, subject).await?;
202 let peer = self.with_subject_entry(upstream, subject, |entry| {
203 ensure_subject_routable(&entry.snapshot)?;
204 Ok(entry.live.peer())
205 })?;
206 let value = live::read_live_resource(upstream, peer, uri.to_owned()).await?;
207 let bytes = serde_json::to_vec(&value).map_or(usize::MAX, |bytes| bytes.len());
208 self.response_caps()
209 .enforce(CapScope::ResourcesRead, bytes)?;
210 Ok(value)
211 }
212
213 pub async fn read_resource_once_for_subject(
214 &self,
215 upstream: &str,
216 uri: &str,
217 round_trip: McpRoundTrip,
218 subject: Option<&str>,
219 ) -> Result<McpRequestOutcome, UpstreamError> {
220 let Some(subject) = subject else {
221 return self.read_resource_once(upstream, uri, round_trip).await;
222 };
223 if !self.config_is_oauth(upstream)? {
224 return self.read_resource_once(upstream, uri, round_trip).await;
225 }
226 self.ensure_subject_connected(upstream, subject).await?;
227 let peer = self.with_subject_entry(upstream, subject, |entry| {
228 ensure_subject_routable(&entry.snapshot)?;
229 Ok(entry.live.peer())
230 })?;
231 let outcome =
232 live::read_live_resource_once(upstream, peer, uri.to_owned(), round_trip).await?;
233 let bytes = serde_json::to_vec(outcome.payload()).map_or(usize::MAX, |bytes| bytes.len());
234 self.response_caps()
235 .enforce(CapScope::ResourcesRead, bytes)?;
236 Ok(outcome)
237 }
238
239 pub async fn list_prompts_for_subject(
240 &self,
241 upstream: &str,
242 subject: Option<&str>,
243 ) -> Result<Vec<PromptDescriptor>, UpstreamError> {
244 let Some(subject) = subject else {
245 return self.list_prompts(upstream).await;
246 };
247 let Some((snapshot, config)) = self.subject_snapshot_and_config(upstream, subject)? else {
248 return self.list_prompts(upstream).await;
249 };
250 if !config.proxy_prompts {
251 return Ok(Vec::new());
252 }
253 let prompts: Vec<PromptDescriptor> = snapshot
254 .prompts
255 .into_iter()
256 .filter(|prompt| matches_filter(config.expose_prompts.as_deref(), &prompt.name))
257 .collect();
258 let bytes = serde_json::to_vec(&prompts).map_or(usize::MAX, |bytes| bytes.len());
259 self.response_caps().enforce(CapScope::PromptsList, bytes)?;
260 Ok(prompts)
261 }
262
263 pub async fn get_prompt_for_subject(
264 &self,
265 upstream: &str,
266 name: &str,
267 arguments: Option<Map<String, Value>>,
268 subject: Option<&str>,
269 ) -> Result<Value, UpstreamError> {
270 let Some(subject) = subject else {
271 return self.get_prompt(upstream, name, arguments).await;
272 };
273 if !self.config_is_oauth(upstream)? {
274 return self.get_prompt(upstream, name, arguments).await;
275 }
276 self.ensure_subject_connected(upstream, subject).await?;
277 let peer = self.with_subject_entry(upstream, subject, |entry| {
278 ensure_subject_routable(&entry.snapshot)?;
279 Ok(entry.live.peer())
280 })?;
281 let value = live::get_live_prompt(upstream, peer, name.to_owned(), arguments).await?;
282 let bytes = serde_json::to_vec(&value).map_or(usize::MAX, |bytes| bytes.len());
283 self.response_caps().enforce(CapScope::PromptsGet, bytes)?;
284 Ok(value)
285 }
286
287 pub async fn get_prompt_once_for_subject(
288 &self,
289 upstream: &str,
290 name: &str,
291 arguments: Option<Map<String, Value>>,
292 round_trip: McpRoundTrip,
293 subject: Option<&str>,
294 ) -> Result<McpRequestOutcome, UpstreamError> {
295 let Some(subject) = subject else {
296 return self
297 .get_prompt_once(upstream, name, arguments, round_trip)
298 .await;
299 };
300 if !self.config_is_oauth(upstream)? {
301 return self
302 .get_prompt_once(upstream, name, arguments, round_trip)
303 .await;
304 }
305 self.ensure_subject_connected(upstream, subject).await?;
306 let peer = self.with_subject_entry(upstream, subject, |entry| {
307 ensure_subject_routable(&entry.snapshot)?;
308 Ok(entry.live.peer())
309 })?;
310 let outcome =
311 live::get_live_prompt_once(upstream, peer, name.to_owned(), arguments, round_trip)
312 .await?;
313 let bytes = serde_json::to_vec(outcome.payload()).map_or(usize::MAX, |bytes| bytes.len());
314 self.response_caps().enforce(CapScope::PromptsGet, bytes)?;
315 Ok(outcome)
316 }
317
318 pub async fn get_task_for_subject(
319 &self,
320 upstream: &str,
321 task_id: &str,
322 subject: Option<&str>,
323 ) -> Result<Value, UpstreamError> {
324 let Some(subject) = subject else {
325 return self.get_task(upstream, task_id).await;
326 };
327 if !self.config_is_oauth(upstream)? {
328 return self.get_task(upstream, task_id).await;
329 }
330 self.ensure_subject_connected(upstream, subject).await?;
331 let peer = self.with_subject_entry(upstream, subject, |entry| {
332 ensure_subject_routable(&entry.snapshot)?;
333 Ok(entry.live.peer())
334 })?;
335 let result = peer
336 .get_task(rmcp::model::GetTaskParams::new(task_id))
337 .await
338 .map_err(|error| UpstreamError::LiveCall {
339 upstream: upstream.to_owned(),
340 operation: "tasks/get",
341 message: error.to_string(),
342 })?;
343 serde_json::to_value(result).map_err(|error| UpstreamError::LiveCall {
344 upstream: upstream.to_owned(),
345 operation: "tasks/get",
346 message: error.to_string(),
347 })
348 }
349
350 pub async fn update_task_for_subject(
351 &self,
352 upstream: &str,
353 task_id: &str,
354 input_responses: std::collections::BTreeMap<String, Value>,
355 subject: Option<&str>,
356 ) -> Result<(), UpstreamError> {
357 let Some(subject) = subject else {
358 return self.update_task(upstream, task_id, input_responses).await;
359 };
360 if !self.config_is_oauth(upstream)? {
361 return self.update_task(upstream, task_id, input_responses).await;
362 }
363 self.ensure_subject_connected(upstream, subject).await?;
364 let peer = self.with_subject_entry(upstream, subject, |entry| {
365 ensure_subject_routable(&entry.snapshot)?;
366 Ok(entry.live.peer())
367 })?;
368 peer.update_task(rmcp::model::UpdateTaskParams::new(task_id, input_responses))
369 .await
370 .map_err(|error| UpstreamError::LiveCall {
371 upstream: upstream.to_owned(),
372 operation: "tasks/update",
373 message: error.to_string(),
374 })
375 }
376
377 pub async fn cancel_task_for_subject(
378 &self,
379 upstream: &str,
380 task_id: &str,
381 subject: Option<&str>,
382 ) -> Result<(), UpstreamError> {
383 let Some(subject) = subject else {
384 return self.cancel_task(upstream, task_id).await;
385 };
386 if !self.config_is_oauth(upstream)? {
387 return self.cancel_task(upstream, task_id).await;
388 }
389 self.ensure_subject_connected(upstream, subject).await?;
390 let peer = self.with_subject_entry(upstream, subject, |entry| {
391 ensure_subject_routable(&entry.snapshot)?;
392 Ok(entry.live.peer())
393 })?;
394 peer.cancel_task(rmcp::model::CancelTaskParams::new(task_id))
395 .await
396 .map_err(|error| UpstreamError::LiveCall {
397 upstream: upstream.to_owned(),
398 operation: "tasks/cancel",
399 message: error.to_string(),
400 })
401 }
402
403 async fn ensure_subject_connected(
404 &self,
405 upstream: &str,
406 subject: &str,
407 ) -> Result<(), UpstreamError> {
408 let key = (upstream.to_owned(), subject.to_owned());
409 if self
410 .subject_entries
411 .read()
412 .expect("subject pool lock poisoned")
413 .contains_key(&key)
414 {
415 return Ok(());
416 }
417 let config = self.config_for_subject(upstream)?;
418 if !config.enabled {
419 return Ok(());
420 }
421 if config.oauth.is_none() {
422 return self.ensure_connected(upstream).await;
423 }
424 let provider = self.oauth_provider()?;
425 let context = live::LiveConnectContext::oauth(self.response_caps(), subject, provider);
426 let (live, snapshot) = live::connect_live(&config, &SpawnGuard::default(), context).await?;
427 self.subject_entries
428 .write()
429 .expect("subject pool lock poisoned")
430 .insert(
431 key,
432 SubjectPoolEntry {
433 snapshot,
434 live: std::sync::Arc::new(live),
435 },
436 );
437 Ok(())
438 }
439
440 fn configured_upstreams(&self) -> Vec<(String, bool)> {
441 self.entries
442 .read()
443 .expect("upstream pool lock poisoned")
444 .iter()
445 .map(|(name, entry)| (name.clone(), entry.config.oauth.is_some()))
446 .collect()
447 }
448
449 fn snapshots_for_subject(&self, subject: &str) -> Vec<UpstreamSnapshot> {
450 let entries = self.entries.read().expect("upstream pool lock poisoned");
451 let subject_entries = self
452 .subject_entries
453 .read()
454 .expect("subject pool lock poisoned");
455 entries
456 .iter()
457 .filter_map(|(name, entry)| {
458 if entry.config.oauth.is_some() {
459 return subject_entries
460 .get(&(name.clone(), subject.to_owned()))
461 .map(|entry| entry.snapshot.clone());
462 }
463 Some(entry.snapshot.clone())
464 })
465 .collect()
466 }
467
468 fn subject_snapshot_and_config(
469 &self,
470 upstream: &str,
471 subject: &str,
472 ) -> Result<Option<(UpstreamSnapshot, crate::config::UpstreamConfig)>, UpstreamError> {
473 if !self.config_is_oauth(upstream)? {
474 return Ok(None);
475 }
476 let snapshot =
477 self.with_subject_entry(upstream, subject, |entry| Ok(entry.snapshot.clone()))?;
478 let config = self.config_for_subject(upstream)?;
479 Ok(Some((snapshot, config)))
480 }
481
482 fn with_subject_entry<T>(
483 &self,
484 upstream: &str,
485 subject: &str,
486 f: impl FnOnce(&SubjectPoolEntry) -> Result<T, UpstreamError>,
487 ) -> Result<T, UpstreamError> {
488 let entries = self
489 .subject_entries
490 .read()
491 .expect("subject pool lock poisoned");
492 let entry = entries
493 .get(&(upstream.to_owned(), subject.to_owned()))
494 .ok_or_else(|| UpstreamError::NotRoutable {
495 upstream: upstream.to_owned(),
496 reason: "subject connection is not established".to_owned(),
497 })?;
498 f(entry)
499 }
500
501 fn oauth_provider(&self) -> Result<std::sync::Arc<dyn UpstreamOAuthProvider>, UpstreamError> {
502 self.oauth_provider
503 .read()
504 .expect("oauth provider lock poisoned")
505 .clone()
506 .ok_or_else(|| UpstreamError::LiveConnect {
507 upstream: "oauth".to_owned(),
508 message: "upstream OAuth runtime is not configured".to_owned(),
509 })
510 }
511
512 fn config_for_subject(
513 &self,
514 upstream: &str,
515 ) -> Result<crate::config::UpstreamConfig, UpstreamError> {
516 self.with_entry(upstream, |entry| Ok(entry.config.clone()))
517 }
518
519 fn config_is_oauth(&self, upstream: &str) -> Result<bool, UpstreamError> {
520 self.with_entry(upstream, |entry| Ok(entry.config.oauth.is_some()))
521 }
522}
523
524fn ensure_subject_routable(snapshot: &UpstreamSnapshot) -> Result<(), UpstreamError> {
525 if snapshot.health == UpstreamHealth::Connected {
526 return Ok(());
527 }
528 Err(UpstreamError::NotRoutable {
529 upstream: snapshot.name.clone(),
530 reason: "subject-scoped upstream is not connected".to_owned(),
531 })
532}
533
534#[cfg(test)]
535#[path = "subject_tests.rs"]
536mod tests;