ares_agent/
context_provider.rs

1//! External context injection for agents.
2//!
3//! The `ContextProvider` trait allows external systems to inject context
4//! into agent calls before LLM invocation. This is the extension point
5//! that separates generic ARES from managed platform features.
6//!
7//! ## OSS Mode
8//!
9//! By default, ARES uses `NoOpContextProvider` which returns `None`.
10//! Agents run with only their system prompt — no external context.
11//!
12//! ## Managed Mode
13//!
14//! Platform extensions (e.g., dirmacs-core) implement `ContextProvider`
15//! to inject knowledge states, gap constraints, or any external context
16//! into the system prompt before every LLM call.
17
18use async_trait::async_trait;
19
20/// Runtime metadata available to managed context providers.
21///
22/// Public ARES treats these fields as caller-supplied metadata. Managed
23/// providers must still validate workspace use against their own binding or
24/// auth policy before using it for external memory fetches.
25#[derive(Debug, Clone, Default, PartialEq, Eq)]
26pub struct AgentRuntimeContext {
27    /// Tenant that owns the agent execution.
28    pub tenant_id: String,
29    /// Registry or tenant-agent name being executed.
30    pub agent_name: String,
31    /// Optional workspace selected by the authenticated upstream runtime.
32    pub workspace_id: Option<String>,
33    /// Optional end-user identifier for user-scoped products.
34    pub user_id: Option<String>,
35    /// Optional session or conversation identifier for this run.
36    pub session_id: Option<String>,
37    /// Logical source of the request, such as an API handler name.
38    pub request_source: String,
39}
40
41impl AgentRuntimeContext {
42    /// Build runtime metadata with required tenant, agent, and source fields.
43    pub fn new(
44        tenant_id: impl Into<String>,
45        agent_name: impl Into<String>,
46        request_source: impl Into<String>,
47    ) -> Self {
48        Self {
49            tenant_id: tenant_id.into(),
50            agent_name: agent_name.into(),
51            request_source: request_source.into(),
52            ..Self::default()
53        }
54    }
55}
56
57/// Trait for injecting external context into agent calls.
58///
59/// Called before every LLM invocation with the agent name and tenant ID.
60/// Returns `None` if no external context is available.
61#[async_trait]
62pub trait ContextProvider: Send + Sync + 'static {
63    /// Get context for a specific agent and tenant.
64    async fn get_context(&self, agent_name: &str, tenant_id: &str) -> Option<String> {
65        let runtime = AgentRuntimeContext::new(tenant_id, agent_name, "legacy_context_provider");
66        self.get_context_for_run(&runtime).await
67    }
68
69    /// Get context using the full runtime metadata when available.
70    async fn get_context_for_run(&self, runtime: &AgentRuntimeContext) -> Option<String>;
71}
72
73/// Default: no external context (pure OSS mode).
74///
75/// Agents run with only their configured system prompt.
76pub struct NoOpContextProvider;
77
78#[async_trait]
79impl ContextProvider for NoOpContextProvider {
80    async fn get_context_for_run(&self, _runtime: &AgentRuntimeContext) -> Option<String> {
81        None
82    }
83}
84
85/// Cordis-native handle for the process-wide ContextProvider.
86#[derive(Clone)]
87pub struct ContextProviderHandle(pub std::sync::Arc<dyn ContextProvider>);
88
89impl cordis::Service for ContextProviderHandle {
90    fn name(&self) -> &'static str {
91        "context_provider"
92    }
93    fn init(
94        &self,
95        _ctx: &std::sync::Arc<cordis::Context>,
96    ) -> cordis::ServiceInitFuture<'_> {
97        Box::pin(async { Ok(None) })
98    }
99    fn check(&self) -> bool {
100        true
101    }
102}
103
104impl ContextProviderHandle {
105    pub fn new(inner: std::sync::Arc<dyn ContextProvider>) -> Self {
106        Self(inner)
107    }
108    pub fn inner(&self) -> &std::sync::Arc<dyn ContextProvider> {
109        &self.0
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use std::sync::Mutex;
117
118    struct RecordingContextProvider {
119        last_runtime: Mutex<Option<AgentRuntimeContext>>,
120        response: Option<String>,
121    }
122
123    #[async_trait]
124    impl ContextProvider for RecordingContextProvider {
125        async fn get_context_for_run(&self, runtime: &AgentRuntimeContext) -> Option<String> {
126            *self.last_runtime.lock().unwrap() = Some(runtime.clone());
127            self.response.clone()
128        }
129    }
130
131    struct SelectiveContextProvider;
132
133    #[async_trait]
134    impl ContextProvider for SelectiveContextProvider {
135        async fn get_context_for_run(&self, runtime: &AgentRuntimeContext) -> Option<String> {
136            match (runtime.tenant_id.as_str(), runtime.agent_name.as_str()) {
137                ("tenant-a", "agent-x") => Some(format!(
138                    "workspace={:?}",
139                    runtime.workspace_id.as_deref().unwrap_or("none")
140                )),
141                ("tenant-b", _) => Some("tenant-b default".into()),
142                _ => None,
143            }
144        }
145    }
146
147    #[tokio::test]
148    async fn test_noop_returns_none() {
149        let provider = NoOpContextProvider;
150        let result = provider.get_context("any_agent", "any_tenant").await;
151        assert!(result.is_none(), "NoOp should always return None");
152    }
153
154    #[test]
155    fn runtime_context_default_has_empty_strings_and_none_optionals() {
156        let runtime = AgentRuntimeContext::default();
157        assert_eq!(runtime.tenant_id, "");
158        assert_eq!(runtime.agent_name, "");
159        assert_eq!(runtime.request_source, "");
160        assert_eq!(runtime.workspace_id, None);
161        assert_eq!(runtime.user_id, None);
162        assert_eq!(runtime.session_id, None);
163    }
164
165    #[test]
166    fn runtime_context_new_sets_required_fields() {
167        let runtime = AgentRuntimeContext::new("tenant-1", "agent-1", "api_v1_chat");
168        assert_eq!(runtime.tenant_id, "tenant-1");
169        assert_eq!(runtime.agent_name, "agent-1");
170        assert_eq!(runtime.request_source, "api_v1_chat");
171        assert_eq!(runtime.workspace_id, None);
172        assert_eq!(runtime.user_id, None);
173        assert_eq!(runtime.session_id, None);
174    }
175
176    #[tokio::test]
177    async fn test_noop_get_context_for_run_returns_none() {
178        let provider = NoOpContextProvider;
179        let runtime = AgentRuntimeContext::new("tenant-1", "agent-1", "unit_test");
180        assert!(provider.get_context_for_run(&runtime).await.is_none());
181    }
182
183    #[test]
184    fn runtime_context_optional_fields_round_trip() {
185        let runtime = AgentRuntimeContext {
186            tenant_id: "tenant-1".into(),
187            agent_name: "agent-1".into(),
188            workspace_id: Some("ws-9".into()),
189            user_id: Some("user-42".into()),
190            session_id: Some("sess-7".into()),
191            request_source: "orchestrator".into(),
192        };
193        assert_eq!(runtime.workspace_id.as_deref(), Some("ws-9"));
194        assert_eq!(runtime.user_id.as_deref(), Some("user-42"));
195        assert_eq!(runtime.session_id.as_deref(), Some("sess-7"));
196    }
197
198    #[tokio::test]
199    async fn get_context_builds_legacy_runtime_for_resolution() {
200        let provider = RecordingContextProvider {
201            last_runtime: Mutex::new(None),
202            response: Some("injected".into()),
203        };
204
205        let resolved = provider.get_context("my-agent", "my-tenant").await;
206        assert_eq!(resolved.as_deref(), Some("injected"));
207
208        let runtime = provider.last_runtime.lock().unwrap().take().unwrap();
209        assert_eq!(runtime.tenant_id, "my-tenant");
210        assert_eq!(runtime.agent_name, "my-agent");
211        assert_eq!(runtime.request_source, "legacy_context_provider");
212    }
213
214    #[tokio::test]
215    async fn context_resolution_uses_runtime_metadata() {
216        let provider = SelectiveContextProvider;
217
218        let mut runtime =
219            AgentRuntimeContext::new("tenant-a", "agent-x", "managed_platform");
220        runtime.workspace_id = Some("ws-1".into());
221
222        let resolved = provider.get_context_for_run(&runtime).await;
223        assert_eq!(resolved.as_deref(), Some("workspace=\"ws-1\""));
224
225        let unknown =
226            AgentRuntimeContext::new("tenant-z", "agent-x", "managed_platform");
227        assert!(provider.get_context_for_run(&unknown).await.is_none());
228
229        let tenant_default =
230            AgentRuntimeContext::new("tenant-b", "any-agent", "managed_platform");
231        assert_eq!(
232            provider.get_context_for_run(&tenant_default).await.as_deref(),
233            Some("tenant-b default")
234        );
235    }
236
237    #[tokio::test]
238    async fn test_noop_is_send_sync() {
239        // Verify the trait object can be shared across threads
240        let provider: Box<dyn ContextProvider> = Box::new(NoOpContextProvider);
241        let arc = std::sync::Arc::new(provider);
242        let _clone = arc.clone();
243    }
244
245    #[test]
246    fn context_provider_handle_readable_via_cordis() {
247        let ctx = cordis::Context::new_root();
248        ctx.provide(ContextProviderHandle::new(std::sync::Arc::new(
249            NoOpContextProvider,
250        )));
251        assert!(ctx.get::<ContextProviderHandle>().is_some());
252    }
253}
ares_agent/context_provider.rs

ares_agent/
context_provider.rs
