perf(vault): cached credentials and safe singleflight refresh

README.md

@@ -305,13 +305,47 @@ AgentGate is deterministic. Agents never see tokens, calls go only to configured
 6. **OAuth state encrypted** with AES-256-GCM and expires after 10 minutes.
 7. **OAuth refresh.** Configured OAuth providers refresh tokens expiring within 5 minutes before dispatch. A failed refresh returns `token_expired` without calling the SaaS API.
 
+### Vault cache and refresh
+
+Vault reads cache decrypted access and refresh tokens in process memory.
+The default TTL is `60s`, with an LRU limit of 10,000 user-service pairs.
+An entry expires earlier when its token reaches 30 seconds before expiry.
+Tokens without an expiry use the TTL alone.
+Tokens already within that 30-second window aren't cached.
+
+Token strings cannot be reliably zeroed. Keep the cache bounded.
+Delete drops cached references. Go may retain string bytes until garbage collection.
+Expired entries aren't reused. Lookup or eviction removes their cached references.
+Put writes SQLite before caching the new token. Failed writes invalidate the local entry.
+All writes must pass through the decorated store.
+
+Concurrent refreshes share one result per store and user-service pair within a process.
+Current waiters share refresh failures. Later calls retry.
+Canceling a waiter doesn't cancel the shared refresh.
+Each shared refresh has a separate `15s` deadline. OAuth HTTP exchanges also have a `15s` client timeout.
+Refresh functions receive that context and must honor cancellation.
+Successful exchanges still persist their results after the refresh deadline. That deadline doesn't cancel persistence retries.
+Shared refresh panics return a fixed error without panic values. The flight is cleared so later requests can retry.
+Timeouts and canceled waits return HTTP `504` with receipt error code `vault_refresh_timeout`, not `token_expired`.
+Refresh writes compare the stored refresh token inside a SQLite transaction.
+Deleted credentials stay deleted. A concurrent re-link keeps its replacement token and scopes.
+On a comparison mismatch, refresh returns the current credential only when it remains valid.
+Rotated refresh tokens get up to 3 conditional write attempts.
+If every write fails, the current request still receives the new access token.
+The failure is logged without token values. Later requests may require re-linking because the rotation wasn't saved.
+Replicas maintain separate caches and refresh independently.
+Rotating refresh tokens still need cross-process coordination when replicas share credentials.
+External writes or deletes can leave stale local credentials until cache expiry.
+Set `AGENTGATE_VAULT_CACHE_TTL=0` to disable plaintext caching. Refresh sharing remains enabled.
+
 ## Configuration
 
 ### Environment Variables
 
 | Variable | Description | Required |
 |----------|-------------|----------|
 | `AGENTGATE_VAULT_KEY` | 32-byte vault key. Also derives receipt encryption keys and the API-key HMAC pepper. | Yes |
+| `AGENTGATE_VAULT_CACHE_TTL` | Decrypted-token cache duration. Default `60s`; `0` disables caching. Negative or invalid durations stop startup. | No |
 | `AGENTGATE_ADMIN_SECRET` | Secret for admin API access | Yes |
 | `AGENTGATE_PUBLIC_URL` | Base URL used to build the OAuth callback (default `http://localhost:8080`) | No |
 | `AGENTGATE_API_KEY_CACHE_TTL` | Verified-key cache duration. Default `30s`; `0` disables caching. | No |

cmd/agentgw/main.go

@@ -96,11 +96,21 @@ func main() {
 		os.Exit(1)
 	}
 
-	vaultStore, err := vault.NewSQLiteStore(database, masterKey)
+	sqliteVault, err := vault.NewSQLiteStore(database, masterKey)
 	if err != nil {
 		logger.Error("init vault", "error", err)
 		os.Exit(1)
 	}
+	vaultCacheTTL, err := loadVaultCacheTTL()
+	if err != nil {
+		logger.Error("load vault cache TTL", "error", err)
+		os.Exit(1)
+	}
+	vaultStore, err := vault.NewCachedStore(sqliteVault, vault.WithCacheTTL(vaultCacheTTL))
+	if err != nil {
+		logger.Error("init vault cache", "error", err)
+		os.Exit(1)
+	}
 
 	pepper, err := auth.DeriveKeyPepper(masterKey)
 	if err != nil {
@@ -275,6 +285,18 @@ func loadAPIKeyCacheTTL() (time.Duration, error) {
 	return ttl, nil
 }
 
+func loadVaultCacheTTL() (time.Duration, error) {
+	value, set := os.LookupEnv("AGENTGATE_VAULT_CACHE_TTL")
+	if !set {
+		return vault.DefaultCacheTTL, nil
+	}
+	ttl, err := time.ParseDuration(value)
+	if err != nil || ttl < 0 {
+		return 0, fmt.Errorf("AGENTGATE_VAULT_CACHE_TTL must be a non-negative duration")
+	}
+	return ttl, nil
+}
+
 // loadMasterKey reads the 32-byte master secret for the vault.
 // It also derives the signer's encryption key and the API-key HMAC pepper.
 func loadMasterKey() ([]byte, error) {

cmd/agentgw/vault_cache_ttl_test.go

@@ -0,0 +1,41 @@
+package main
+
+import (
+	"os"
+	"testing"
+	"time"
+)
+
+func TestLoadVaultCacheTTL(t *testing.T) {
+	for _, tc := range []struct {
+		value string
+		want  time.Duration
+		fail  bool
+	}{
+		{"60s", time.Minute, false},
+		{"5m", 5 * time.Minute, false},
+		{"0", 0, false},
+		{"0s", 0, false},
+		{"-1s", 0, true},
+		{"invalid", 0, true},
+		{"", 0, true},
+	} {
+		t.Run(tc.value, func(t *testing.T) {
+			t.Setenv("AGENTGATE_VAULT_CACHE_TTL", tc.value)
+			got, err := loadVaultCacheTTL()
+			if (err != nil) != tc.fail || got != tc.want {
+				t.Fatalf("TTL = %v, err = %v", got, err)
+			}
+		})
+	}
+	t.Run("unset", func(t *testing.T) {
+		t.Setenv("AGENTGATE_VAULT_CACHE_TTL", "")
+		if err := os.Unsetenv("AGENTGATE_VAULT_CACHE_TTL"); err != nil {
+			t.Fatal(err)
+		}
+		got, err := loadVaultCacheTTL()
+		if err != nil || got != time.Minute {
+			t.Fatalf("default TTL = %v, err = %v", got, err)
+		}
+	})
+}

internal/gateway/gateway.go

@@ -457,8 +457,16 @@ func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome {
 		}
 	}
 	if refresher := s.cfg.Refreshers[att.req.Service]; refresher != nil {
-		tok, err = vault.GetOrRefresh(s.cfg.Vault, att.req.OnBehalfOf, att.req.Service, refresher)
+		tok, err = vault.GetOrRefresh(ctx, s.cfg.Vault, att.req.OnBehalfOf, att.req.Service, refresher)
 		if err != nil {
+			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+				return &outcome{
+					status:         http.StatusGatewayTimeout,
+					body:           ErrorResponse{Error: "vault refresh timed out or was canceled", Code: "vault_refresh_timeout"},
+					policyDecision: "allow",
+					errorCode:      "vault_refresh_timeout",
+				}
+			}
 			return &outcome{
 				status:         http.StatusForbidden,
 				body:           ErrorResponse{Error: "token expired — user must re-authenticate", Code: "token_expired"},

internal/gateway/gateway_test.go

@@ -468,7 +468,7 @@ services:
 		}},
 		Receipts: &fakeReceipts{},
 		Refreshers: map[string]vault.RefreshFunc{
-			"github": func(refreshToken string) (string, string, time.Duration, error) {
+			"github": func(_ context.Context, refreshToken string) (string, string, time.Duration, error) {
 				if refreshToken != "refresh-token" {
 					t.Errorf("refresh token = %q, want stored refresh token", refreshToken)
 				}
@@ -508,7 +508,7 @@ func TestAct_RefreshFailureBlocksDispatchAndIsReceipted(t *testing.T) {
 		ExpiresAt:    time.Now().Add(-time.Minute),
 	})
 	srv.cfg.Refreshers = map[string]vault.RefreshFunc{
-		"github": func(string) (string, string, time.Duration, error) {
+		"github": func(context.Context, string) (string, string, time.Duration, error) {
 			return "", "", 0, errors.New("provider rejected refresh token")
 		},
 	}

internal/gateway/path_test.go

@@ -1,6 +1,7 @@
 package gateway
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/fs"
@@ -333,7 +334,7 @@ func TestActInvalidPathParamsBeforeVault(t *testing.T) {
 			srv.cfg.Vault = counted
 			var refreshCalls atomic.Int32
 			srv.cfg.Refreshers = map[string]vault.RefreshFunc{
-				"github": func(string) (string, string, time.Duration, error) {
+				"github": func(context.Context, string) (string, string, time.Duration, error) {
 					refreshCalls.Add(1)
 					return "new-token", "new-refresh", time.Hour, nil
 				},

internal/gateway/refresh_cancellation_test.go

@@ -0,0 +1,78 @@
+package gateway
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/Clawdlinux/agentgate/internal/vault"
+)
+
+func TestActRefreshCancellationIsNotTokenExpired(t *testing.T) {
+	for _, deadline := range []bool{false, true} {
+		t.Run(map[bool]string{false: "canceled", true: "deadline"}[deadline], func(t *testing.T) {
+			srv, store, receipts := testSetup(t)
+			_ = store.Put("user-1", "github", vault.Token{
+				AccessToken: "old", RefreshToken: "refresh", ExpiresAt: time.Now().Add(-time.Minute),
+			})
+			started, release := make(chan struct{}), make(chan struct{})
+			refresh := func(context.Context, string) (string, string, time.Duration, error) {
+				close(started)
+				<-release
+				return "new", "rotated", time.Hour, nil
+			}
+			srv.cfg.Refreshers = map[string]vault.RefreshFunc{"github": refresh}
+			ctx, cancel := context.WithCancel(context.Background())
+			if deadline {
+				cancel()
+				ctx, cancel = context.WithTimeout(context.Background(), 100*time.Millisecond)
+			}
+			defer cancel()
+			req := httptest.NewRequest("POST", "/v1/act", strings.NewReader(
+				`{"service":"github","action":"list_repos","on_behalf_of":"user-1"}`)).WithContext(ctx)
+			req.Header.Set("Authorization", "Bearer test-key")
+			w := httptest.NewRecorder()
+			done := make(chan struct{})
+			go func() {
+				defer close(done)
+				srv.ServeHTTP(w, req)
+			}()
+			select {
+			case <-started:
+			case <-time.After(time.Second):
+				close(release)
+				t.Fatal("refresh never started")
+			}
+			if !deadline {
+				cancel()
+			}
+			select {
+			case <-done:
+			case <-time.After(time.Second):
+				close(release)
+				t.Fatal("gateway didn't honor waiter cancellation")
+			}
+			close(release)
+			if _, err := vault.GetOrRefresh(context.Background(), store, "user-1", "github", refresh); err != nil {
+				t.Fatal(err)
+			}
+			var body ErrorResponse
+			if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
+				t.Fatal(err)
+			}
+			if w.Code != http.StatusGatewayTimeout || body.Code != "vault_refresh_timeout" {
+				t.Fatalf("status = %d, body = %s", w.Code, w.Body)
+			}
+			if body.Error != "vault refresh timed out or was canceled" {
+				t.Fatalf("unexpected cancellation message: %q", body.Error)
+			}
+			if receipts.count() != 1 || receipts.last().Error != "vault_refresh_timeout" {
+				t.Fatalf("receipt = %+v", receipts.last())
+			}
+		})
+	}
+}

internal/oauth/handler.go

@@ -1,8 +1,10 @@
 package oauth
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
+	"html/template"
 	"io"
 	"log/slog"
 	"net/http"
@@ -33,6 +35,12 @@ type CallbackHandler struct {
 	logger       *slog.Logger
 }
 
+var linkedPage = template.Must(template.New("linked").Parse(`<!DOCTYPE html>
+<html><body>
+<h2>Account linked!</h2>
+<p>Your {{.}} account has been connected. You can close this tab.</p>
+</body></html>`))
+
 // NewCallbackHandler creates a new OAuth callback handler.
 func NewCallbackHandler(providers map[string]*Provider, v vault.Store, encKey []byte, callbackBase string, logger *slog.Logger) *CallbackHandler {
 	if logger == nil {
@@ -118,7 +126,7 @@ func (h *CallbackHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Exchange code for tokens.
-	tok, err := h.exchangeCode(provider, code, service)
+	tok, err := h.exchangeCode(r.Context(), provider, code, service)
 	if err != nil {
 		h.logger.Error("oauth code exchange failed", "service", service, "error", err)
 		http.Error(w, "token exchange failed", http.StatusInternalServerError)
@@ -135,31 +143,29 @@ func (h *CallbackHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	h.logger.Info("oauth linked", "user", userID, "service", service)
 
 	w.Header().Set("Content-Type", "text/html")
-	fmt.Fprintf(w, `<!DOCTYPE html>
-<html><body>
-<h2>Account linked!</h2>
-<p>Your %s account has been connected. You can close this tab.</p>
-</body></html>`, service)
+	if err := linkedPage.Execute(w, service); err != nil {
+		h.logger.Error("render oauth callback", "error", err)
+	}
 }
 
-func (h *CallbackHandler) exchangeCode(provider *Provider, code, service string) (*vault.Token, error) {
+func (h *CallbackHandler) exchangeCode(ctx context.Context, provider *Provider, code, service string) (*vault.Token, error) {
 	data := url.Values{
 		"grant_type":    {"authorization_code"},
 		"code":          {code},
 		"client_id":     {provider.ClientID},
 		"client_secret": {provider.ClientSecret},
 		"redirect_uri":  {h.callbackBase + "/auth/callback/" + service},
 	}
-	tok, _, err := exchangeToken(provider, data)
+	tok, _, err := exchangeToken(ctx, provider, data)
 	return tok, err
 }
 
 // NewRefreshFunc returns the refresh-token exchange for one configured OAuth
 // provider. It returns an error when the provider omits an expiry because the
 // caller must not persist a token that will immediately be treated as expired.
 func NewRefreshFunc(provider *Provider) vault.RefreshFunc {
-	return func(refreshToken string) (string, string, time.Duration, error) {
-		tok, expiresIn, err := exchangeToken(provider, url.Values{
+	return func(ctx context.Context, refreshToken string) (string, string, time.Duration, error) {
+		tok, expiresIn, err := exchangeToken(ctx, provider, url.Values{
 			"grant_type":    {"refresh_token"},
 			"refresh_token": {refreshToken},
 			"client_id":     {provider.ClientID},
@@ -175,14 +181,24 @@ func NewRefreshFunc(provider *Provider) vault.RefreshFunc {
 	}
 }
 
-func exchangeToken(provider *Provider, data url.Values) (*vault.Token, time.Duration, error) {
-	resp, err := http.PostForm(provider.TokenURL, data)
+var tokenHTTPClient = &http.Client{Timeout: vault.DefaultRefreshTimeout}
+
+func exchangeToken(ctx context.Context, provider *Provider, data url.Values) (*vault.Token, time.Duration, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.TokenURL, strings.NewReader(data.Encode()))
+	if err != nil {
+		return nil, 0, fmt.Errorf("oauth: create exchange request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+	resp, err := tokenHTTPClient.Do(req)
 	if err != nil {
 		return nil, 0, fmt.Errorf("oauth: exchange: %w", err)
 	}
 	defer resp.Body.Close()
 
-	body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
+	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
+	if err != nil {
+		return nil, 0, fmt.Errorf("oauth: read token response: %w", err)
+	}
 
 	if resp.StatusCode != 200 {
 		return nil, 0, fmt.Errorf("oauth: exchange returned %d: %s", resp.StatusCode, string(body))

internal/oauth/linked_page_test.go

@@ -0,0 +1,17 @@
+package oauth
+
+import (
+	"bytes"
+	"strings"
+	"testing"
+)
+
+func TestLinkedPageEscapesService(t *testing.T) {
+	var body bytes.Buffer
+	if err := linkedPage.Execute(&body, "<script>alert(1)</script>"); err != nil {
+		t.Fatal(err)
+	}
+	if strings.Contains(body.String(), "<script>") || !strings.Contains(body.String(), "&lt;script&gt;") {
+		t.Fatal("callback page did not escape the service name")
+	}
+}