Merge pull request #502 from Milhouszhang/fix/501-pr495-followups

fix(providers): PR #495 follow-ups — copilot 401 retry, real SSE streaming, device-flow errors, gemini env keys (#501)

Author: Milhouszhang

crates/puffer-cli/src/copilot_login.rs

@@ -7,7 +7,7 @@
 //! provider credential; the runtime later exchanges it for a short-lived
 //! Copilot bearer (see `puffer-core/runtime/copilot.rs`).
 
-use anyhow::{bail, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use puffer_provider_registry::COPILOT_USER_AGENT;
 use serde::Deserialize;
 use std::time::Duration;
@@ -76,6 +76,7 @@ pub(crate) fn start_device_flow() -> Result<DeviceFlowStart> {
 }
 
 /// Outcome of a single poll of the device-flow token endpoint.
+#[derive(Debug)]
 pub(crate) enum DeviceFlowPoll {
     /// User has not authorized yet — keep polling.
     Pending,
@@ -87,39 +88,32 @@ pub(crate) enum DeviceFlowPoll {
     Failed(String),
 }
 
-/// Polls the token endpoint once with the device code.
-pub(crate) fn poll_device_flow(device_code: &str) -> Result<DeviceFlowPoll> {
-    #[derive(Deserialize)]
-    struct Resp {
-        #[serde(default)]
-        access_token: Option<String>,
-        #[serde(default)]
-        error: Option<String>,
+struct DevicePollHttpResponse {
+    status: reqwest::StatusCode,
+    body: String,
+}
+
+#[derive(Deserialize)]
+struct DevicePollResponse {
+    #[serde(default)]
+    access_token: Option<String>,
+    #[serde(default)]
+    error: Option<String>,
+}
+
+fn classify_device_flow_poll_response(
+    response: Result<DevicePollHttpResponse>,
+) -> Result<DeviceFlowPoll> {
+    let response = response?;
+    if !response.status.is_success() {
+        bail!(
+            "GitHub device-flow token poll failed ({}): {}",
+            response.status,
+            response.body
+        );
     }
-    let client = http_client()?;
-    // A single poll must not abort the whole login on a transient blip. Network
-    // errors and non-JSON/unknown bodies (e.g. a 5xx HTML error page from an
-    // infra hiccup) are treated as Pending so the caller keeps polling until the
-    // device code genuinely expires; only GitHub's documented terminal device-
-    // flow errors end the flow.
-    let response = match client
-        .post(ACCESS_TOKEN_URL)
-        .header("Accept", "application/json")
-        .header("User-Agent", COPILOT_USER_AGENT)
-        .form(&[
-            ("client_id", COPILOT_CLIENT_ID),
-            ("device_code", device_code),
-            ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"),
-        ])
-        .send()
-    {
-        Ok(response) => response,
-        Err(_) => return Ok(DeviceFlowPoll::Pending),
-    };
-    let body = response.text().unwrap_or_default();
-    let Ok(parsed) = serde_json::from_str::<Resp>(&body) else {
-        return Ok(DeviceFlowPoll::Pending);
-    };
+    let parsed: DevicePollResponse =
+        serde_json::from_str(&response.body).context("parsing GitHub device-flow poll response")?;
     if let Some(token) = parsed.access_token.filter(|t| !t.is_empty()) {
         return Ok(DeviceFlowPoll::Done(token));
     }
@@ -136,7 +130,80 @@ pub(crate) fn poll_device_flow(device_code: &str) -> Result<DeviceFlowPoll> {
             | "incorrect_device_code"
             | "device_flow_disabled"),
         ) => Ok(DeviceFlowPoll::Failed(err.to_string())),
-        // Unknown error code — treat as transient rather than aborting.
+        // Unknown GitHub error code — treat as transient rather than aborting.
         Some(_) => Ok(DeviceFlowPoll::Pending),
     }
 }
+
+/// Polls the token endpoint once with the device code.
+pub(crate) fn poll_device_flow(device_code: &str) -> Result<DeviceFlowPoll> {
+    let client = http_client()?;
+    // GitHub's device-flow protocol has explicit non-terminal states
+    // (`authorization_pending`, `slow_down`). Transport failures are not one of
+    // them: surface those as poll errors so desktop callers can use their
+    // consecutive-error guard instead of waiting until device-code expiry.
+    let response = match client
+        .post(ACCESS_TOKEN_URL)
+        .header("Accept", "application/json")
+        .header("User-Agent", COPILOT_USER_AGENT)
+        .form(&[
+            ("client_id", COPILOT_CLIENT_ID),
+            ("device_code", device_code),
+            ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"),
+        ])
+        .send()
+    {
+        Ok(response) => {
+            let status = response.status();
+            let body = response.text().unwrap_or_default();
+            Ok(DevicePollHttpResponse { status, body })
+        }
+        Err(error) => Err(anyhow!("GitHub device-flow poll network error: {error}")),
+    };
+    classify_device_flow_poll_response(response)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use anyhow::anyhow;
+    use reqwest::StatusCode;
+
+    fn classify_ok(body: &str) -> Result<DeviceFlowPoll> {
+        classify_device_flow_poll_response(Ok(DevicePollHttpResponse {
+            status: StatusCode::OK,
+            body: body.to_string(),
+        }))
+    }
+
+    #[test]
+    fn authorization_pending_remains_pending() {
+        let result = classify_ok(r#"{"error":"authorization_pending"}"#).unwrap();
+        assert!(matches!(result, DeviceFlowPoll::Pending));
+    }
+
+    #[test]
+    fn slow_down_remains_slow_down() {
+        let result = classify_ok(r#"{"error":"slow_down"}"#).unwrap();
+        assert!(matches!(result, DeviceFlowPoll::SlowDown));
+    }
+
+    #[test]
+    fn transport_errors_are_not_mapped_to_pending() {
+        let error = classify_device_flow_poll_response(Err(anyhow!("connection refused")))
+            .expect_err("transport failures must reject the poll RPC");
+        assert!(error.to_string().contains("connection refused"));
+    }
+
+    #[test]
+    fn malformed_poll_response_is_not_mapped_to_pending() {
+        let error = classify_ok("<html>bad gateway</html>")
+            .expect_err("malformed poll responses are not protocol pending states");
+        assert!(
+            error
+                .to_string()
+                .contains("parsing GitHub device-flow poll response"),
+            "{error:#}"
+        );
+    }
+}

crates/puffer-cli/src/non_interactive.rs

@@ -600,14 +600,20 @@ fn compose_prompt(
 }
 
 fn hydrate_env_auth(auth_store: &mut AuthStore) {
-    for (provider, env_name) in [
-        ("openai", "OPENAI_API_KEY"),
-        ("anthropic", "ANTHROPIC_API_KEY"),
-    ] {
-        if let Ok(value) = std::env::var(env_name) {
-            let trimmed = value.trim();
-            if !trimmed.is_empty() {
-                auth_store.set_api_key(provider, trimmed.to_string());
+    let mappings: &[(&str, &[&str])] = &[
+        ("openai", &["OPENAI_API_KEY"]),
+        ("anthropic", &["ANTHROPIC_API_KEY"]),
+        ("google", &["GEMINI_API_KEY", "GOOGLE_API_KEY"]),
+    ];
+
+    for (provider, env_names) in mappings {
+        for env_name in *env_names {
+            if let Ok(value) = std::env::var(env_name) {
+                let trimmed = value.trim();
+                if !trimmed.is_empty() {
+                    auth_store.set_api_key(*provider, trimmed.to_string());
+                    break;
+                }
             }
         }
     }
@@ -798,6 +804,47 @@ impl ReplayArtifact {
 mod tests {
     use super::*;
     use puffer_provider_registry::ProviderDescriptor;
+    use std::sync::{Mutex, OnceLock};
+
+    fn env_lock() -> &'static Mutex<()> {
+        static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
+        LOCK.get_or_init(|| Mutex::new(()))
+    }
+
+    struct EnvGuard {
+        name: &'static str,
+        prior: Option<String>,
+    }
+
+    impl EnvGuard {
+        fn set(name: &'static str, value: &str) -> Self {
+            let prior = std::env::var(name).ok();
+            std::env::set_var(name, value);
+            Self { name, prior }
+        }
+
+        fn remove(name: &'static str) -> Self {
+            let prior = std::env::var(name).ok();
+            std::env::remove_var(name);
+            Self { name, prior }
+        }
+    }
+
+    impl Drop for EnvGuard {
+        fn drop(&mut self) {
+            match self.prior.take() {
+                Some(value) => std::env::set_var(self.name, value),
+                None => std::env::remove_var(self.name),
+            }
+        }
+    }
+
+    fn api_key(auth_store: &AuthStore, provider: &str) -> Option<String> {
+        match auth_store.get(provider) {
+            Some(StoredCredential::ApiKey { key }) => Some(key.clone()),
+            _ => None,
+        }
+    }
 
     #[test]
     fn compose_prompt_includes_transcript_and_skill() {
@@ -833,6 +880,44 @@ mod tests {
         assert!(!looks_like_jsonl_transcript(r#"{"role":"user"}"#));
     }
 
+    #[test]
+    fn hydrate_env_auth_uses_gemini_api_key_for_google() {
+        let _lock = env_lock()
+            .lock()
+            .unwrap_or_else(|poison| poison.into_inner());
+        let _openai = EnvGuard::remove("OPENAI_API_KEY");
+        let _anthropic = EnvGuard::remove("ANTHROPIC_API_KEY");
+        let _google = EnvGuard::remove("GOOGLE_API_KEY");
+        let _gemini = EnvGuard::set("GEMINI_API_KEY", " gemini-key ");
+        let mut auth_store = AuthStore::default();
+
+        hydrate_env_auth(&mut auth_store);
+
+        assert_eq!(
+            api_key(&auth_store, "google").as_deref(),
+            Some("gemini-key")
+        );
+    }
+
+    #[test]
+    fn hydrate_env_auth_uses_google_api_key_alias_for_google() {
+        let _lock = env_lock()
+            .lock()
+            .unwrap_or_else(|poison| poison.into_inner());
+        let _openai = EnvGuard::remove("OPENAI_API_KEY");
+        let _anthropic = EnvGuard::remove("ANTHROPIC_API_KEY");
+        let _gemini = EnvGuard::remove("GEMINI_API_KEY");
+        let _google = EnvGuard::set("GOOGLE_API_KEY", " google-key ");
+        let mut auth_store = AuthStore::default();
+
+        hydrate_env_auth(&mut auth_store);
+
+        assert_eq!(
+            api_key(&auth_store, "google").as_deref(),
+            Some("google-key")
+        );
+    }
+
     #[test]
     fn custom_model_selector_registers_unknown_provider_model() {
         let mut providers = ProviderRegistry::new();

crates/puffer-cli/tests/tmux_agent_loop.rs

@@ -366,11 +366,11 @@ fn tmux_agent_loop_renders_assistant_reply_from_mock_anthropic() {
         "sh",
         &[
             "-lc",
-            // HOME=workspace makes `.puffer/` in the workspace be the
-            // workspace config dir. Tracing PUFFER_HTTP_TRACE_PATH lets
+            // PUFFER_HOME=workspace makes `.puffer/` in the workspace be the
+            // user config dir. Tracing PUFFER_HTTP_TRACE_PATH lets
             // post-mortem inspection see the wire bytes if the test fails.
             &format!(
-                "HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'",
+                "PUFFER_HOME='{ws}' HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'",
                 ws = workspace.display(),
                 bin = binary
             ),
@@ -440,7 +440,7 @@ fn tmux_agent_loop_accepts_codex_default_provider_alias() {
         &[
             "-lc",
             &format!(
-                "HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'",
+                "PUFFER_HOME='{ws}' HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'",
                 ws = workspace.display(),
                 bin = binary
             ),
@@ -548,7 +548,7 @@ fn tmux_agent_loop_drives_tool_round_trip_in_tui() {
         &[
             "-lc",
             &format!(
-                "HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'",
+                "PUFFER_HOME='{ws}' HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'",
                 ws = workspace.display(),
                 bin = binary
             ),
@@ -641,7 +641,7 @@ fn tmux_agent_loop_validates_workflow_shorthand_in_tui() {
         &[
             "-lc",
             &format!(
-                "HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'",
+                "PUFFER_HOME='{ws}' HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'",
                 ws = workspace.display(),
                 bin = binary
             ),
@@ -743,7 +743,7 @@ fn tmux_agent_loop_answers_ask_user_question_with_keyboard_selection() {
         &[
             "-lc",
             &format!(
-                "HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'",
+                "PUFFER_HOME='{ws}' HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'",
                 ws = workspace.display(),
                 bin = binary
             ),

crates/puffer-core/runtime/copilot.rs

@@ -11,6 +11,8 @@ use anyhow::{bail, Context, Result};
 use puffer_provider_registry::{apply_copilot_client_identity, COPILOT_TOKEN_URL};
 use reqwest::blocking::Client;
 use std::collections::HashMap;
+#[cfg(test)]
+use std::sync::Arc;
 use std::sync::{Mutex, OnceLock};
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
@@ -79,6 +81,34 @@ fn cache() -> &'static Mutex<HashMap<String, CachedToken>> {
     CACHE.get_or_init(|| Mutex::new(HashMap::new()))
 }
 
+#[cfg(test)]
+type TestBearerExchange = Arc<dyn Fn(&str) -> Result<CopilotAuth> + Send + Sync>;
+
+#[cfg(test)]
+fn test_bearer_exchange() -> &'static Mutex<Option<TestBearerExchange>> {
+    static EXCHANGE: OnceLock<Mutex<Option<TestBearerExchange>>> = OnceLock::new();
+    EXCHANGE.get_or_init(|| Mutex::new(None))
+}
+
+#[cfg(test)]
+pub(super) struct TestBearerExchangeGuard;
+
+#[cfg(test)]
+impl Drop for TestBearerExchangeGuard {
+    fn drop(&mut self) {
+        *test_bearer_exchange().lock().unwrap() = None;
+    }
+}
+
+#[cfg(test)]
+pub(super) fn install_test_bearer_exchange<F>(exchange: F) -> TestBearerExchangeGuard
+where
+    F: Fn(&str) -> Result<CopilotAuth> + Send + Sync + 'static,
+{
+    *test_bearer_exchange().lock().unwrap() = Some(Arc::new(exchange));
+    TestBearerExchangeGuard
+}
+
 /// Drops the cached exchanged bearer for a GitHub token. Called when the chat
 /// endpoint rejects the bearer with 401 (e.g. it was invalidated before its
 /// cached expiry — revoked Copilot seat, server-side rotation), and on
@@ -109,6 +139,11 @@ fn now_secs() -> u64 {
 /// small burst at first use / expiry; a shared in-flight map would add locking
 /// complexity for little gain, so we accept it.
 pub(crate) fn copilot_bearer_token(github_token: &str) -> Result<CopilotAuth> {
+    #[cfg(test)]
+    if let Some(exchange) = test_bearer_exchange().lock().unwrap().clone() {
+        return exchange(github_token);
+    }
+
     let now = now_secs();
     if let Some(cached) = cache().lock().unwrap().get(github_token) {
         if cached.expires_at_secs > now + EXPIRY_SKEW_SECS {