Skip to main content

forge_core/mcp/
context.rs

1use std::sync::Arc;
2
3use crate::Result;
4use crate::env::{EnvAccess, EnvProvider, RealEnvProvider};
5use crate::function::{AuthContext, JobDispatch, RequestMetadata, WorkflowDispatch};
6use uuid::Uuid;
7
8/// Context for MCP tool execution.
9pub struct McpToolContext {
10    /// Authentication context.
11    pub auth: AuthContext,
12    /// Request metadata.
13    pub request: RequestMetadata,
14    db_pool: sqlx::PgPool,
15    job_dispatch: Option<Arc<dyn JobDispatch>>,
16    workflow_dispatch: Option<Arc<dyn WorkflowDispatch>>,
17    env_provider: Arc<dyn EnvProvider>,
18}
19
20impl McpToolContext {
21    /// Create a new MCP tool context.
22    pub fn new(db_pool: sqlx::PgPool, auth: AuthContext, request: RequestMetadata) -> Self {
23        Self::with_dispatch(db_pool, auth, request, None, None)
24    }
25
26    /// Create a context with dispatch capabilities.
27    pub fn with_dispatch(
28        db_pool: sqlx::PgPool,
29        auth: AuthContext,
30        request: RequestMetadata,
31        job_dispatch: Option<Arc<dyn JobDispatch>>,
32        workflow_dispatch: Option<Arc<dyn WorkflowDispatch>>,
33    ) -> Self {
34        Self::with_env(
35            db_pool,
36            auth,
37            request,
38            job_dispatch,
39            workflow_dispatch,
40            Arc::new(RealEnvProvider::new()),
41        )
42    }
43
44    /// Create a context with a custom environment provider.
45    pub fn with_env(
46        db_pool: sqlx::PgPool,
47        auth: AuthContext,
48        request: RequestMetadata,
49        job_dispatch: Option<Arc<dyn JobDispatch>>,
50        workflow_dispatch: Option<Arc<dyn WorkflowDispatch>>,
51        env_provider: Arc<dyn EnvProvider>,
52    ) -> Self {
53        Self {
54            auth,
55            request,
56            db_pool,
57            job_dispatch,
58            workflow_dispatch,
59            env_provider,
60        }
61    }
62
63    pub fn db(&self) -> &sqlx::PgPool {
64        &self.db_pool
65    }
66
67    pub fn require_user_id(&self) -> Result<Uuid> {
68        self.auth.require_user_id()
69    }
70
71    pub fn require_subject(&self) -> Result<&str> {
72        self.auth.require_subject()
73    }
74
75    /// Dispatch a background job.
76    pub async fn dispatch_job<T: serde::Serialize>(&self, job_type: &str, args: T) -> Result<Uuid> {
77        let dispatcher = self.job_dispatch.as_ref().ok_or_else(|| {
78            crate::error::ForgeError::Internal("Job dispatch not available".to_string())
79        })?;
80
81        let args_json = serde_json::to_value(args)?;
82        dispatcher
83            .dispatch_by_name(job_type, args_json, self.auth.principal_id())
84            .await
85    }
86
87    /// Start a workflow.
88    pub async fn start_workflow<T: serde::Serialize>(
89        &self,
90        workflow_name: &str,
91        input: T,
92    ) -> Result<Uuid> {
93        let dispatcher = self.workflow_dispatch.as_ref().ok_or_else(|| {
94            crate::error::ForgeError::Internal("Workflow dispatch not available".to_string())
95        })?;
96
97        let input_json = serde_json::to_value(input)?;
98        dispatcher
99            .start_by_name(workflow_name, input_json, self.auth.principal_id())
100            .await
101    }
102}
103
104impl EnvAccess for McpToolContext {
105    fn env_provider(&self) -> &dyn EnvProvider {
106        self.env_provider.as_ref()
107    }
108}