a3s_code_core/security/
mod.rs

1//! Security Module
2//!
3//! Provides a trait-based security interface for A3S Code sessions.
4//! External consumers implement `SecurityProvider` to plug in their own
5//! security logic (sanitization, taint tracking, injection detection, etc.).
6
7pub mod config;
8pub mod default;
9mod event_sanitizer;
10
11pub use config::{RedactionStrategy, SecurityConfig, SensitivityLevel};
12pub use default::{DefaultSecurityConfig, DefaultSecurityProvider, SensitivePattern};
13pub(crate) use event_sanitizer::AgentEventStreamSanitizer;
14
15use crate::hooks::HookEngine;
16
17/// Sanitize every data-bearing field of an agent event while preserving the
18/// identifiers and discriminants used to correlate the event stream.
19///
20/// Runtime boundaries call this immediately before events are observed,
21/// persisted, or exposed to an SDK. Implementations of [`SecurityProvider`]
22/// should therefore make [`SecurityProvider::sanitize_output`] idempotent.
23pub fn sanitize_agent_event(
24    provider: &dyn SecurityProvider,
25    event: &crate::agent::AgentEvent,
26) -> crate::agent::AgentEvent {
27    use crate::agent::AgentEvent;
28
29    fn text(provider: &dyn SecurityProvider, value: &str) -> String {
30        provider.sanitize_output(value)
31    }
32
33    fn optional_text(provider: &dyn SecurityProvider, value: &Option<String>) -> Option<String> {
34        value.as_deref().map(|value| text(provider, value))
35    }
36
37    fn strings(provider: &dyn SecurityProvider, values: &[String]) -> Vec<String> {
38        values.iter().map(|value| text(provider, value)).collect()
39    }
40
41    fn json(provider: &dyn SecurityProvider, value: &serde_json::Value) -> serde_json::Value {
42        match value {
43            serde_json::Value::String(value) => serde_json::Value::String(text(provider, value)),
44            serde_json::Value::Array(values) => {
45                serde_json::Value::Array(values.iter().map(|value| json(provider, value)).collect())
46            }
47            serde_json::Value::Object(values) => serde_json::Value::Object(
48                values
49                    .iter()
50                    .map(|(key, value)| (key.clone(), json(provider, value)))
51                    .collect(),
52            ),
53            value => value.clone(),
54        }
55    }
56
57    fn optional_json(
58        provider: &dyn SecurityProvider,
59        value: &Option<serde_json::Value>,
60    ) -> Option<serde_json::Value> {
61        value.as_ref().map(|value| json(provider, value))
62    }
63
64    fn tool_error_kind(
65        provider: &dyn SecurityProvider,
66        value: &Option<crate::tools::ToolErrorKind>,
67    ) -> Option<crate::tools::ToolErrorKind> {
68        use crate::tools::ToolErrorKind;
69
70        value.as_ref().map(|kind| match kind {
71            ToolErrorKind::VersionConflict {
72                path,
73                expected,
74                actual,
75            } => ToolErrorKind::VersionConflict {
76                path: text(provider, path),
77                expected: text(provider, expected),
78                actual: optional_text(provider, actual),
79            },
80            ToolErrorKind::RemoteGitConflict { code, message } => {
81                ToolErrorKind::RemoteGitConflict {
82                    code: text(provider, code),
83                    message: text(provider, message),
84                }
85            }
86            ToolErrorKind::NotFound { path } => ToolErrorKind::NotFound {
87                path: text(provider, path),
88            },
89            ToolErrorKind::InvalidArgument { message } => ToolErrorKind::InvalidArgument {
90                message: text(provider, message),
91            },
92            ToolErrorKind::Unsupported { message } => ToolErrorKind::Unsupported {
93                message: text(provider, message),
94            },
95            ToolErrorKind::Timeout { op, duration_ms } => ToolErrorKind::Timeout {
96                op: text(provider, op),
97                duration_ms: *duration_ms,
98            },
99        })
100    }
101
102    fn task(
103        provider: &dyn SecurityProvider,
104        task: &crate::planning::Task,
105    ) -> crate::planning::Task {
106        let mut task = task.clone();
107        task.content = text(provider, &task.content);
108        task.success_criteria = optional_text(provider, &task.success_criteria);
109        task
110    }
111
112    fn plan(
113        provider: &dyn SecurityProvider,
114        plan: &crate::planning::ExecutionPlan,
115    ) -> crate::planning::ExecutionPlan {
116        let mut plan = plan.clone();
117        plan.goal = text(provider, &plan.goal);
118        plan.steps = plan.steps.iter().map(|item| task(provider, item)).collect();
119        plan
120    }
121
122    fn goal(
123        provider: &dyn SecurityProvider,
124        goal: &crate::planning::AgentGoal,
125    ) -> crate::planning::AgentGoal {
126        let mut goal = goal.clone();
127        goal.description = text(provider, &goal.description);
128        goal.success_criteria = strings(provider, &goal.success_criteria);
129        goal
130    }
131
132    fn verification_summary(
133        provider: &dyn SecurityProvider,
134        summary: &crate::verification::VerificationSummary,
135    ) -> crate::verification::VerificationSummary {
136        let mut summary = summary.clone();
137        summary.pending_subjects = strings(provider, &summary.pending_subjects);
138        summary.failed_subjects = strings(provider, &summary.failed_subjects);
139        summary
140    }
141
142    fn response_meta(
143        provider: &dyn SecurityProvider,
144        meta: &Option<crate::llm::LlmResponseMeta>,
145    ) -> Option<crate::llm::LlmResponseMeta> {
146        meta.as_ref().map(|meta| {
147            let mut meta = meta.clone();
148            meta.request_url = optional_text(provider, &meta.request_url);
149            meta
150        })
151    }
152
153    match event {
154        AgentEvent::Start { prompt } => AgentEvent::Start {
155            prompt: text(provider, prompt),
156        },
157        AgentEvent::AgentModeChanged {
158            mode,
159            agent,
160            description,
161        } => AgentEvent::AgentModeChanged {
162            mode: mode.clone(),
163            agent: agent.clone(),
164            description: text(provider, description),
165        },
166        AgentEvent::TextDelta { text: value } => AgentEvent::TextDelta {
167            text: text(provider, value),
168        },
169        AgentEvent::ReasoningDelta { text: value } => AgentEvent::ReasoningDelta {
170            text: text(provider, value),
171        },
172        AgentEvent::ToolInputDelta { id, delta } => AgentEvent::ToolInputDelta {
173            id: id.clone(),
174            delta: text(provider, delta),
175        },
176        AgentEvent::ToolExecutionStart { id, name, args } => AgentEvent::ToolExecutionStart {
177            id: id.clone(),
178            name: name.clone(),
179            args: json(provider, args),
180        },
181        AgentEvent::ToolEnd {
182            id,
183            name,
184            args,
185            output,
186            exit_code,
187            metadata,
188            error_kind,
189        } => AgentEvent::ToolEnd {
190            id: id.clone(),
191            name: name.clone(),
192            args: optional_json(provider, args),
193            output: text(provider, output),
194            exit_code: *exit_code,
195            metadata: optional_json(provider, metadata),
196            error_kind: tool_error_kind(provider, error_kind),
197        },
198        AgentEvent::ToolOutputDelta { id, name, delta } => AgentEvent::ToolOutputDelta {
199            id: id.clone(),
200            name: name.clone(),
201            delta: text(provider, delta),
202        },
203        AgentEvent::End {
204            text: value,
205            usage,
206            verification_summary: summary,
207            meta,
208        } => AgentEvent::End {
209            text: text(provider, value),
210            usage: usage.clone(),
211            verification_summary: Box::new(verification_summary(provider, summary)),
212            meta: response_meta(provider, meta),
213        },
214        AgentEvent::Error { message } => AgentEvent::Error {
215            message: text(provider, message),
216        },
217        AgentEvent::ConfirmationRequired {
218            tool_id,
219            tool_name,
220            args,
221            timeout_ms,
222        } => AgentEvent::ConfirmationRequired {
223            tool_id: tool_id.clone(),
224            tool_name: tool_name.clone(),
225            args: json(provider, args),
226            timeout_ms: *timeout_ms,
227        },
228        AgentEvent::ConfirmationReceived {
229            tool_id,
230            approved,
231            reason,
232        } => AgentEvent::ConfirmationReceived {
233            tool_id: tool_id.clone(),
234            approved: *approved,
235            reason: optional_text(provider, reason),
236        },
237        AgentEvent::ConfirmationTimeout {
238            tool_id,
239            action_taken,
240        } => AgentEvent::ConfirmationTimeout {
241            tool_id: tool_id.clone(),
242            action_taken: action_taken.clone(),
243        },
244        AgentEvent::ExternalTaskPending {
245            task_id,
246            session_id,
247            lane,
248            command_type,
249            payload,
250            timeout_ms,
251        } => AgentEvent::ExternalTaskPending {
252            task_id: task_id.clone(),
253            session_id: session_id.clone(),
254            lane: *lane,
255            command_type: command_type.clone(),
256            payload: json(provider, payload),
257            timeout_ms: *timeout_ms,
258        },
259        AgentEvent::PermissionDenied {
260            tool_id,
261            tool_name,
262            args,
263            reason,
264        } => AgentEvent::PermissionDenied {
265            tool_id: tool_id.clone(),
266            tool_name: tool_name.clone(),
267            args: json(provider, args),
268            reason: text(provider, reason),
269        },
270        AgentEvent::CommandDeadLettered {
271            command_id,
272            command_type,
273            lane,
274            error,
275            attempts,
276        } => AgentEvent::CommandDeadLettered {
277            command_id: command_id.clone(),
278            command_type: command_type.clone(),
279            lane: lane.clone(),
280            error: text(provider, error),
281            attempts: *attempts,
282        },
283        AgentEvent::QueueAlert {
284            level,
285            alert_type,
286            message,
287        } => AgentEvent::QueueAlert {
288            level: level.clone(),
289            alert_type: alert_type.clone(),
290            message: text(provider, message),
291        },
292        AgentEvent::TaskUpdated { session_id, tasks } => AgentEvent::TaskUpdated {
293            session_id: session_id.clone(),
294            tasks: tasks.iter().map(|item| task(provider, item)).collect(),
295        },
296        AgentEvent::MemoryStored {
297            memory_id,
298            memory_type,
299            importance,
300            tags,
301        } => AgentEvent::MemoryStored {
302            memory_id: memory_id.clone(),
303            memory_type: memory_type.clone(),
304            importance: *importance,
305            tags: strings(provider, tags),
306        },
307        AgentEvent::MemoryRecalled {
308            memory_id,
309            content,
310            relevance,
311        } => AgentEvent::MemoryRecalled {
312            memory_id: memory_id.clone(),
313            content: text(provider, content),
314            relevance: *relevance,
315        },
316        AgentEvent::MemoriesSearched {
317            query,
318            tags,
319            result_count,
320        } => AgentEvent::MemoriesSearched {
321            query: optional_text(provider, query),
322            tags: strings(provider, tags),
323            result_count: *result_count,
324        },
325        AgentEvent::SubagentStart {
326            task_id,
327            session_id,
328            parent_session_id,
329            agent,
330            description,
331            started_ms,
332        } => AgentEvent::SubagentStart {
333            task_id: task_id.clone(),
334            session_id: session_id.clone(),
335            parent_session_id: parent_session_id.clone(),
336            agent: agent.clone(),
337            description: text(provider, description),
338            started_ms: *started_ms,
339        },
340        AgentEvent::SubagentProgress {
341            task_id,
342            session_id,
343            status,
344            metadata,
345        } => AgentEvent::SubagentProgress {
346            task_id: task_id.clone(),
347            session_id: session_id.clone(),
348            status: text(provider, status),
349            metadata: json(provider, metadata),
350        },
351        AgentEvent::SubagentEnd {
352            task_id,
353            session_id,
354            agent,
355            output,
356            success,
357            finished_ms,
358        } => AgentEvent::SubagentEnd {
359            task_id: task_id.clone(),
360            session_id: session_id.clone(),
361            agent: agent.clone(),
362            output: text(provider, output),
363            success: *success,
364            finished_ms: *finished_ms,
365        },
366        AgentEvent::PlanningStart { prompt } => AgentEvent::PlanningStart {
367            prompt: text(provider, prompt),
368        },
369        AgentEvent::PlanningEnd {
370            plan: value,
371            estimated_steps,
372        } => AgentEvent::PlanningEnd {
373            plan: plan(provider, value),
374            estimated_steps: *estimated_steps,
375        },
376        AgentEvent::StepStart {
377            step_id,
378            description,
379            step_number,
380            total_steps,
381        } => AgentEvent::StepStart {
382            step_id: step_id.clone(),
383            description: text(provider, description),
384            step_number: *step_number,
385            total_steps: *total_steps,
386        },
387        AgentEvent::GoalExtracted { goal: value } => AgentEvent::GoalExtracted {
388            goal: goal(provider, value),
389        },
390        AgentEvent::GoalProgress {
391            goal,
392            progress,
393            completed_steps,
394            total_steps,
395        } => AgentEvent::GoalProgress {
396            goal: text(provider, goal),
397            progress: *progress,
398            completed_steps: *completed_steps,
399            total_steps: *total_steps,
400        },
401        AgentEvent::GoalAchieved {
402            goal,
403            total_steps,
404            duration_ms,
405        } => AgentEvent::GoalAchieved {
406            goal: text(provider, goal),
407            total_steps: *total_steps,
408            duration_ms: *duration_ms,
409        },
410        AgentEvent::PersistenceFailed {
411            session_id,
412            operation,
413            error,
414        } => AgentEvent::PersistenceFailed {
415            session_id: session_id.clone(),
416            operation: operation.clone(),
417            error: text(provider, error),
418        },
419        AgentEvent::BudgetThresholdHit {
420            resource,
421            kind,
422            consumed,
423            limit,
424            message,
425        } => AgentEvent::BudgetThresholdHit {
426            resource: resource.clone(),
427            kind: kind.clone(),
428            consumed: *consumed,
429            limit: *limit,
430            message: optional_text(provider, message),
431        },
432        AgentEvent::PassivationRequested {
433            reason,
434            deadline_ms,
435        } => AgentEvent::PassivationRequested {
436            reason: text(provider, reason),
437            deadline_ms: *deadline_ms,
438        },
439        _ => event.clone(),
440    }
441}
442
443/// Trait for pluggable security providers.
444///
445/// Implement this trait to provide custom security logic for sessions.
446/// The default `NoOpSecurityProvider` passes everything through unchanged.
447pub trait SecurityProvider: Send + Sync {
448    /// Classify and register sensitive data found in input text
449    fn taint_input(&self, _text: &str) {}
450
451    /// Sanitize output text by redacting sensitive data.
452    /// Returns the sanitized text.
453    fn sanitize_output(&self, text: &str) -> String {
454        text.to_string()
455    }
456
457    /// Securely wipe all session security state
458    fn wipe(&self) {}
459
460    /// Register security hooks with the given engine
461    fn register_hooks(&self, _hook_engine: &HookEngine) {}
462
463    /// Unregister all hooks from the engine
464    fn teardown(&self, _hook_engine: &HookEngine) {}
465}
466
467/// No-op security provider (default when security is disabled)
468pub struct NoOpSecurityProvider;
469
470impl SecurityProvider for NoOpSecurityProvider {}
471
472#[cfg(test)]
473mod tests {
474    use super::*;
475
476    #[test]
477    fn test_noop_provider_passthrough() {
478        let provider = NoOpSecurityProvider;
479        provider.taint_input("SSN: 123-45-6789");
480        let output = provider.sanitize_output("SSN: 123-45-6789");
481        assert_eq!(output, "SSN: 123-45-6789");
482    }
483
484    #[test]
485    fn test_noop_provider_wipe() {
486        let provider = NoOpSecurityProvider;
487        provider.wipe(); // Should not panic
488    }
489
490    #[test]
491    fn test_noop_provider_hooks() {
492        let engine = HookEngine::new();
493        let provider = NoOpSecurityProvider;
494        provider.register_hooks(&engine);
495        provider.teardown(&engine);
496        assert_eq!(engine.hook_count(), 0);
497    }
498
499    #[test]
500    fn agent_event_sanitization_redacts_streams_arguments_and_outputs() {
501        let provider = DefaultSecurityProvider::new();
502        let secret = "user@example.com";
503        let events = [
504            crate::agent::AgentEvent::TextDelta {
505                text: secret.to_string(),
506            },
507            crate::agent::AgentEvent::ReasoningDelta {
508                text: secret.to_string(),
509            },
510            crate::agent::AgentEvent::ToolInputDelta {
511                id: Some("tool-1".to_string()),
512                delta: format!(r#"{{"token":"{secret}"}}"#),
513            },
514            crate::agent::AgentEvent::ToolExecutionStart {
515                id: "tool-1".to_string(),
516                name: "bash".to_string(),
517                args: serde_json::json!({"command": format!("echo {secret}")}),
518            },
519            crate::agent::AgentEvent::ToolOutputDelta {
520                id: "tool-1".to_string(),
521                name: "bash".to_string(),
522                delta: secret.to_string(),
523            },
524            crate::agent::AgentEvent::ToolEnd {
525                id: "tool-1".to_string(),
526                name: "bash".to_string(),
527                args: Some(serde_json::json!({"command": format!("echo {secret}")})),
528                output: secret.to_string(),
529                exit_code: 0,
530                metadata: Some(serde_json::json!({"contact": secret})),
531                error_kind: None,
532            },
533        ];
534
535        for event in &events {
536            let sanitized = sanitize_agent_event(&provider, event);
537            let json = serde_json::to_string(&sanitized).unwrap();
538            assert!(!json.contains(secret), "unsanitized event: {json}");
539            assert!(json.contains("REDACTED:EMAIL"));
540        }
541
542        let sanitized = sanitize_agent_event(&provider, &events[3]);
543        assert!(matches!(
544            sanitized,
545            crate::agent::AgentEvent::ToolExecutionStart { id, name, .. }
546                if id == "tool-1" && name == "bash"
547        ));
548    }
549
550    #[test]
551    fn agent_event_sanitization_redacts_every_tool_error_kind_string() {
552        use crate::agent::AgentEvent;
553        use crate::tools::ToolErrorKind;
554
555        let provider = DefaultSecurityProvider::new();
556        let secret = "user@example.com";
557        let redacted = "[REDACTED:EMAIL]";
558        let cases = [
559            (
560                ToolErrorKind::VersionConflict {
561                    path: format!("path/{secret}"),
562                    expected: format!("expected: {secret}"),
563                    actual: Some(format!("actual: {secret}")),
564                },
565                ToolErrorKind::VersionConflict {
566                    path: format!("path/{redacted}"),
567                    expected: format!("expected: {redacted}"),
568                    actual: Some(format!("actual: {redacted}")),
569                },
570            ),
571            (
572                ToolErrorKind::RemoteGitConflict {
573                    code: format!("code: {secret}"),
574                    message: format!("message: {secret}"),
575                },
576                ToolErrorKind::RemoteGitConflict {
577                    code: format!("code: {redacted}"),
578                    message: format!("message: {redacted}"),
579                },
580            ),
581            (
582                ToolErrorKind::NotFound {
583                    path: format!("path/{secret}"),
584                },
585                ToolErrorKind::NotFound {
586                    path: format!("path/{redacted}"),
587                },
588            ),
589            (
590                ToolErrorKind::InvalidArgument {
591                    message: format!("message: {secret}"),
592                },
593                ToolErrorKind::InvalidArgument {
594                    message: format!("message: {redacted}"),
595                },
596            ),
597            (
598                ToolErrorKind::Unsupported {
599                    message: format!("message: {secret}"),
600                },
601                ToolErrorKind::Unsupported {
602                    message: format!("message: {redacted}"),
603                },
604            ),
605            (
606                ToolErrorKind::Timeout {
607                    op: format!("operation: {secret}"),
608                    duration_ms: 42,
609                },
610                ToolErrorKind::Timeout {
611                    op: format!("operation: {redacted}"),
612                    duration_ms: 42,
613                },
614            ),
615        ];
616
617        for (error_kind, expected) in cases {
618            let event = AgentEvent::ToolEnd {
619                id: "tool-1".to_string(),
620                name: "test".to_string(),
621                args: None,
622                output: String::new(),
623                exit_code: 1,
624                metadata: None,
625                error_kind: Some(error_kind),
626            };
627
628            let sanitized = sanitize_agent_event(&provider, &event);
629            let AgentEvent::ToolEnd {
630                id,
631                name,
632                exit_code,
633                error_kind,
634                ..
635            } = sanitized
636            else {
637                panic!("sanitization changed the event variant");
638            };
639
640            assert_eq!(id, "tool-1");
641            assert_eq!(name, "test");
642            assert_eq!(exit_code, 1);
643            assert_eq!(error_kind, Some(expected));
644        }
645    }
646}
a3s_code_core/security/mod.rs

a3s_code_core/security/
mod.rs
