aft/bash_rewrite/
dispatch.rs

1use std::collections::HashMap;
2use std::fs::OpenOptions;
3use std::io::Write;
4use std::sync::{Mutex, OnceLock};
5
6use serde::{Deserialize, Serialize};
7
8use crate::bash_rewrite::catalog::ControlRole;
9use crate::bash_rewrite::rules::{
10    CatAppendRule, CatRule, FindRule, GrepRule, LsRule, RgRule, SedRule,
11};
12use crate::bash_rewrite::RewriteRule;
13use crate::context::AppContext;
14use crate::protocol::Response;
15
16#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
17#[serde(tag = "kind", rename_all = "snake_case")]
18pub enum DispatchRoute {
19    Rewritten {
20        rule_id: String,
21        branch_id: String,
22        decision_class_id: String,
23    },
24    Native {
25        role: String,
26        branch_id: String,
27        reason: String,
28    },
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
32pub struct DispatchRecord {
33    pub request_id: String,
34    pub route: DispatchRoute,
35}
36
37static ROUTE_RECORDS: OnceLock<Mutex<HashMap<String, DispatchRecord>>> = OnceLock::new();
38
39fn route_records() -> &'static Mutex<HashMap<String, DispatchRecord>> {
40    ROUTE_RECORDS.get_or_init(|| Mutex::new(HashMap::new()))
41}
42
43fn store_record(record: DispatchRecord) {
44    if let Ok(mut records) = route_records().lock() {
45        records.insert(record.request_id.clone(), record.clone());
46    }
47
48    // The sidecar is opt-in for the differential child process. Normal AFT
49    // responses remain byte-for-byte free of test metadata.
50    if let Some(path) = std::env::var_os("AFT_BASH_REWRITE_ROUTE_RECORD") {
51        if let Ok(mut file) = OpenOptions::new().create(true).append(true).open(path) {
52            if let Ok(line) = serde_json::to_string(&record) {
53                let _ = writeln!(file, "{line}");
54            }
55        }
56    }
57}
58
59pub fn route_record(request_id: &str) -> Option<DispatchRecord> {
60    route_records()
61        .lock()
62        .ok()
63        .and_then(|records| records.get(request_id).cloned())
64}
65
66pub fn take_route_record(request_id: &str) -> Option<DispatchRecord> {
67    route_records()
68        .lock()
69        .ok()
70        .and_then(|mut records| records.remove(request_id))
71}
72
73pub fn record_native(request_id: &str, role: ControlRole, branch_id: &str, reason: &str) {
74    store_record(DispatchRecord {
75        request_id: request_id.to_string(),
76        route: DispatchRoute::Native {
77            role: role.id().to_string(),
78            branch_id: branch_id.to_string(),
79            reason: reason.to_string(),
80        },
81    });
82}
83
84pub fn dispatch(command: &str, session_id: Option<&str>, ctx: &AppContext) -> Option<Response> {
85    dispatch_for_request(command, "bash_rewrite", session_id, ctx)
86}
87
88pub fn dispatch_for_request(
89    command: &str,
90    request_id: &str,
91    session_id: Option<&str>,
92    ctx: &AppContext,
93) -> Option<Response> {
94    if !ctx.config().experimental_bash_rewrite {
95        record_native(
96            request_id,
97            ControlRole::Native,
98            "dispatch.native.no_rule",
99            "experimental bash rewriting is disabled",
100        );
101        return None;
102    }
103
104    let rules: [&dyn RewriteRule; 7] = [
105        &GrepRule,
106        &RgRule,
107        &FindRule,
108        &CatRule,
109        &CatAppendRule,
110        &SedRule,
111        &LsRule,
112    ];
113
114    for rule in rules {
115        let decision = rule.decide(command, request_id, session_id, ctx);
116        match decision {
117            crate::bash_rewrite::RewriteDecision::Accept(request) => {
118                store_record(DispatchRecord {
119                    request_id: request_id.to_string(),
120                    route: DispatchRoute::Rewritten {
121                        rule_id: request.rule_id.to_string(),
122                        branch_id: request.branch_id.to_string(),
123                        decision_class_id: request.decision_class_id.to_string(),
124                    },
125                });
126                // Do not turn a handler failure into a native execution. The
127                // handler has already begun and a second execution could
128                // duplicate a mutation such as cat_append.
129                return Some(rule.execute(&request, ctx));
130            }
131            crate::bash_rewrite::RewriteDecision::Decline(reason) => {
132                if reason.rule_id.is_some() {
133                    crate::slog_debug!(
134                        "bash rewrite rule {} declined before execution: {}",
135                        reason.rule_id.unwrap_or("unknown"),
136                        reason.reason
137                    );
138                }
139            }
140        }
141    }
142
143    record_native(
144        request_id,
145        ControlRole::Native,
146        "dispatch.native.no_rule",
147        "no rewrite rule accepted the command shape",
148    );
149    None
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use crate::bash_rewrite::catalog::ControlRole;
156
157    #[test]
158    fn route_records_are_request_correlated_and_replaceable() {
159        record_native(
160            "route-test",
161            ControlRole::Native,
162            "dispatch.native.no_rule",
163            "test",
164        );
165        assert_eq!(
166            route_record("route-test").expect("route record").request_id,
167            "route-test"
168        );
169        record_native(
170            "route-test",
171            ControlRole::Sandbox,
172            "dispatch.native.sandbox",
173            "test",
174        );
175        assert_eq!(
176            route_record("route-test").expect("replacement").route,
177            DispatchRoute::Native {
178                role: "sandbox".to_string(),
179                branch_id: "dispatch.native.sandbox".to_string(),
180                reason: "test".to_string(),
181            }
182        );
183        let _ = take_route_record("route-test");
184    }
185}
aft/bash_rewrite/dispatch.rs

aft/bash_rewrite/
dispatch.rs
