Skip to main content

ai_agent/utils/
session_environment.rs

1//! Session environment utilities.
2
3use crate::constants::env::ai;
4use std::collections::HashMap;
5
6/// Get environment variables for a session
7pub fn get_session_environment() -> HashMap<String, String> {
8    let mut env = HashMap::new();
9
10    // Add session-specific environment variables
11    if let Ok(session_id) = std::env::var(ai::CODE_SESSION_ID) {
12        env.insert(ai::CODE_SESSION_ID.to_string(), session_id);
13    }
14
15    if let Ok(model) = std::env::var(ai::MODEL) {
16        env.insert(ai::MODEL.to_string(), model);
17    }
18
19    // Add all environment variables starting with AI_CODE_
20    for (key, value) in std::env::vars() {
21        if key.starts_with("AI_CODE_") {
22            env.insert(key, value);
23        }
24    }
25
26    env
27}
28
29/// Set session environment variables
30pub fn set_session_environment(env: &HashMap<String, String>) {
31    for (key, value) in env {
32        unsafe { std::env::set_var(key, value) };
33    }
34}
35
36/// Clear session environment variables
37pub fn clear_session_environment() {
38    let keys: Vec<String> = std::env::vars()
39        .filter(|(key, _)| key.starts_with("AI_CODE_"))
40        .map(|(key, _)| key)
41        .collect();
42
43    for key in keys {
44        unsafe { std::env::remove_var(&key) };
45    }
46}