vtcode_auth/
mcp_oauth.rs

1//! OAuth support for HTTP MCP providers.
2
3use anyhow::{Context, Result, anyhow, bail};
4use base64::Engine;
5use base64::engine::general_purpose::URL_SAFE_NO_PAD;
6use reqwest::{Client, Url};
7use ring::rand::{SecureRandom, SystemRandom};
8use serde::{Deserialize, Serialize};
9use std::collections::BTreeMap;
10
11use crate::credentials::{AuthCredentialsStoreMode, CredentialStorage};
12use crate::pkce::{PkceChallenge, generate_pkce_challenge};
13
14const DEFAULT_CALLBACK_PORT: u16 = 8768;
15const DEFAULT_FLOW_TIMEOUT_SECS: u64 = 300;
16const REFRESH_SKEW_SECS: u64 = 60;
17
18/// Configuration for OAuth-enabled MCP HTTP providers.
19#[derive(Debug, Clone, Serialize, Deserialize)]
20#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
21#[serde(default)]
22pub struct McpOAuthConfig {
23    /// OAuth authorization endpoint.
24    pub authorization_url: String,
25    /// OAuth token endpoint.
26    pub token_url: String,
27    /// OAuth client identifier.
28    pub client_id: String,
29    /// Requested scopes.
30    #[serde(default)]
31    pub scopes: Vec<String>,
32    /// Optional audience/resource hint sent with the auth and token requests.
33    #[serde(default)]
34    pub audience: Option<String>,
35    /// Local callback server port.
36    pub callback_port: u16,
37    /// Browser-flow timeout in seconds.
38    pub flow_timeout_secs: u64,
39    /// Credential storage backend for this provider's token.
40    #[serde(default)]
41    pub credentials_store_mode: AuthCredentialsStoreMode,
42    /// Extra query parameters appended to the authorization URL.
43    #[serde(default)]
44    pub extra_auth_params: BTreeMap<String, String>,
45    /// Extra form fields appended to token exchanges and refreshes.
46    #[serde(default)]
47    pub extra_token_params: BTreeMap<String, String>,
48}
49
50impl Default for McpOAuthConfig {
51    fn default() -> Self {
52        Self {
53            authorization_url: String::new(),
54            token_url: String::new(),
55            client_id: String::new(),
56            scopes: Vec::new(),
57            audience: None,
58            callback_port: DEFAULT_CALLBACK_PORT,
59            flow_timeout_secs: DEFAULT_FLOW_TIMEOUT_SECS,
60            credentials_store_mode: AuthCredentialsStoreMode::default(),
61            extra_auth_params: BTreeMap::new(),
62            extra_token_params: BTreeMap::new(),
63        }
64    }
65}
66
67impl McpOAuthConfig {
68    pub fn validate(&self, provider_name: &str) -> Result<()> {
69        if self.authorization_url.trim().is_empty() {
70            bail!("MCP provider '{provider_name}' is missing oauth.authorization_url");
71        }
72        if self.token_url.trim().is_empty() {
73            bail!("MCP provider '{provider_name}' is missing oauth.token_url");
74        }
75        if self.client_id.trim().is_empty() {
76            bail!("MCP provider '{provider_name}' is missing oauth.client_id");
77        }
78        Ok(())
79    }
80
81    fn callback_url(&self) -> String {
82        format!("http://localhost:{}/auth/callback", self.callback_port)
83    }
84}
85
86/// Stored OAuth token for an MCP HTTP provider.
87#[derive(Debug, Clone, Serialize, Deserialize)]
88pub struct McpOAuthToken {
89    pub access_token: String,
90    pub refresh_token: Option<String>,
91    pub token_type: Option<String>,
92    pub scope: Option<String>,
93    pub obtained_at: u64,
94    pub expires_at: Option<u64>,
95}
96
97impl McpOAuthToken {
98    pub fn is_refresh_due(&self) -> bool {
99        self.expires_at
100            .is_some_and(|expires_at| now_secs().saturating_add(REFRESH_SKEW_SECS) >= expires_at)
101    }
102}
103
104/// Status for an MCP provider's stored OAuth token.
105#[derive(Debug, Clone)]
106pub enum McpOAuthStatus {
107    Authenticated {
108        age_seconds: u64,
109        expires_in: Option<u64>,
110    },
111    NotAuthenticated,
112}
113
114/// Prepared browser-login flow for an MCP OAuth provider.
115#[derive(Debug, Clone)]
116pub struct McpOAuthPreparedLogin {
117    pub auth_url: String,
118    pub callback_port: u16,
119    pub timeout_secs: u64,
120    pkce: PkceChallenge,
121    state: String,
122}
123
124impl McpOAuthPreparedLogin {
125    #[must_use]
126    pub fn expected_state(&self) -> &str {
127        &self.state
128    }
129}
130
131/// Completion payload kept intentionally close to Codex app-server.
132#[derive(Debug, Clone, PartialEq, Eq)]
133pub struct McpOAuthLoginCompletion {
134    pub name: String,
135    pub success: bool,
136    pub error: Option<String>,
137}
138
139/// Service for loading, refreshing, and persisting MCP OAuth tokens.
140#[derive(Debug, Clone, Default)]
141pub struct McpOAuthService;
142
143impl McpOAuthService {
144    #[must_use]
145    pub fn new() -> Self {
146        Self
147    }
148
149    pub fn prepare_login(
150        &self,
151        provider_name: &str,
152        config: &McpOAuthConfig,
153    ) -> Result<McpOAuthPreparedLogin> {
154        config.validate(provider_name)?;
155        let pkce = generate_pkce_challenge()?;
156        let state = generate_state()?;
157        let auth_url = build_auth_url(config, &pkce, &state)?;
158        Ok(McpOAuthPreparedLogin {
159            auth_url,
160            callback_port: config.callback_port,
161            timeout_secs: config.flow_timeout_secs,
162            pkce,
163            state,
164        })
165    }
166
167    pub async fn complete_login(
168        &self,
169        provider_name: &str,
170        config: &McpOAuthConfig,
171        prepared: &McpOAuthPreparedLogin,
172        code: &str,
173    ) -> Result<McpOAuthLoginCompletion> {
174        config.validate(provider_name)?;
175        let token = exchange_code_for_token(config, code, &prepared.pkce).await?;
176        save_token(provider_name, &token, config.credentials_store_mode)?;
177        Ok(McpOAuthLoginCompletion {
178            name: provider_name.to_string(),
179            success: true,
180            error: None,
181        })
182    }
183
184    pub fn status(
185        &self,
186        provider_name: &str,
187        storage_mode: AuthCredentialsStoreMode,
188    ) -> Result<McpOAuthStatus> {
189        let Some(token) = load_token(provider_name, storage_mode)? else {
190            return Ok(McpOAuthStatus::NotAuthenticated);
191        };
192        let now = now_secs();
193        Ok(McpOAuthStatus::Authenticated {
194            age_seconds: now.saturating_sub(token.obtained_at),
195            expires_in: token
196                .expires_at
197                .map(|expires_at| expires_at.saturating_sub(now)),
198        })
199    }
200
201    pub fn load_token(
202        &self,
203        provider_name: &str,
204        storage_mode: AuthCredentialsStoreMode,
205    ) -> Result<Option<McpOAuthToken>> {
206        load_token(provider_name, storage_mode)
207    }
208
209    pub async fn resolve_access_token(
210        &self,
211        provider_name: &str,
212        config: &McpOAuthConfig,
213    ) -> Result<Option<String>> {
214        let Some(mut token) = load_token(provider_name, config.credentials_store_mode)? else {
215            return Ok(None);
216        };
217
218        if token.is_refresh_due() {
219            if token.refresh_token.is_some() {
220                token = refresh_token(config, &token).await?;
221                save_token(provider_name, &token, config.credentials_store_mode)?;
222            } else {
223                bail!(
224                    "Stored MCP OAuth token for '{provider_name}' expired and cannot be refreshed. Run `vtcode mcp login {provider_name}` again."
225                );
226            }
227        }
228
229        Ok(Some(token.access_token))
230    }
231
232    pub fn logout(
233        &self,
234        provider_name: &str,
235        storage_mode: AuthCredentialsStoreMode,
236    ) -> Result<McpOAuthLoginCompletion> {
237        clear_token(provider_name, storage_mode)?;
238        Ok(McpOAuthLoginCompletion {
239            name: provider_name.to_string(),
240            success: true,
241            error: None,
242        })
243    }
244}
245
246fn build_auth_url(
247    config: &McpOAuthConfig,
248    challenge: &PkceChallenge,
249    state: &str,
250) -> Result<String> {
251    let mut url =
252        Url::parse(&config.authorization_url).context("invalid oauth.authorization_url")?;
253    {
254        let mut query = url.query_pairs_mut();
255        query.append_pair("response_type", "code");
256        query.append_pair("client_id", &config.client_id);
257        query.append_pair("redirect_uri", &config.callback_url());
258        query.append_pair("code_challenge", &challenge.code_challenge);
259        query.append_pair("code_challenge_method", &challenge.code_challenge_method);
260        query.append_pair("state", state);
261        if !config.scopes.is_empty() {
262            query.append_pair("scope", &config.scopes.join(" "));
263        }
264        if let Some(audience) = config.audience.as_deref()
265            && !audience.trim().is_empty()
266        {
267            query.append_pair("audience", audience);
268        }
269        for (key, value) in &config.extra_auth_params {
270            if !key.trim().is_empty() {
271                query.append_pair(key, value);
272            }
273        }
274    }
275    Ok(url.to_string())
276}
277
278async fn exchange_code_for_token(
279    config: &McpOAuthConfig,
280    code: &str,
281    challenge: &PkceChallenge,
282) -> Result<McpOAuthToken> {
283    let mut form = vec![
284        ("grant_type".to_string(), "authorization_code".to_string()),
285        ("client_id".to_string(), config.client_id.clone()),
286        ("code".to_string(), code.to_string()),
287        ("redirect_uri".to_string(), config.callback_url()),
288        (
289            "code_verifier".to_string(),
290            challenge.code_verifier.to_string(),
291        ),
292    ];
293    if let Some(audience) = config.audience.as_deref()
294        && !audience.trim().is_empty()
295    {
296        form.push(("audience".to_string(), audience.to_string()));
297    }
298    form.extend(
299        config
300            .extra_token_params
301            .iter()
302            .map(|(key, value)| (key.clone(), value.clone())),
303    );
304    send_token_request(&config.token_url, &form).await
305}
306
307async fn refresh_token(config: &McpOAuthConfig, current: &McpOAuthToken) -> Result<McpOAuthToken> {
308    let refresh_token = current
309        .refresh_token
310        .as_deref()
311        .filter(|value| !value.trim().is_empty())
312        .ok_or_else(|| anyhow!("Stored MCP OAuth token does not include a refresh token"))?;
313    let mut form = vec![
314        ("grant_type".to_string(), "refresh_token".to_string()),
315        ("client_id".to_string(), config.client_id.clone()),
316        ("refresh_token".to_string(), refresh_token.to_string()),
317    ];
318    if let Some(audience) = config.audience.as_deref()
319        && !audience.trim().is_empty()
320    {
321        form.push(("audience".to_string(), audience.to_string()));
322    }
323    form.extend(
324        config
325            .extra_token_params
326            .iter()
327            .map(|(key, value)| (key.clone(), value.clone())),
328    );
329
330    let refreshed = send_token_request(&config.token_url, &form).await?;
331    Ok(McpOAuthToken {
332        refresh_token: refreshed
333            .refresh_token
334            .or_else(|| current.refresh_token.clone()),
335        ..refreshed
336    })
337}
338
339async fn send_token_request(token_url: &str, form: &[(String, String)]) -> Result<McpOAuthToken> {
340    let response = Client::new()
341        .post(token_url)
342        .header("Content-Type", "application/x-www-form-urlencoded")
343        .form(form)
344        .send()
345        .await
346        .with_context(|| format!("failed to send MCP OAuth request to {token_url}"))?;
347    let status = response.status();
348    let body = response
349        .text()
350        .await
351        .context("failed to read MCP OAuth response body")?;
352
353    if !status.is_success() {
354        bail!("MCP OAuth request failed (HTTP {status}): {body}");
355    }
356
357    let payload: TokenResponse =
358        serde_json::from_str(&body).context("failed to parse MCP OAuth token response")?;
359    let now = now_secs();
360    Ok(McpOAuthToken {
361        access_token: payload.access_token,
362        refresh_token: payload.refresh_token,
363        token_type: payload.token_type,
364        scope: payload.scope,
365        obtained_at: now,
366        expires_at: payload.expires_in.map(|secs| now.saturating_add(secs)),
367    })
368}
369
370#[derive(Debug, Deserialize)]
371struct TokenResponse {
372    access_token: String,
373    #[serde(default)]
374    refresh_token: Option<String>,
375    #[serde(default)]
376    token_type: Option<String>,
377    #[serde(default)]
378    scope: Option<String>,
379    #[serde(default)]
380    expires_in: Option<u64>,
381}
382
383fn generate_state() -> Result<String> {
384    let mut state_bytes = [0_u8; 32];
385    SystemRandom::new()
386        .fill(&mut state_bytes)
387        .map_err(|_| anyhow!("failed to generate MCP OAuth state"))?;
388    Ok(URL_SAFE_NO_PAD.encode(state_bytes))
389}
390
391fn save_token(
392    provider_name: &str,
393    token: &McpOAuthToken,
394    storage_mode: AuthCredentialsStoreMode,
395) -> Result<()> {
396    let serialized = serde_json::to_string(token).context("failed to serialize MCP OAuth token")?;
397    token_storage(provider_name).store_with_mode(&serialized, storage_mode)
398}
399
400fn load_token(
401    provider_name: &str,
402    storage_mode: AuthCredentialsStoreMode,
403) -> Result<Option<McpOAuthToken>> {
404    let Some(serialized) = token_storage(provider_name).load_with_mode(storage_mode)? else {
405        return Ok(None);
406    };
407    serde_json::from_str(&serialized)
408        .context("failed to parse stored MCP OAuth token")
409        .map(Some)
410}
411
412fn clear_token(provider_name: &str, storage_mode: AuthCredentialsStoreMode) -> Result<()> {
413    token_storage(provider_name).clear_with_mode(storage_mode)
414}
415
416fn token_storage(provider_name: &str) -> CredentialStorage {
417    let normalized_provider = provider_name
418        .chars()
419        .map(|ch| {
420            if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
421                ch
422            } else {
423                '_'
424            }
425        })
426        .collect::<String>();
427    CredentialStorage::new("vtcode", format!("mcp_oauth_{normalized_provider}"))
428}
429
430fn now_secs() -> u64 {
431    std::time::SystemTime::now()
432        .duration_since(std::time::UNIX_EPOCH)
433        .map(|duration| duration.as_secs())
434        .unwrap_or(0)
435}
436
437#[cfg(test)]
438mod tests {
439    use super::*;
440    use assert_fs::TempDir;
441    use serial_test::serial;
442    use std::path::PathBuf;
443
444    struct TestAuthDirGuard {
445        previous: Option<PathBuf>,
446        temp_dir: Option<TempDir>,
447    }
448
449    impl TestAuthDirGuard {
450        fn new() -> Self {
451            let temp_dir = TempDir::new().expect("temp dir");
452            let previous = crate::storage_paths::auth_storage_dir_override_for_tests()
453                .expect("read previous auth dir override");
454            crate::storage_paths::set_auth_storage_dir_override_for_tests(Some(
455                temp_dir.path().to_path_buf(),
456            ))
457            .expect("set auth dir override");
458            Self {
459                previous,
460                temp_dir: Some(temp_dir),
461            }
462        }
463    }
464
465    impl Drop for TestAuthDirGuard {
466        fn drop(&mut self) {
467            crate::storage_paths::set_auth_storage_dir_override_for_tests(self.previous.clone())
468                .expect("restore auth dir override");
469            if let Some(temp_dir) = self.temp_dir.take() {
470                let _ = temp_dir.close();
471            }
472        }
473    }
474
475    fn sample_config() -> McpOAuthConfig {
476        McpOAuthConfig {
477            authorization_url: "https://example.com/oauth/authorize".to_string(),
478            token_url: "https://example.com/oauth/token".to_string(),
479            client_id: "client-123".to_string(),
480            scopes: vec!["mcp:read".to_string(), "mcp:write".to_string()],
481            audience: Some("mcp-api".to_string()),
482            callback_port: 8123,
483            flow_timeout_secs: 120,
484            credentials_store_mode: AuthCredentialsStoreMode::File,
485            extra_auth_params: BTreeMap::from([("prompt".to_string(), "consent".to_string())]),
486            extra_token_params: BTreeMap::new(),
487        }
488    }
489
490    #[test]
491    fn prepare_login_builds_expected_auth_url() {
492        let service = McpOAuthService::new();
493        let prepared = service
494            .prepare_login("demo", &sample_config())
495            .expect("prepare login");
496
497        assert!(prepared.auth_url.contains("response_type=code"));
498        assert!(prepared.auth_url.contains("client_id=client-123"));
499        assert!(prepared.auth_url.contains("scope=mcp%3Aread+mcp%3Awrite"));
500        assert!(prepared.auth_url.contains("audience=mcp-api"));
501        assert!(prepared.auth_url.contains("prompt=consent"));
502        assert!(prepared.auth_url.contains("code_challenge="));
503        assert!(prepared.auth_url.contains("state="));
504        assert_eq!(prepared.callback_port, 8123);
505        assert_eq!(prepared.timeout_secs, 120);
506    }
507
508    #[test]
509    #[serial]
510    fn status_reflects_stored_token() {
511        let _guard = TestAuthDirGuard::new();
512        let service = McpOAuthService::new();
513        let storage_mode = AuthCredentialsStoreMode::File;
514        assert!(matches!(
515            service.status("demo", storage_mode).expect("status"),
516            McpOAuthStatus::NotAuthenticated
517        ));
518
519        save_token(
520            "demo",
521            &McpOAuthToken {
522                access_token: "access".to_string(),
523                refresh_token: Some("refresh".to_string()),
524                token_type: Some("Bearer".to_string()),
525                scope: Some("mcp:read".to_string()),
526                obtained_at: now_secs(),
527                expires_at: Some(now_secs() + 3600),
528            },
529            storage_mode,
530        )
531        .expect("save token");
532
533        let status = service.status("demo", storage_mode).expect("status");
534        assert!(matches!(
535            status,
536            McpOAuthStatus::Authenticated {
537                expires_in: Some(_),
538                ..
539            }
540        ));
541    }
542
543    #[test]
544    #[serial]
545    fn logout_clears_stored_token() {
546        let _guard = TestAuthDirGuard::new();
547        let service = McpOAuthService::new();
548        let storage_mode = AuthCredentialsStoreMode::File;
549        save_token(
550            "demo",
551            &McpOAuthToken {
552                access_token: "access".to_string(),
553                refresh_token: None,
554                token_type: Some("Bearer".to_string()),
555                scope: None,
556                obtained_at: now_secs(),
557                expires_at: None,
558            },
559            storage_mode,
560        )
561        .expect("save token");
562
563        service.logout("demo", storage_mode).expect("logout");
564        assert!(load_token("demo", storage_mode).expect("load").is_none());
565    }
566}
vtcode_auth/mcp_oauth.rs

vtcode_auth/
mcp_oauth.rs
