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    /// Returns a `DbConn` wrapping the pool, allowing shared helper functions
68    /// that accept `DbConn` to work across all context types.
69    pub fn db_conn(&self) -> crate::function::DbConn<'_> {
70        crate::function::DbConn::Pool(&self.db_pool)
71    }
72
73    /// Acquire a connection compatible with sqlx compile-time checked macros.
74    pub async fn conn(&self) -> sqlx::Result<crate::function::ForgeConn<'static>> {
75        Ok(crate::function::ForgeConn::Pool(
76            self.db_pool.acquire().await?,
77        ))
78    }
79
80    pub fn require_user_id(&self) -> Result<Uuid> {
81        self.auth.require_user_id()
82    }
83
84    pub fn require_subject(&self) -> Result<&str> {
85        self.auth.require_subject()
86    }
87
88    /// Dispatch a background job.
89    pub async fn dispatch_job<T: serde::Serialize>(&self, job_type: &str, args: T) -> Result<Uuid> {
90        let dispatcher = self.job_dispatch.as_ref().ok_or_else(|| {
91            crate::error::ForgeError::Internal("Job dispatch not available".to_string())
92        })?;
93
94        let args_json = serde_json::to_value(args)?;
95        dispatcher
96            .dispatch_by_name(job_type, args_json, self.auth.principal_id())
97            .await
98    }
99
100    /// Start a workflow.
101    pub async fn start_workflow<T: serde::Serialize>(
102        &self,
103        workflow_name: &str,
104        input: T,
105    ) -> Result<Uuid> {
106        let dispatcher = self.workflow_dispatch.as_ref().ok_or_else(|| {
107            crate::error::ForgeError::Internal("Workflow dispatch not available".to_string())
108        })?;
109
110        let input_json = serde_json::to_value(input)?;
111        dispatcher
112            .start_by_name(workflow_name, input_json, self.auth.principal_id())
113            .await
114    }
115}
116
117impl EnvAccess for McpToolContext {
118    fn env_provider(&self) -> &dyn EnvProvider {
119        self.env_provider.as_ref()
120    }
121}