tenant.openagents/omega

1use std::{
2    io::{ErrorKind, Read, Write},
3    net::{Ipv4Addr, TcpListener, TcpStream},
4    sync::{Arc, Mutex},
5    time::{Duration, Instant},
6};
7
8use anyhow::{Context as _, Result, anyhow};
9use base64::Engine as _;
10use credentials_provider::CredentialsProvider;
11use gpui::{App, AsyncApp, Global};
12use http_client::{AsyncBody, HttpClient, Method, Request, StatusCode};
13use rand::RngCore as _;
14use serde::{Deserialize, Serialize};
15use sha2::{Digest as _, Sha256};
16use smol::io::AsyncReadExt as _;
17use url::{Url, form_urlencoded};
18
19pub const OPENAGENTS_DESKTOP_CLIENT_ID: &str = "openagents-desktop";
20pub const OPENAGENTS_AUTHORIZE_URL: &str = "https://auth.openagents.com/authorize";
21pub const OPENAGENTS_TOKEN_URL: &str = "https://auth.openagents.com/token";
22pub const OPENAGENTS_BASE_URL: &str = "https://openagents.com";
23pub const OPENAGENTS_AUTH_SESSION_URL: &str = "https://openagents.com/api/mobile/auth/session";
24pub const OPENAGENTS_CALLBACK_PATH: &str = "/auth/callback";
25pub const OPENAGENTS_SESSION_KEY: &str = "omega://openagents/native-session/v1";
26
27const OPENAGENTS_REFRESH_HEADER: &str = "x-openagents-refresh-token";
28const CALLBACK_TIMEOUT: Duration = Duration::from_secs(120);
29const MAX_HTTP_BODY_BYTES: u64 = 64 * 1024;
30const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024;
31
32#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize)]
33#[serde(rename_all = "snake_case")]
34pub enum OpenAgentsSessionPhase {
35    #[default]
36    SignedOut,
37    Connecting,
38    Ready,
39    Denied,
40    Unavailable,
41    Disconnecting,
42}
43
44impl OpenAgentsSessionPhase {
45    pub fn label(self) -> &'static str {
46        match self {
47            Self::SignedOut => "Not connected",
48            Self::Connecting => "Connecting in your browser…",
49            Self::Ready => "Connected",
50            Self::Denied => "Session expired — reconnect required",
51            Self::Unavailable => "OpenAgents account unavailable",
52            Self::Disconnecting => "Disconnecting…",
53        }
54    }
55}
56
57#[derive(Clone, Debug, Eq, PartialEq)]
58pub struct VerifiedOpenAgentsSession {
59    pub base_url: String,
60    pub access_token: String,
61}
62
63#[derive(Clone)]
64pub struct OpenAgentsSession {
65    credentials: Arc<dyn CredentialsProvider>,
66    http_client: Arc<dyn HttpClient>,
67    phase: Arc<Mutex<OpenAgentsSessionPhase>>,
68}
69
70struct OpenAgentsSessionGlobal(OpenAgentsSession);
71impl Global for OpenAgentsSessionGlobal {}
72
73#[derive(Debug, Deserialize, Serialize)]
74#[serde(rename_all = "camelCase", deny_unknown_fields)]
75struct StoredCredential {
76    schema_version: u8,
77    owner_user_id: String,
78    access_token: String,
79    refresh_token: String,
80}
81
82#[derive(Deserialize)]
83struct TokenResponse {
84    access_token: String,
85    refresh_token: String,
86    expires_in: Option<f64>,
87}
88
89#[derive(Deserialize)]
90struct VerifiedSessionResponse {
91    authenticated: bool,
92    user: VerifiedSessionUser,
93    tokens: Option<RotatedTokens>,
94}
95
96#[derive(Deserialize)]
97#[serde(rename_all = "camelCase")]
98struct VerifiedSessionUser {
99    user_id: String,
100}
101
102#[derive(Deserialize)]
103#[serde(rename_all = "camelCase")]
104struct RotatedTokens {
105    access: String,
106    refresh: String,
107    expires_in: f64,
108}
109
110#[derive(Deserialize)]
111#[serde(rename_all = "camelCase")]
112struct RevokedSessionResponse {
113    signed_out: bool,
114    access_revoked: bool,
115    refresh_revoked: bool,
116}
117
118enum CallbackResult {
119    Code(String),
120    Cancelled,
121    Unavailable,
122}
123
124enum VerificationResult {
125    Verified(StoredCredential),
126    Denied,
127    Unavailable,
128}
129
130pub fn init_openagents_session(cx: &mut App) {
131    if cx.try_global::<OpenAgentsSessionGlobal>().is_some() {
132        return;
133    }
134    cx.set_global(OpenAgentsSessionGlobal(OpenAgentsSession {
135        credentials: zed_credentials_provider::system_keychain(cx),
136        http_client: cx.http_client(),
137        phase: Arc::new(Mutex::new(OpenAgentsSessionPhase::SignedOut)),
138    }));
139}
140
141pub fn openagents_session(cx: &App) -> OpenAgentsSession {
142    cx.global::<OpenAgentsSessionGlobal>().0.clone()
143}
144
145impl OpenAgentsSession {
146    pub fn phase(&self) -> OpenAgentsSessionPhase {
147        self.phase.lock().map(|phase| *phase).unwrap_or_default()
148    }
149
150    fn set_phase(&self, phase: OpenAgentsSessionPhase) {
151        if let Ok(mut current) = self.phase.lock() {
152            *current = phase;
153        }
154    }
155
156    pub async fn connect(&self, cx: &mut AsyncApp) -> OpenAgentsSessionPhase {
157        self.set_phase(OpenAgentsSessionPhase::Connecting);
158        let phase = match self.connect_inner(cx).await {
159            Ok(Some(credential)) if self.save_credential(&credential, cx).await.is_ok() => {
160                OpenAgentsSessionPhase::Ready
161            }
162            Ok(None) => OpenAgentsSessionPhase::SignedOut,
163            Ok(Some(_)) | Err(_) => OpenAgentsSessionPhase::Unavailable,
164        };
165        self.set_phase(phase);
166        phase
167    }
168
169    pub async fn disconnect(&self, cx: &mut AsyncApp) -> OpenAgentsSessionPhase {
170        self.set_phase(OpenAgentsSessionPhase::Disconnecting);
171        let phase = match self.load_credential(cx).await {
172            Ok(None) => OpenAgentsSessionPhase::SignedOut,
173            Ok(Some(credential)) => match self.revoke_credential(&credential).await {
174                Ok(true) if self.delete_credential(cx).await.is_ok() => {
175                    OpenAgentsSessionPhase::SignedOut
176                }
177                Ok(true) | Ok(false) | Err(_) => OpenAgentsSessionPhase::Unavailable,
178            },
179            Err(_) => OpenAgentsSessionPhase::Unavailable,
180        };
181        self.set_phase(phase);
182        phase
183    }
184
185    pub async fn resolve_verified(&self, cx: &mut AsyncApp) -> Option<VerifiedOpenAgentsSession> {
186        let credential = match self.load_credential(cx).await {
187            Ok(Some(credential)) => credential,
188            Ok(None) => {
189                self.set_phase(OpenAgentsSessionPhase::SignedOut);
190                return None;
191            }
192            Err(_) => {
193                self.set_phase(OpenAgentsSessionPhase::Unavailable);
194                return None;
195            }
196        };
197        match self.verify_credential(&credential).await {
198            VerificationResult::Verified(verified) => {
199                if verified.access_token.len() > MAX_ACCESS_TOKEN_BYTES
200                    || self.save_credential(&verified, cx).await.is_err()
201                {
202                    self.set_phase(OpenAgentsSessionPhase::Unavailable);
203                    return None;
204                }
205                self.set_phase(OpenAgentsSessionPhase::Ready);
206                Some(VerifiedOpenAgentsSession {
207                    base_url: OPENAGENTS_BASE_URL.to_string(),
208                    access_token: verified.access_token,
209                })
210            }
211            VerificationResult::Denied => {
212                let phase = if self.delete_credential(cx).await.is_ok() {
213                    OpenAgentsSessionPhase::Denied
214                } else {
215                    OpenAgentsSessionPhase::Unavailable
216                };
217                self.set_phase(phase);
218                None
219            }
220            VerificationResult::Unavailable => {
221                self.set_phase(OpenAgentsSessionPhase::Unavailable);
222                None
223            }
224        }
225    }
226
227    async fn connect_inner(&self, cx: &mut AsyncApp) -> Result<Option<StoredCredential>> {
228        let state = random_base64_url();
229        let verifier = random_base64_url();
230        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0))
231            .context("OpenAgents callback listener unavailable")?;
232        listener.set_nonblocking(true)?;
233        let port = listener.local_addr()?.port();
234        if port < 1024 {
235            return Err(anyhow!("OpenAgents callback listener unavailable"));
236        }
237        let redirect_uri = format!("http://127.0.0.1:{port}{OPENAGENTS_CALLBACK_PATH}");
238        let authorize_url = build_authorize_url(&redirect_uri, &state, &verifier)?;
239        cx.update(|cx| cx.open_url(authorize_url.as_str()));
240        let background_executor = cx.background_executor().clone();
241        let code =
242            match wait_for_callback(&listener, &state, CALLBACK_TIMEOUT, &background_executor)
243                .await?
244            {
245                CallbackResult::Code(code) => code,
246                CallbackResult::Cancelled => return Ok(None),
247                CallbackResult::Unavailable => {
248                    return Err(anyhow!("OpenAgents authorization timed out"));
249                }
250            };
251        let tokens = self.exchange_code(&code, &verifier, &redirect_uri).await?;
252        let credential = StoredCredential {
253            schema_version: 1,
254            owner_user_id: String::new(),
255            access_token: tokens.access_token.trim().to_string(),
256            refresh_token: tokens.refresh_token.trim().to_string(),
257        };
258        if credential.access_token.is_empty()
259            || credential.refresh_token.is_empty()
260            || credential.access_token.len() > MAX_ACCESS_TOKEN_BYTES
261            || tokens
262                .expires_in
263                .is_some_and(|expires| !expires.is_finite() || expires <= 0.0)
264        {
265            return Err(anyhow!("OpenAgents token exchange was invalid"));
266        }
267        match self.verify_credential(&credential).await {
268            VerificationResult::Verified(credential) => Ok(Some(credential)),
269            VerificationResult::Denied | VerificationResult::Unavailable => {
270                Err(anyhow!("OpenAgents session verification failed"))
271            }
272        }
273    }
274
275    async fn exchange_code(
276        &self,
277        code: &str,
278        verifier: &str,
279        redirect_uri: &str,
280    ) -> Result<TokenResponse> {
281        let body = form_urlencoded::Serializer::new(String::new())
282            .append_pair("client_id", OPENAGENTS_DESKTOP_CLIENT_ID)
283            .append_pair("code", code)
284            .append_pair("code_verifier", verifier)
285            .append_pair("grant_type", "authorization_code")
286            .append_pair("redirect_uri", redirect_uri)
287            .finish();
288        let request = Request::builder()
289            .method(Method::POST)
290            .uri(OPENAGENTS_TOKEN_URL)
291            .header("content-type", "application/x-www-form-urlencoded")
292            .body(AsyncBody::from(body))?;
293        let (status, body) = send_json(&self.http_client, request).await?;
294        if !status.is_success() {
295            return Err(anyhow!("OpenAgents token exchange failed ({status})"));
296        }
297        serde_json::from_slice(&body).context("OpenAgents token exchange response was invalid")
298    }
299
300    async fn verify_credential(&self, credential: &StoredCredential) -> VerificationResult {
301        let request = match authenticated_session_request(Method::GET, credential) {
302            Ok(request) => request,
303            Err(_) => return VerificationResult::Unavailable,
304        };
305        let (status, body) = match send_json(&self.http_client, request).await {
306            Ok(response) => response,
307            Err(_) => return VerificationResult::Unavailable,
308        };
309        if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN {
310            return VerificationResult::Denied;
311        }
312        if !status.is_success() {
313            return VerificationResult::Unavailable;
314        }
315        let Ok(session) = serde_json::from_slice::<VerifiedSessionResponse>(&body) else {
316            return VerificationResult::Unavailable;
317        };
318        let owner_user_id = session.user.user_id.trim();
319        if !session.authenticated
320            || owner_user_id.is_empty()
321            || (!credential.owner_user_id.is_empty() && credential.owner_user_id != owner_user_id)
322        {
323            return VerificationResult::Denied;
324        }
325        let mut verified = StoredCredential {
326            schema_version: 1,
327            owner_user_id: owner_user_id.to_string(),
328            access_token: credential.access_token.clone(),
329            refresh_token: credential.refresh_token.clone(),
330        };
331        if let Some(tokens) = session.tokens {
332            if tokens.access.trim().is_empty()
333                || tokens.refresh.trim().is_empty()
334                || tokens.access.len() > MAX_ACCESS_TOKEN_BYTES
335                || !tokens.expires_in.is_finite()
336                || tokens.expires_in <= 0.0
337            {
338                return VerificationResult::Unavailable;
339            }
340            verified.access_token = tokens.access.trim().to_string();
341            verified.refresh_token = tokens.refresh.trim().to_string();
342        }
343        VerificationResult::Verified(verified)
344    }
345
346    async fn revoke_credential(&self, credential: &StoredCredential) -> Result<bool> {
347        let request = authenticated_session_request(Method::DELETE, credential)?;
348        let (status, body) = send_json(&self.http_client, request).await?;
349        if !status.is_success() {
350            return Ok(false);
351        }
352        let proof: RevokedSessionResponse =
353            serde_json::from_slice(&body).context("OpenAgents revoke response was invalid")?;
354        Ok(proof.signed_out && proof.access_revoked && proof.refresh_revoked)
355    }
356
357    async fn load_credential(&self, cx: &AsyncApp) -> Result<Option<StoredCredential>> {
358        let Some((username, secret)) = self
359            .credentials
360            .read_credentials(OPENAGENTS_SESSION_KEY, cx)
361            .await?
362        else {
363            return Ok(None);
364        };
365        let credential: StoredCredential = serde_json::from_slice(&secret)
366            .context("OpenAgents keychain credential was invalid")?;
367        if credential.schema_version != 1
368            || credential.owner_user_id != username
369            || credential.owner_user_id.trim().is_empty()
370            || credential.access_token.trim().is_empty()
371            || credential.refresh_token.trim().is_empty()
372            || credential.access_token.len() > MAX_ACCESS_TOKEN_BYTES
373        {
374            return Err(anyhow!("OpenAgents keychain credential was invalid"));
375        }
376        Ok(Some(credential))
377    }
378
379    async fn save_credential(&self, credential: &StoredCredential, cx: &AsyncApp) -> Result<()> {
380        if credential.owner_user_id.trim().is_empty()
381            || credential.access_token.trim().is_empty()
382            || credential.refresh_token.trim().is_empty()
383        {
384            return Err(anyhow!("OpenAgents credential was incomplete"));
385        }
386        let secret = serde_json::to_vec(credential)?;
387        self.credentials
388            .write_credentials(
389                OPENAGENTS_SESSION_KEY,
390                &credential.owner_user_id,
391                &secret,
392                cx,
393            )
394            .await
395    }
396
397    async fn delete_credential(&self, cx: &AsyncApp) -> Result<()> {
398        self.credentials
399            .delete_credentials(OPENAGENTS_SESSION_KEY, cx)
400            .await
401    }
402}
403
404fn authenticated_session_request(
405    method: Method,
406    credential: &StoredCredential,
407) -> Result<http_client::Request<AsyncBody>> {
408    Ok(Request::builder()
409        .method(method)
410        .uri(OPENAGENTS_AUTH_SESSION_URL)
411        .header(
412            "authorization",
413            format!("Bearer {}", credential.access_token),
414        )
415        .header(OPENAGENTS_REFRESH_HEADER, &credential.refresh_token)
416        .body(AsyncBody::empty())?)
417}
418
419async fn send_json(
420    client: &Arc<dyn HttpClient>,
421    request: http_client::Request<AsyncBody>,
422) -> Result<(StatusCode, Vec<u8>)> {
423    let mut response = client.send(request).await?;
424    let status = response.status();
425    let mut body = Vec::new();
426    response
427        .body_mut()
428        .take(MAX_HTTP_BODY_BYTES)
429        .read_to_end(&mut body)
430        .await?;
431    Ok((status, body))
432}
433
434fn random_base64_url() -> String {
435    let mut bytes = [0_u8; 32];
436    rand::rng().fill_bytes(&mut bytes);
437    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
438}
439
440fn pkce_challenge(verifier: &str) -> String {
441    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()))
442}
443
444fn build_authorize_url(redirect_uri: &str, state: &str, verifier: &str) -> Result<Url> {
445    let mut authorize = Url::parse(OPENAGENTS_AUTHORIZE_URL)?;
446    authorize
447        .query_pairs_mut()
448        .append_pair("client_id", OPENAGENTS_DESKTOP_CLIENT_ID)
449        .append_pair("code_challenge", &pkce_challenge(verifier))
450        .append_pair("code_challenge_method", "S256")
451        .append_pair("provider", "github")
452        .append_pair("redirect_uri", redirect_uri)
453        .append_pair("response_type", "code")
454        .append_pair("state", state);
455    Ok(authorize)
456}
457
458async fn wait_for_callback(
459    listener: &TcpListener,
460    expected_state: &str,
461    timeout: Duration,
462    background_executor: &gpui::BackgroundExecutor,
463) -> Result<CallbackResult> {
464    let deadline = Instant::now() + timeout;
465    loop {
466        match listener.accept() {
467            Ok((mut stream, _)) => {
468                if let Some(result) = read_callback(&mut stream, expected_state)? {
469                    return Ok(result);
470                }
471            }
472            Err(error) if error.kind() == ErrorKind::WouldBlock => {}
473            Err(error) => return Err(error.into()),
474        }
475        if Instant::now() >= deadline {
476            return Ok(CallbackResult::Unavailable);
477        }
478        background_executor.timer(Duration::from_millis(50)).await;
479    }
480}
481
482fn read_callback(stream: &mut TcpStream, expected_state: &str) -> Result<Option<CallbackResult>> {
483    stream.set_read_timeout(Some(Duration::from_secs(2)))?;
484    let mut buffer = [0_u8; 8192];
485    let size = stream.read(&mut buffer)?;
486    let request = std::str::from_utf8(&buffer[..size])?;
487    let Some(request_line) = request.lines().next() else {
488        write_callback_response(stream, 400)?;
489        return Ok(None);
490    };
491    let mut parts = request_line.split_whitespace();
492    if parts.next().unwrap_or_default() != "GET" {
493        write_callback_response(stream, 400)?;
494        return Ok(None);
495    }
496    let target = parts.next().unwrap_or_default();
497    let url = Url::parse(&format!("http://127.0.0.1{target}"))?;
498    if url.path() != OPENAGENTS_CALLBACK_PATH {
499        write_callback_response(stream, 404)?;
500        return Ok(None);
501    }
502    let parameter = |name: &str| {
503        url.query_pairs()
504            .find_map(|(key, value)| (key == name).then(|| value.into_owned()))
505    };
506    if parameter("state").as_deref() != Some(expected_state) {
507        write_callback_response(stream, 400)?;
508        return Ok(None);
509    }
510    if parameter("error").is_some_and(|error| !error.trim().is_empty()) {
511        write_callback_response(stream, 200)?;
512        return Ok(Some(CallbackResult::Cancelled));
513    }
514    let code = parameter("code").unwrap_or_default();
515    if code.trim().is_empty() {
516        write_callback_response(stream, 400)?;
517        return Ok(None);
518    }
519    write_callback_response(stream, 200)?;
520    Ok(Some(CallbackResult::Code(code)))
521}
522
523fn write_callback_response(stream: &mut TcpStream, status: u16) -> Result<()> {
524    let reason = match status {
525        200 => "OK",
526        404 => "Not Found",
527        _ => "Bad Request",
528    };
529    let body = "<!doctype html><meta charset=utf-8><meta name=referrer content=no-referrer><title>Omega</title><p>You can return to Omega.</p>";
530    write!(
531        stream,
532        "HTTP/1.1 {status} {reason}\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nCache-Control: no-store\r\nContent-Security-Policy: default-src 'none'\r\nReferrer-Policy: no-referrer\r\nX-Content-Type-Options: nosniff\r\nConnection: close\r\n\r\n{body}",
533        body.len()
534    )?;
535    stream.flush()?;
536    Ok(())
537}
538
539#[cfg(test)]
540mod tests {
541    use super::*;
542
543    #[test]
544    fn authorize_url_uses_admitted_desktop_pkce_contract() {
545        let url = build_authorize_url(
546            "http://127.0.0.1:49152/auth/callback",
547            "state-fixture",
548            "verifier-fixture-with-enough-entropy-for-the-test",
549        )
550        .expect("authorize URL");
551        assert_eq!(
552            url.as_str().split('?').next(),
553            Some(OPENAGENTS_AUTHORIZE_URL)
554        );
555        let values = url
556            .query_pairs()
557            .collect::<std::collections::HashMap<_, _>>();
558        assert_eq!(
559            values.get("client_id").map(|value| value.as_ref()),
560            Some("openagents-desktop")
561        );
562        assert_eq!(
563            values.get("provider").map(|value| value.as_ref()),
564            Some("github")
565        );
566        assert_eq!(
567            values
568                .get("code_challenge_method")
569                .map(|value| value.as_ref()),
570            Some("S256")
571        );
572        assert_eq!(
573            values.get("state").map(|value| value.as_ref()),
574            Some("state-fixture")
575        );
576        assert_eq!(
577            values.get("redirect_uri").map(|value| value.as_ref()),
578            Some("http://127.0.0.1:49152/auth/callback")
579        );
580        assert_eq!(
581            values.get("code_challenge").map(|value| value.len()),
582            Some(43)
583        );
584    }
585
586    #[test]
587    fn keychain_namespace_and_public_phase_never_contain_tokens() {
588        assert_eq!(
589            OPENAGENTS_SESSION_KEY,
590            "omega://openagents/native-session/v1"
591        );
592        assert!(!OPENAGENTS_SESSION_KEY.contains("token"));
593        let projection =
594            serde_json::to_string(&OpenAgentsSessionPhase::Ready).expect("serialize phase");
595        assert_eq!(projection, "\"ready\"");
596        assert!(!projection.contains("access-fixture"));
597    }
598
599    #[test]
600    fn callback_requires_exact_path_state_and_hides_code() {
601        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).expect("listener");
602        let address = listener.local_addr().expect("address");
603        let client = std::thread::spawn(move || {
604            let mut stream = TcpStream::connect(address).expect("connect");
605            write!(stream, "GET /auth/callback?state=exact&code=private-code HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n").expect("write");
606            let mut response = String::new();
607            stream.read_to_string(&mut response).expect("response");
608            response
609        });
610        let (mut stream, _) = listener.accept().expect("accept");
611        let result = read_callback(&mut stream, "exact").expect("callback");
612        assert!(matches!(result, Some(CallbackResult::Code(code)) if code == "private-code"));
613        drop(stream);
614        assert!(!client.join().expect("client").contains("private-code"));
615    }
616
617    #[test]
618    fn rotation_and_revoke_decoders_fail_closed() {
619        let session: VerifiedSessionResponse = serde_json::from_value(serde_json::json!({
620            "authenticated": true,
621            "user": { "userId": "owner.fixture" },
622            "tokens": { "access": "rotated-access", "refresh": "rotated-refresh", "expiresIn": 3600 }
623        })).expect("verified session");
624        assert_eq!(session.user.user_id, "owner.fixture");
625        assert_eq!(session.tokens.expect("rotation").access, "rotated-access");
626
627        let incomplete: RevokedSessionResponse = serde_json::from_value(serde_json::json!({
628            "signedOut": true, "accessRevoked": true, "refreshRevoked": false
629        }))
630        .expect("revoke proof");
631        assert!(
632            !(incomplete.signed_out && incomplete.access_revoked && incomplete.refresh_revoked)
633        );
634    }
635}
636