From c8ffe76688f540d0582254707a17c1822cee9fd3 Mon Sep 17 00:00:00 2001
From: Michael Freno <michaelt.freno@gmail.com>
Date: Sun, 10 May 2026 18:06:16 -0400
Subject: [PATCH] FRE-4763: Fix Clone() context argument and apply Code
 Reviewer P0-P3 fixes

- P0: Update auth header after token refresh via GetSession() + SetAuthHeader()
- P2: Unconditional req.WithContext(ctx) instead of fragile context.Background() check
- Fix: req.Clone(ctx) takes context.Context, not *http.Request (req.WithContext(ctx))
- Remove unused checkAuthenticated() and NewRequestWithContext() helpers
---
 cmd/mail.go            | 16 ----------------
 internal/api/client.go | 24 +++++++++---------------
 2 files changed, 9 insertions(+), 31 deletions(-)

diff --git a/cmd/mail.go b/cmd/mail.go
index f1514b5..51d80a4 100644
--- a/cmd/mail.go
+++ b/cmd/mail.go
@@ -15,22 +15,6 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func checkAuthenticated() (*auth.Session, error) {
-	sessionMgr, err := auth.NewSessionManager()
-	if err != nil {
-		return nil, fmt.Errorf("failed to create session manager: %w", err)
-	}
-	authenticated, err := sessionMgr.IsAuthenticated()
-	if err != nil || !authenticated {
-		return nil, fmt.Errorf("not authenticated (run 'pop login' first): %w", err)
-	}
-	session, err := sessionMgr.GetSession()
-	if err != nil {
-		return nil, fmt.Errorf("not authenticated: %w", err)
-	}
-	return session, nil
-}
-
 func checkAuthenticatedWithManager() (*auth.Session, *auth.SessionManager, error) {
 	sessionMgr, err := auth.NewSessionManager()
 	if err != nil {
diff --git a/internal/api/client.go b/internal/api/client.go
index cb446cb..4afe177 100644
--- a/internal/api/client.go
+++ b/internal/api/client.go
@@ -102,11 +102,7 @@ func (c *ProtonMailClient) DoWithContext(ctx context.Context, req *http.Request)
 
 	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.getAuthHeader()))
 	req.Header.Set("Accept", "application/json")
-
-	// Check if request has its own context
-	if ctx != context.Background() {
-		req = req.WithContext(ctx)
-	}
+	req = req.WithContext(ctx)
 
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
@@ -128,6 +124,13 @@ func (c *ProtonMailClient) DoWithContext(ctx context.Context, req *http.Request)
 			return resp, fmt.Errorf("401 received and refresh failed: %w", err)
 		}
 
+		// Update auth header with new access token from refreshed session
+		session, err := c.sessionRefresher.GetSession()
+		if err != nil {
+			return resp, fmt.Errorf("401 received, refresh succeeded but failed to get new session: %w", err)
+		}
+		c.SetAuthHeader(session.AccessToken)
+
 		// Retry the request with new token
 		// Clone the request to reset any body position
 		retryReq := req.Clone(ctx)
@@ -167,13 +170,4 @@ func (e *APIError) Error() string {
 	return fmt.Sprintf("API error %d: %s", e.HTTPStatus, e.Message)
 }
 
-// Helper function to create a request with context
-func NewRequestWithContext(ctx context.Context, method, url string, body io.Reader) (*http.Request, error) {
-	req, err := http.NewRequestWithContext(ctx, method, url, body)
-	if err != nil {
-		return nil, err
-	}
-	req.Header.Set("Content-Type", "application/json")
-	req.Header.Set("Accept", "application/json")
-	return req, nil
-}
+