From 2b583e08a6412a3af7806b90b37623c97e4893f0 Mon Sep 17 00:00:00 2001 From: Ben Lovell Date: Wed, 24 Dec 2025 19:34:04 +0000 Subject: [PATCH] fix: always read the JWT from disk per command, fixing team switching in MCP --- Cargo.lock | 1 + crates/config/Cargo.toml | 3 +- crates/config/src/lib.rs | 6 ++-- crates/config/src/session.rs | 55 ++++++++++++++++++++++++++---------- 4 files changed, 45 insertions(+), 20 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6f26a44f..483dbbc0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -482,6 +482,7 @@ dependencies = [ name = "config" version = "0.3.37" dependencies = [ + "base64", "chrono", "clap", "dirs", diff --git a/crates/config/Cargo.toml b/crates/config/Cargo.toml index a16d2f0e..39ec33a7 100644 --- a/crates/config/Cargo.toml +++ b/crates/config/Cargo.toml @@ -7,9 +7,10 @@ rust-version = { workspace = true } license = { workspace = true } [dependencies] +base64 = { workspace = true } chrono = { workspace = true } clap = { workspace = true } -dirs = { workspace = true } +dirs = { workspace = true } futures = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 5890e72a..a923147e 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -181,17 +181,15 @@ impl Config { configuration.base_path = base_path.to_string(); - if let Some(session) = &self.session { + // Always read from disk to pick up team switches + if let Ok(session) = Session::from_config_dir() { if let Some(active_team) = &session.active_team { - // Use the active team's JWT token configuration.bearer_access_token = Some(active_team.token.jwt.clone()); } else { - // Fall back to session token if no active team configuration.bearer_access_token = Some(session.token.jwt.clone()); } } - // Store the configuration in self configuration } } diff --git a/crates/config/src/session.rs b/crates/config/src/session.rs index 9148bb22..53830ccf 100644 --- a/crates/config/src/session.rs +++ b/crates/config/src/session.rs @@ -1,3 +1,4 @@ +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; use chrono::{DateTime, TimeZone, Utc}; use serde::{Deserialize, Serialize}; use std::fs; @@ -9,6 +10,20 @@ use crate::error::Error; use tower_api::apis::default_api::describe_session; use tower_telemetry::debug; +/// Extracts the account ID (aid) from a Tower JWT token. +/// Returns None if the JWT is malformed or doesn't contain an aid. +fn extract_aid_from_jwt(jwt: &str) -> Option { + let parts: Vec<&str> = jwt.split('.').collect(); + if parts.len() != 3 { + return None; + } + + let payload = parts[1]; + let decoded = URL_SAFE_NO_PAD.decode(payload).ok()?; + let json: serde_json::Value = serde_json::from_slice(&decoded).ok()?; + json.get("https://tower.dev/aid")?.as_str().map(String::from) +} + const DEFAULT_TOWER_URL: &str = "https://api.tower.dev"; pub fn default_tower_url() -> Url { @@ -163,6 +178,22 @@ impl Session { } } + /// Sets the active team based on an account ID (aid) extracted from a JWT. + /// Returns true if a matching team was found and set as active, false otherwise. + pub fn set_active_team_by_aid(&mut self, aid: &str) -> bool { + // Find the team whose JWT contains the matching aid + if let Some(team) = self + .teams + .iter() + .find(|team| extract_aid_from_jwt(&team.token.jwt).as_deref() == Some(aid)) + { + self.active_team = Some(team.clone()); + true + } else { + false + } + } + /// Updates the session with data from the API response pub fn update_from_api_response( &mut self, @@ -263,36 +294,30 @@ impl Session { } pub fn from_jwt(jwt: &str) -> Result { - // We need to instantiate our own configuration object here, instead of the typical thing - // that we do which is turn a Config into a Configuration. + let jwt_aid = extract_aid_from_jwt(jwt); + let mut config = tower_api::apis::configuration::Configuration::new(); config.bearer_access_token = Some(jwt.to_string()); - // We only pull TOWER_URL out of the environment here because we only ever use the JWT and - // all that in programmatic contexts (when TOWER_URL is set). - let tower_url = if let Ok(val) = std::env::var("TOWER_URL") { - val - } else { - DEFAULT_TOWER_URL.to_string() - }; - - // Setup the base path to point to the /v1 API endpoint as expected. + let tower_url = std::env::var("TOWER_URL").unwrap_or(DEFAULT_TOWER_URL.to_string()); let mut base_path = Url::parse(&tower_url).unwrap(); base_path.set_path("/v1"); - config.base_path = base_path.to_string(); - // This is a bit of a hairy thing: I didn't want to pull in too much from the Tower API - // client, so we're using the raw bindings here. match run_future_sync(describe_session(&config)) { Ok(resp) => { - // Now we need to extract the session from the response. let entity = resp.entity.unwrap(); match entity { tower_api::apis::default_api::DescribeSessionSuccess::Status200(resp) => { let mut session = Session::from_api_session(&resp.session); session.tower_url = base_path; + + if let Some(aid) = jwt_aid { + session.set_active_team_by_aid(&aid); + } + + session.save()?; Ok(session) } tower_api::apis::default_api::DescribeSessionSuccess::UnknownValue(val) => {