agent_framework_purview/
middleware.rs

1//! [`PurviewAgentMiddleware`] and [`PurviewChatMiddleware`]: enforce Purview
2//! policy on both the outgoing prompt and the model/agent response. Mirrors
3//! Python's `PurviewPolicyMiddleware` (agent-level) and
4//! `PurviewChatPolicyMiddleware` (chat-client-level) — see the crate docs
5//! for which hook points these attach to and why both directions evaluate
6//! with [`Activity::UploadText`](crate::models::Activity::UploadText).
7
8use async_trait::async_trait;
9
10use agent_framework_core::error::Result;
11use agent_framework_core::middleware::{AgentContext, ChatContext, Middleware, Next};
12use agent_framework_core::types::{AgentResponse, ChatResponse, Message, Role};
13
14use crate::auth::TokenProvider;
15use crate::client::PurviewClient;
16use crate::processor::ContentProcessor;
17use crate::settings::PurviewSettings;
18
19/// Shared evaluation core for [`PurviewAgentMiddleware`] and
20/// [`PurviewChatMiddleware`]: identical policy logic, only the surrounding
21/// context type differs — matching how the two Python middleware classes
22/// each build their own `PurviewClient`/`ScopedContentProcessor` but run the
23/// exact same `process_messages` calls.
24struct PurviewPolicyCore {
25    processor: ContentProcessor,
26    settings: PurviewSettings,
27}
28
29impl PurviewPolicyCore {
30    fn new(token_provider: impl TokenProvider + 'static, settings: PurviewSettings) -> Self {
31        let client = PurviewClient::new(token_provider, &settings);
32        Self {
33            processor: ContentProcessor::new(client),
34            settings,
35        }
36    }
37
38    /// Evaluate `messages`, folding a suppressed error into an "allow, carry
39    /// the previously-known user id forward" outcome.
40    ///
41    /// Mirrors the two-tier `except PurviewPaymentRequiredError: ... except
42    /// Exception: ...` structure in both Python middleware classes: a 402
43    /// is checked against [`PurviewSettings::ignore_payment_required`]
44    /// *independently of* [`PurviewSettings::ignore_exceptions`], which
45    /// governs every other error. An unsuppressed error propagates (`?` at
46    /// the call site), same as Python re-raising.
47    async fn check(
48        &self,
49        messages: &[Message],
50        provided_user_id: Option<&str>,
51        phase: &'static str,
52    ) -> Result<(bool, Option<String>)> {
53        match self
54            .processor
55            .evaluate(messages, &self.settings, provided_user_id)
56            .await
57        {
58            Ok(outcome) => Ok(outcome),
59            Err(e) => {
60                let suppress = if e.status() == Some(402) {
61                    self.settings.ignore_payment_required
62                } else {
63                    self.settings.ignore_exceptions
64                };
65                if suppress {
66                    tracing::error!(error = %e, phase, "Purview policy check failed; continuing (ignored per settings)");
67                    Ok((false, provided_user_id.map(str::to_string)))
68                } else {
69                    Err(e)
70                }
71            }
72        }
73    }
74
75    fn blocked_prompt_message(&self) -> Message {
76        Message::new(Role::system(), self.settings.blocked_prompt_message.clone())
77    }
78
79    fn blocked_response_message(&self) -> Message {
80        Message::new(
81            Role::system(),
82            self.settings.blocked_response_message.clone(),
83        )
84    }
85}
86
87/// SupportsAgentRun middleware enforcing Purview policy on both the outgoing prompt and
88/// the agent's response. Mirrors Python's `PurviewPolicyMiddleware`.
89///
90/// - **Prompt (pre) check**: evaluates `ctx.messages`. If blocked,
91///   short-circuits with a single `system`-role message
92///   ([`PurviewSettings::blocked_prompt_message`]) and `ctx.terminate =
93///   true` — the wrapped agent/chat-client is never called, matching this
94///   crate's `ShortCircuitChat`-style test pattern in
95///   `agent-framework-core`.
96/// - **Response (post) check**: only runs when `!ctx.is_streaming` and a
97///   result was produced (mirrors Python: "Streaming responses are not
98///   supported for post-checks"). If blocked, `ctx.result` is *replaced*
99///   with a single `system`-role message
100///   ([`PurviewSettings::blocked_response_message`]) — unlike the prompt
101///   check, `ctx.terminate` is not set here (there's nothing left to
102///   terminate; the underlying call already happened), matching Python.
103///
104/// ```no_run
105/// use agent_framework_core::prelude::*;
106/// use agent_framework_purview::{PurviewAgentMiddleware, PurviewSettings, PurviewAppLocation, PurviewLocationType, StaticTokenProvider};
107/// use std::sync::Arc;
108///
109/// # fn demo(client: impl ChatClient + 'static) {
110/// let settings = PurviewSettings::new("My App")
111///     .with_tenant_id("00000000-0000-0000-0000-000000000000")
112///     .with_purview_app_location(PurviewAppLocation::new(
113///         PurviewLocationType::Application,
114///         "00000000-0000-0000-0000-000000000001",
115///     ));
116/// let middleware = PurviewAgentMiddleware::new(StaticTokenProvider::new("<graph-bearer-token>"), settings);
117///
118/// let agent = Agent::builder(client)
119///     .instructions("You are a helpful assistant.")
120///     .middleware(Arc::new(middleware))
121///     .build();
122/// # let _ = agent;
123/// # }
124/// ```
125pub struct PurviewAgentMiddleware(PurviewPolicyCore);
126
127impl PurviewAgentMiddleware {
128    pub fn new(token_provider: impl TokenProvider + 'static, settings: PurviewSettings) -> Self {
129        Self(PurviewPolicyCore::new(token_provider, settings))
130    }
131}
132
133#[async_trait]
134impl Middleware<AgentContext> for PurviewAgentMiddleware {
135    async fn process(
136        &self,
137        mut ctx: AgentContext,
138        next: Next<AgentContext>,
139    ) -> Result<AgentContext> {
140        let (should_block, resolved_user_id) = self.0.check(&ctx.messages, None, "prompt").await?;
141        if should_block {
142            ctx.result = Some(AgentResponse {
143                messages: vec![self.0.blocked_prompt_message()],
144                ..Default::default()
145            });
146            ctx.terminate = true;
147            return Ok(ctx);
148        }
149
150        let mut ctx = next.run(ctx).await?;
151
152        if !ctx.is_streaming {
153            let post_messages = ctx.result.as_ref().map(|r| r.messages.clone());
154            if let Some(messages) = post_messages {
155                let (should_block, _) = self
156                    .0
157                    .check(&messages, resolved_user_id.as_deref(), "response")
158                    .await?;
159                if should_block {
160                    ctx.result = Some(AgentResponse {
161                        messages: vec![self.0.blocked_response_message()],
162                        ..Default::default()
163                    });
164                }
165            }
166        }
167        Ok(ctx)
168    }
169}
170
171/// Chat-client middleware variant of [`PurviewAgentMiddleware`], for
172/// attaching Purview enforcement directly to a [`ChatClient`](agent_framework_core::client::ChatClient)
173/// rather than an agent's middleware pipeline. Mirrors Python's
174/// `PurviewChatPolicyMiddleware`; the policy logic is identical (see
175/// the internal `PurviewPolicyCore`) — only the hook point (and result/message types)
176/// differ.
177pub struct PurviewChatMiddleware(PurviewPolicyCore);
178
179impl PurviewChatMiddleware {
180    pub fn new(token_provider: impl TokenProvider + 'static, settings: PurviewSettings) -> Self {
181        Self(PurviewPolicyCore::new(token_provider, settings))
182    }
183}
184
185#[async_trait]
186impl Middleware<ChatContext> for PurviewChatMiddleware {
187    async fn process(&self, mut ctx: ChatContext, next: Next<ChatContext>) -> Result<ChatContext> {
188        let (should_block, resolved_user_id) = self.0.check(&ctx.messages, None, "prompt").await?;
189        if should_block {
190            ctx.result = Some(ChatResponse {
191                messages: vec![self.0.blocked_prompt_message()],
192                ..Default::default()
193            });
194            ctx.terminate = true;
195            return Ok(ctx);
196        }
197
198        let mut ctx = next.run(ctx).await?;
199
200        if !ctx.is_streaming {
201            let post_messages = ctx.result.as_ref().map(|r| r.messages.clone());
202            if let Some(messages) = post_messages {
203                let (should_block, _) = self
204                    .0
205                    .check(&messages, resolved_user_id.as_deref(), "response")
206                    .await?;
207                if should_block {
208                    ctx.result = Some(ChatResponse {
209                        messages: vec![self.0.blocked_response_message()],
210                        ..Default::default()
211                    });
212                }
213            }
214        }
215        Ok(ctx)
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222    use crate::auth::StaticTokenProvider;
223    use crate::settings::{PurviewAppLocation, PurviewLocationType};
224    use agent_framework_core::middleware::{MiddlewarePipeline, Terminal};
225    use agent_framework_core::tools::BoxFuture;
226    use std::sync::atomic::{AtomicBool, Ordering};
227    use std::sync::Arc;
228
229    fn valid_settings() -> PurviewSettings {
230        PurviewSettings::new("Test App")
231            .with_tenant_id("12345678-1234-1234-1234-123456789012")
232            .with_purview_app_location(PurviewAppLocation::new(
233                PurviewLocationType::Application,
234                "app-1",
235            ))
236    }
237
238    // These tests never reach the network: a missing tenant id / app
239    // location fails deterministically before any HTTP call is attempted,
240    // and "no resolvable user id" short-circuits to an allow verdict before
241    // one too (see `ContentProcessor::evaluate`). That's enough to exercise
242    // `PurviewAgentMiddleware`/`PurviewChatMiddleware::process`'s
243    // short-circuit and error-propagation logic hermetically, using a "mock
244    // next" continuation built from `agent-framework-core`'s own
245    // `MiddlewarePipeline`/`Terminal` — the same machinery `Agent` uses
246    // internally — per this work package's "core test patterns" note. The
247    // happy-path HTTP call itself (an actual `processContent` round trip) is
248    // covered by `tests/loopback.rs`.
249
250    fn agent_middleware(settings: PurviewSettings) -> PurviewAgentMiddleware {
251        PurviewAgentMiddleware::new(StaticTokenProvider::new("token"), settings)
252    }
253
254    fn chat_middleware(settings: PurviewSettings) -> PurviewChatMiddleware {
255        PurviewChatMiddleware::new(StaticTokenProvider::new("token"), settings)
256    }
257
258    fn agent_terminal(called: Arc<AtomicBool>, text: &'static str) -> Terminal<AgentContext> {
259        Box::new(move |mut ctx: AgentContext| {
260            called.store(true, Ordering::SeqCst);
261            Box::pin(async move {
262                ctx.result = Some(AgentResponse {
263                    messages: vec![Message::assistant(text)],
264                    ..Default::default()
265                });
266                Ok(ctx)
267            }) as BoxFuture<Result<AgentContext>>
268        })
269    }
270
271    fn chat_terminal(called: Arc<AtomicBool>, text: &'static str) -> Terminal<ChatContext> {
272        Box::new(move |mut ctx: ChatContext| {
273            called.store(true, Ordering::SeqCst);
274            Box::pin(async move {
275                ctx.result = Some(ChatResponse::from_text(text));
276                Ok(ctx)
277            }) as BoxFuture<Result<ChatContext>>
278        })
279    }
280
281    /// `Result::unwrap_err` requires the `Ok` type to implement `Debug`,
282    /// which `AgentContext`/`ChatContext` deliberately don't (they carry
283    /// non-`Debug` middleware/result trait objects). Same shape, without
284    /// that bound.
285    fn expect_err<T>(result: Result<T>) -> agent_framework_core::error::Error {
286        match result {
287            Ok(_) => panic!("expected Err, got Ok"),
288            Err(e) => e,
289        }
290    }
291
292    // -- config-error propagation (no network: fails before any HTTP call) -
293
294    #[tokio::test]
295    async fn agent_middleware_propagates_config_error_without_ignore_exceptions() {
296        let settings = PurviewSettings::new("Test App"); // no tenant_id/app_location
297        let middleware = agent_middleware(settings);
298        let pipeline = MiddlewarePipeline::new(vec![Arc::new(middleware)]);
299        let called = Arc::new(AtomicBool::new(false));
300        let ctx = AgentContext::new(vec![Message::user("hi")], false);
301
302        let err = expect_err(
303            pipeline
304                .execute(ctx, agent_terminal(called.clone(), "should not be reached"))
305                .await,
306        );
307        assert!(err.to_string().contains("tenant_id"));
308        assert!(
309            !called.load(Ordering::SeqCst),
310            "next must not run on a pre-check error"
311        );
312    }
313
314    #[tokio::test]
315    async fn agent_middleware_ignore_exceptions_suppresses_config_error_and_allows() {
316        let settings = PurviewSettings::new("Test App").with_ignore_exceptions(true);
317        let middleware = agent_middleware(settings);
318        let pipeline = MiddlewarePipeline::new(vec![Arc::new(middleware)]);
319        let called = Arc::new(AtomicBool::new(false));
320        let ctx = AgentContext::new(vec![Message::user("hi")], false);
321
322        let result_ctx = pipeline
323            .execute(ctx, agent_terminal(called.clone(), "real response"))
324            .await
325            .unwrap();
326        assert!(
327            called.load(Ordering::SeqCst),
328            "next must run once the error is suppressed"
329        );
330        assert!(!result_ctx.terminate);
331        assert_eq!(result_ctx.result.unwrap().text(), "real response");
332    }
333
334    // -- resolve_user_id short-circuits allow (no network) ------------------
335
336    #[tokio::test]
337    async fn agent_middleware_allows_when_no_resolvable_user_id() {
338        // Valid settings, but the message carries no GUID-shaped user_id or
339        // author_name -- resolve_user_id returns None, which the processor
340        // treats as "cannot evaluate; allow" *without* an HTTP call. If it
341        // did attempt one, this test would hang/fail trying to reach
342        // graph.microsoft.com.
343        let middleware = agent_middleware(valid_settings());
344        let pipeline = MiddlewarePipeline::new(vec![Arc::new(middleware)]);
345        let called = Arc::new(AtomicBool::new(false));
346        let ctx = AgentContext::new(
347            vec![Message::user("hello, nothing identifying here")],
348            false,
349        );
350
351        let result_ctx = pipeline
352            .execute(ctx, agent_terminal(called.clone(), "real response"))
353            .await
354            .unwrap();
355        assert!(called.load(Ordering::SeqCst));
356        assert!(!result_ctx.terminate);
357        assert_eq!(result_ctx.result.unwrap().text(), "real response");
358    }
359
360    #[tokio::test]
361    async fn agent_middleware_skips_post_check_when_streaming() {
362        // is_streaming = true -> the response-phase check must never run,
363        // so even a config-broken response phase can't surface an error.
364        let middleware = agent_middleware(valid_settings());
365        let pipeline = MiddlewarePipeline::new(vec![Arc::new(middleware)]);
366        let called = Arc::new(AtomicBool::new(false));
367        let ctx = AgentContext::new(vec![Message::user("hello, nothing identifying here")], true);
368
369        let result_ctx = pipeline
370            .execute(ctx, agent_terminal(called.clone(), "streamed response"))
371            .await
372            .unwrap();
373        assert!(called.load(Ordering::SeqCst));
374        assert_eq!(result_ctx.result.unwrap().text(), "streamed response");
375    }
376
377    #[tokio::test]
378    async fn chat_middleware_allows_when_no_resolvable_user_id() {
379        let middleware = chat_middleware(valid_settings());
380        let pipeline = MiddlewarePipeline::new(vec![Arc::new(middleware)]);
381        let called = Arc::new(AtomicBool::new(false));
382        let ctx = ChatContext::new(
383            vec![Message::user("hello, nothing identifying here")],
384            agent_framework_core::types::ChatOptions::new(),
385            false,
386        );
387
388        let result_ctx = pipeline
389            .execute(ctx, chat_terminal(called.clone(), "real response"))
390            .await
391            .unwrap();
392        assert!(called.load(Ordering::SeqCst));
393        assert!(!result_ctx.terminate);
394        assert_eq!(result_ctx.result.unwrap().text(), "real response");
395    }
396
397    #[tokio::test]
398    async fn chat_middleware_propagates_config_error_without_ignore_exceptions() {
399        let settings = PurviewSettings::new("Test App");
400        let middleware = chat_middleware(settings);
401        let pipeline = MiddlewarePipeline::new(vec![Arc::new(middleware)]);
402        let called = Arc::new(AtomicBool::new(false));
403        let ctx = ChatContext::new(
404            vec![Message::user("hi")],
405            agent_framework_core::types::ChatOptions::new(),
406            false,
407        );
408
409        let err = expect_err(
410            pipeline
411                .execute(ctx, chat_terminal(called.clone(), "should not be reached"))
412                .await,
413        );
414        assert!(err.to_string().contains("tenant_id"));
415        assert!(!called.load(Ordering::SeqCst));
416    }
417
418    // -- blocked-message text/role ------------------------------------------
419
420    #[test]
421    fn blocked_messages_use_system_role_and_configured_text() {
422        let core = PurviewPolicyCore::new(StaticTokenProvider::new("t"), valid_settings());
423        let prompt_msg = core.blocked_prompt_message();
424        let response_msg = core.blocked_response_message();
425        assert_eq!(prompt_msg.role, Role::system());
426        assert_eq!(prompt_msg.text(), "Prompt blocked by policy");
427        assert_eq!(response_msg.role, Role::system());
428        assert_eq!(response_msg.text(), "Response blocked by policy");
429    }
430
431    #[test]
432    fn blocked_messages_honor_custom_text() {
433        let settings = valid_settings()
434            .with_blocked_prompt_message("custom prompt block")
435            .with_blocked_response_message("custom response block");
436        let core = PurviewPolicyCore::new(StaticTokenProvider::new("t"), settings);
437        assert_eq!(core.blocked_prompt_message().text(), "custom prompt block");
438        assert_eq!(
439            core.blocked_response_message().text(),
440            "custom response block"
441        );
442    }
443}
agent_framework_purview/middleware.rs

agent_framework_purview/
middleware.rs
