ai_agent/bridge/
work_secret.rs

1//! Work secret handling and session ID utilities.
2//!
3//! Translated from openclaudecode/src/bridge/workSecret.ts
4
5use crate::utils::http::get_user_agent;
6use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
7use serde::{Deserialize, Serialize};
8
9#[cfg(feature = "reqwest")]
10use reqwest;
11
12/// Work secret structure decoded from base64url-encoded JSON.
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct WorkSecret {
15    pub version: u32,
16    pub session_ingress_token: String,
17    pub api_base_url: String,
18    pub sources: Vec<WorkSource>,
19    pub auth: Vec<WorkAuth>,
20    #[serde(default)]
21    pub claude_code_args: Option<std::collections::HashMap<String, String>>,
22    #[serde(default)]
23    pub mcp_config: Option<serde_json::Value>,
24    #[serde(default)]
25    pub environment_variables: Option<std::collections::HashMap<String, String>>,
26    /// Server-driven CCR v2 selector. Set when the session was created
27    /// via the v2 compat layer.
28    #[serde(default)]
29    pub use_code_sessions: Option<bool>,
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct WorkSource {
34    #[serde(rename = "type")]
35    pub source_type: String,
36    #[serde(default)]
37    pub git_info: Option<GitInfo>,
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct GitInfo {
42    #[serde(rename = "type")]
43    pub git_type: String,
44    pub repo: String,
45    #[serde(default)]
46    pub r#ref: Option<String>,
47    #[serde(default)]
48    pub token: Option<String>,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct WorkAuth {
53    #[serde(rename = "type")]
54    pub auth_type: String,
55    pub token: String,
56}
57
58/// Decode a base64url-encoded work secret and validate its version.
59pub fn decode_work_secret(secret: &str) -> Result<WorkSecret, String> {
60    let json = URL_SAFE_NO_PAD
61        .decode(secret)
62        .map_err(|e| format!("Failed to decode base64url: {}", e))?;
63
64    let parsed: serde_json::Value =
65        serde_json::from_slice(&json).map_err(|e| format!("Failed to parse JSON: {}", e))?;
66
67    if let Some(obj) = parsed.as_object() {
68        let version = obj.get("version").and_then(|v| v.as_u64()).unwrap_or(0);
69
70        if version != 1 {
71            return Err(format!(
72                "Unsupported work secret version: {}",
73                obj.get("version")
74                    .map(|v| v.to_string())
75                    .unwrap_or_else(|| "unknown".to_string())
76            ));
77        }
78
79        // Validate required fields
80        let session_ingress_token = obj
81            .get("session_ingress_token")
82            .and_then(|v| v.as_str())
83            .filter(|s| !s.is_empty())
84            .ok_or("Invalid work secret: missing or empty session_ingress_token")?;
85
86        let api_base_url = obj
87            .get("api_base_url")
88            .and_then(|v| v.as_str())
89            .ok_or("Invalid work secret: missing api_base_url")?;
90
91        let work_secret = WorkSecret {
92            version: version as u32,
93            session_ingress_token: session_ingress_token.to_string(),
94            api_base_url: api_base_url.to_string(),
95            sources: obj
96                .get("sources")
97                .and_then(|v| serde_json::from_value(v.clone()).ok())
98                .unwrap_or_default(),
99            auth: obj
100                .get("auth")
101                .and_then(|v| serde_json::from_value(v.clone()).ok())
102                .unwrap_or_default(),
103            claude_code_args: obj
104                .get("claude_code_args")
105                .and_then(|v| serde_json::from_value(v.clone()).ok()),
106            mcp_config: obj.get("mcp_config").cloned(),
107            environment_variables: obj
108                .get("environment_variables")
109                .and_then(|v| serde_json::from_value(v.clone()).ok()),
110            use_code_sessions: obj.get("use_code_sessions").and_then(|v| v.as_bool()),
111        };
112
113        Ok(work_secret)
114    } else {
115        Err("Invalid work secret: not an object".to_string())
116    }
117}
118
119/// Build a WebSocket SDK URL from the API base URL and session ID.
120/// Strips the HTTP(S) protocol and constructs a ws(s):// ingress URL.
121///
122/// Uses /v2/ for localhost (direct to session-ingress, no Envoy rewrite)
123/// and /v1/ for production (Envoy rewrites /v1/ -> /v2/).
124pub fn build_sdk_url(api_base_url: &str, session_id: &str) -> String {
125    let is_localhost = api_base_url.contains("localhost") || api_base_url.contains("127.0.0.1");
126    let protocol = if is_localhost { "ws" } else { "wss" };
127    let version = if is_localhost { "v2" } else { "v1" };
128    let host = api_base_url
129        .trim_start_matches("https://")
130        .trim_start_matches("http://")
131        .trim_end_matches('/');
132
133    format!(
134        "{}://{}/{}/session_ingress/ws/{}",
135        protocol, host, version, session_id
136    )
137}
138
139/// Compare two session IDs regardless of their tagged-ID prefix.
140///
141/// Tagged IDs have the form {tag}_{body} or {tag}_staging_{body}, where the
142/// body encodes a UUID. CCR v2's compat layer returns `session_*` to v1 API
143/// clients but the infrastructure layer uses `cse_*`. Both have the same
144/// underlying UUID.
145pub fn same_session_id(a: &str, b: &str) -> bool {
146    if a == b {
147        return true;
148    }
149
150    // The body is everything after the last underscore — this handles both
151    // `{tag}_{body}` and `{tag}_staging_{body}`.
152    let a_body = a.split('_').last().unwrap_or("");
153    let b_body = b.split('_').last().unwrap_or("");
154
155    // Guard against IDs with no underscore (bare UUIDs).
156    // Require a minimum length to avoid accidental matches on short suffixes.
157    a_body.len() >= 4 && a_body == b_body
158}
159
160/// Build a CCR v2 session URL from the API base URL and session ID.
161/// Returns an HTTP(S) URL (not ws://) and points at /v1/code/sessions/{id}.
162pub fn build_ccr_v2_sdk_url(api_base_url: &str, session_id: &str) -> String {
163    let base = api_base_url.trim_end_matches('/');
164    format!("{}/v1/code/sessions/{}", base, session_id)
165}
166
167/// Register this bridge as the worker for a CCR v2 session.
168/// Returns the worker_epoch, which must be passed to the child CC process.
169pub async fn register_worker(session_url: &str, access_token: &str) -> Result<u64, String> {
170    let client = reqwest::Client::new();
171
172    let response = client
173        .post(&format!("{}/worker/register", session_url))
174        .header("Authorization", format!("Bearer {}", access_token))
175        .header("Content-Type", "application/json")
176        .header("anthropic-version", "2023-06-01")
177        .header("User-Agent", get_user_agent())
178        .timeout(std::time::Duration::from_secs(10))
179        .send()
180        .await
181        .map_err(|e| format!("Request failed: {}", e))?;
182
183    let data: serde_json::Value = response
184        .json()
185        .await
186        .map_err(|e| format!("Failed to parse response: {}", e))?;
187
188    let raw = data.get("worker_epoch");
189
190    let epoch = match raw {
191        Some(v) if v.is_number() => v.as_u64(),
192        Some(v) if v.is_string() => v.as_str().and_then(|s| s.parse().ok()),
193        _ => None,
194    };
195
196    epoch.ok_or_else(|| {
197        format!(
198            "register_worker: invalid worker_epoch in response: {}",
199            data
200        )
201    })
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    #[test]
209    fn test_build_sdk_url() {
210        // Production
211        assert_eq!(
212            build_sdk_url("https://api.anthropic.com", "session_abc"),
213            "wss://api.anthropic.com/v1/session_ingress/ws/session_abc"
214        );
215
216        // Localhost
217        assert_eq!(
218            build_sdk_url("http://localhost:8080", "session_abc"),
219            "ws://localhost:8080/v2/session_ingress/ws/session_abc"
220        );
221    }
222
223    #[test]
224    fn test_same_session_id() {
225        // Same ID
226        assert!(same_session_id("session_abc123", "session_abc123"));
227
228        // Same UUID with different tags
229        assert!(same_session_id("cse_abc123", "session_abc123"));
230
231        // Different UUIDs
232        assert!(!same_session_id("session_abc123", "session_xyz789"));
233
234        // Staging format
235        assert!(same_session_id(
236            "cse_staging_abc123",
237            "session_staging_abc123"
238        ));
239    }
240
241    #[test]
242    fn test_build_ccr_v2_sdk_url() {
243        assert_eq!(
244            build_ccr_v2_sdk_url("https://api.anthropic.com", "session_abc"),
245            "https://api.anthropic.com/v1/code/sessions/session_abc"
246        );
247    }
248
249    #[test]
250    fn test_decode_work_secret() {
251        let secret = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
252            r#"{"version":1,"session_ingress_token":"tok123","api_base_url":"https://api.example.com","sources":[],"auth":[]}"#
253        );
254
255        let decoded = decode_work_secret(&secret).unwrap();
256        assert_eq!(decoded.version, 1);
257        assert_eq!(decoded.session_ingress_token, "tok123");
258        assert_eq!(decoded.api_base_url, "https://api.example.com");
259    }
260}
ai_agent/bridge/work_secret.rs

ai_agent/bridge/
work_secret.rs
