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

a3s_code_core/security/
mod.rs
