soma_gateway/gateway/manager/
mcp_tasks.rs

1use std::{
2    collections::BTreeMap,
3    sync::{Arc, atomic::Ordering},
4};
5
6use serde_json::Value;
7
8use crate::upstream::{McpRequestOutcome, UpstreamError};
9
10use super::{GatewayManager, GatewayManagerError, TaskRoute};
11
12impl GatewayManager {
13    pub(super) fn register_task_outcome(
14        &self,
15        outcome: McpRequestOutcome,
16        upstream: &str,
17        subject: Option<&str>,
18    ) -> Result<McpRequestOutcome, GatewayManagerError> {
19        let McpRequestOutcome::Task(mut value) = outcome else {
20            return Ok(outcome);
21        };
22        let native_task_id = value
23            .get("taskId")
24            .and_then(Value::as_str)
25            .filter(|task_id| !task_id.is_empty())
26            .ok_or_else(|| {
27                GatewayManagerError::InvalidTaskResult(
28                    "task response did not contain a non-empty taskId".to_owned(),
29                )
30            })?
31            .to_owned();
32        let public_task_id = format!(
33            "soma-task-{:016x}",
34            self.task_sequence.fetch_add(1, Ordering::Relaxed)
35        );
36        self.task_routes
37            .write()
38            .expect("gateway task routes poisoned")
39            .insert(
40                public_task_id.clone(),
41                TaskRoute {
42                    upstream: upstream.to_owned(),
43                    native_task_id,
44                    subject: subject.map(str::to_owned),
45                },
46            );
47        value["taskId"] = Value::String(public_task_id);
48        Ok(McpRequestOutcome::Task(value))
49    }
50
51    pub async fn get_mcp_task_for_subject(
52        &self,
53        task_id: &str,
54        subject: Option<&str>,
55    ) -> Result<Value, GatewayManagerError> {
56        let route = self.resolve_task_route(task_id, subject)?;
57        let pool = self.task_pool()?;
58        let mut value = get_native_task(&pool, &route).await?;
59        value["taskId"] = Value::String(task_id.to_owned());
60        Ok(value)
61    }
62
63    pub async fn update_mcp_task_for_subject(
64        &self,
65        task_id: &str,
66        input_responses: BTreeMap<String, Value>,
67        subject: Option<&str>,
68    ) -> Result<(), GatewayManagerError> {
69        let route = self.resolve_task_route(task_id, subject)?;
70        let pool = self.task_pool()?;
71        update_native_task(&pool, &route, input_responses).await?;
72        Ok(())
73    }
74
75    pub async fn cancel_mcp_task_for_subject(
76        &self,
77        task_id: &str,
78        subject: Option<&str>,
79    ) -> Result<(), GatewayManagerError> {
80        let route = self.resolve_task_route(task_id, subject)?;
81        let pool = self.task_pool()?;
82        cancel_native_task(&pool, &route).await?;
83        Ok(())
84    }
85
86    fn task_pool(&self) -> Result<Arc<crate::upstream::pool::UpstreamPool>, GatewayManagerError> {
87        self.ensure_ready()?;
88        Ok(self.pool.read().expect("gateway pool poisoned").clone())
89    }
90
91    fn resolve_task_route(
92        &self,
93        task_id: &str,
94        subject: Option<&str>,
95    ) -> Result<TaskRoute, GatewayManagerError> {
96        let route = self
97            .task_routes
98            .read()
99            .expect("gateway task routes poisoned")
100            .get(task_id)
101            .cloned()
102            .ok_or_else(|| GatewayManagerError::TaskMissing(task_id.to_owned()))?;
103        if route.subject.as_deref() != subject {
104            return Err(GatewayManagerError::TaskMissing(task_id.to_owned()));
105        }
106        Ok(route)
107    }
108}
109
110async fn get_native_task(
111    pool: &crate::upstream::pool::UpstreamPool,
112    route: &TaskRoute,
113) -> Result<Value, UpstreamError> {
114    #[cfg(feature = "oauth")]
115    if route.subject.is_some() {
116        return pool
117            .get_task_for_subject(
118                &route.upstream,
119                &route.native_task_id,
120                route.subject.as_deref(),
121            )
122            .await;
123    }
124    pool.get_task(&route.upstream, &route.native_task_id).await
125}
126
127async fn update_native_task(
128    pool: &crate::upstream::pool::UpstreamPool,
129    route: &TaskRoute,
130    input_responses: BTreeMap<String, Value>,
131) -> Result<(), UpstreamError> {
132    #[cfg(feature = "oauth")]
133    if route.subject.is_some() {
134        return pool
135            .update_task_for_subject(
136                &route.upstream,
137                &route.native_task_id,
138                input_responses,
139                route.subject.as_deref(),
140            )
141            .await;
142    }
143    pool.update_task(&route.upstream, &route.native_task_id, input_responses)
144        .await
145}
146
147async fn cancel_native_task(
148    pool: &crate::upstream::pool::UpstreamPool,
149    route: &TaskRoute,
150) -> Result<(), UpstreamError> {
151    #[cfg(feature = "oauth")]
152    if route.subject.is_some() {
153        return pool
154            .cancel_task_for_subject(
155                &route.upstream,
156                &route.native_task_id,
157                route.subject.as_deref(),
158            )
159            .await;
160    }
161    pool.cancel_task(&route.upstream, &route.native_task_id)
162        .await
163}
164
165#[cfg(test)]
166#[path = "mcp_tasks_tests.rs"]
167mod tests;
soma_gateway/gateway/manager/mcp_tasks.rs

soma_gateway/gateway/manager/
mcp_tasks.rs
