Go AI agent API key management: vault keys via context and http.Client

The Go AI agent tool backend pattern

The most common Go pattern for AI agent tool backends initializes the vendor client once and uses it everywhere:

package main

import (
	"context"
	"encoding/json"
	"log"
	"net/http"
	"os"

	stripe "github.com/stripe/stripe-go/v76"
	"github.com/stripe/stripe-go/v76/paymentintent"
)

func init() {
	// Package-level init — shared across ALL goroutines
	stripe.Key = os.Getenv("STRIPE_SECRET_KEY")
}

func chargeHandler(w http.ResponseWriter, r *http.Request) {
	// This uses the package-level stripe.Key for every request
	params := &stripe.PaymentIntentParams{
		Amount:   stripe.Int64(5000),
		Currency: stripe.String(string(stripe.CurrencyUSD)),
	}
	pi, err := paymentintent.New(params)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	json.NewEncoder(w).Encode(pi)
}

The stripe-go library's package-level stripe.Key is a global variable shared across all goroutines. Even if you initialize a per-handler stripe.Client, you typically pass the same key to all instances. Two agents running simultaneously in the same process share credentials at the Stripe API level — Stripe cannot distinguish their spend.

Context-propagated vault keys in Go

Go's context.Context is designed for exactly this problem: passing request-scoped values down the call stack without changing function signatures. Define a typed key to avoid collisions:

package vault

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
)

// Unexported type prevents key collisions from other packages
type contextKey struct{}

type VaultKey struct {
	ID    string
	Token string
}

// Issue creates a new vault key and stores it in the returned context
func Issue(ctx context.Context, label, vendor string, dailyCapUSD int, allowedEndpoints []string, ttl string) (context.Context, error) {
	body, _ := json.Marshal(map[string]any{
		"label":             label,
		"vendor":            vendor,
		"daily_usd_cap":     dailyCapUSD,
		"allowed_endpoints": allowedEndpoints,
		"expires_in":        ttl,
	})

	req, _ := http.NewRequestWithContext(ctx, "POST", "https://api.keybrake.com/v1/keys", bytes.NewReader(body))
	req.Header.Set("Authorization", "Bearer "+os.Getenv("KEYBRAKE_TOKEN"))
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return ctx, fmt.Errorf("keybrake issue: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 201 {
		return ctx, fmt.Errorf("keybrake issue: status %d", resp.StatusCode)
	}

	var vk VaultKey
	if err := json.NewDecoder(resp.Body).Decode(&vk); err != nil {
		return ctx, err
	}

	return context.WithValue(ctx, contextKey{}, vk), nil
}

// FromContext retrieves the vault key from context
func FromContext(ctx context.Context) (VaultKey, bool) {
	vk, ok := ctx.Value(contextKey{}).(VaultKey)
	return vk, ok
}

// Revoke deletes the vault key; best called via defer
func Revoke(ctx context.Context, id string) {
	req, _ := http.NewRequestWithContext(ctx, "DELETE", "https://api.keybrake.com/v1/keys/"+id, nil)
	req.Header.Set("Authorization", "Bearer "+os.Getenv("KEYBRAKE_TOKEN"))
	http.DefaultClient.Do(req) // non-critical: TTL is the safety net
}

Middleware integration

Wrap your HTTP mux with middleware that issues the vault key before the handler and revokes it after:

package main

import (
	"log"
	"net/http"

	"yourapp/vault"
)

func VaultMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Only apply to agent tool routes
		if !strings.HasPrefix(r.URL.Path, "/tools/") {
			next.ServeHTTP(w, r)
			return
		}

		userID := r.Header.Get("X-User-ID")
		runID := r.Header.Get("X-Agent-Run-ID")
		label := fmt.Sprintf("go-%s-%s", userID, runID)

		ctx, err := vault.Issue(r.Context(), label, "stripe", 500,
			[]string{"/v1/payment_intents", "/v1/payment_intents/*"}, "5m")
		if err != nil {
			// Fail open with a warning; vault key is defense-in-depth
			log.Printf("WARN: vault key issuance failed: %v", err)
			next.ServeHTTP(w, r)
			return
		}

		vk, _ := vault.FromContext(ctx)
		defer vault.Revoke(context.Background(), vk.ID)

		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func main() {
	mux := http.NewServeMux()
	mux.HandleFunc("/tools/charge", chargeHandler)

	// Wrap mux with vault middleware
	http.ListenAndServe(":8080", VaultMiddleware(mux))
}

Tool function reading vault key from context

Tool functions read the vault key from context and call the Keybrake proxy instead of Stripe directly:

func chargeHandler(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	vk, ok := vault.FromContext(ctx)
	if !ok {
		// Fallback: vault key unavailable (issuance failed)
		http.Error(w, "vault key not available", 503)
		return
	}

	// Call Keybrake proxy with vault token instead of stripe.Key
	reqBody, _ := json.Marshal(map[string]any{
		"amount":   5000,
		"currency": "usd",
	})

	proxyReq, _ := http.NewRequestWithContext(ctx, "POST",
		"https://proxy.keybrake.com/stripe/v1/payment_intents", bytes.NewReader(reqBody))
	proxyReq.Header.Set("Authorization", "Bearer "+vk.Token)
	proxyReq.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(proxyReq)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer resp.Body.Close()

	// 429 with cap_exhausted = spend cap hit, not a transient error
	if resp.StatusCode == 429 {
		var errBody struct {
			Code string `json:"code"`
		}
		json.NewDecoder(resp.Body).Decode(&errBody)
		if errBody.Code == "cap_exhausted" {
			http.Error(w, "spend_cap_exceeded", http.StatusPaymentRequired)
			return
		}
	}

	w.Header().Set("Content-Type", "application/json")
	io.Copy(w, resp.Body)
}

Temporal and workflow patterns in Go

Go AI agent backends frequently use Temporal for durable workflows. Vault keys integrate at the activity level — each activity is the right granularity for a single vault key:

// Temporal activity: one vault key per activity execution
func StripeChargeActivity(ctx context.Context, params ChargeParams) (string, error) {
	info := activity.GetInfo(ctx)
	label := fmt.Sprintf("temporal-%s-%s", info.WorkflowExecution.ID, info.ActivityID)

	ctx, err := vault.Issue(ctx, label, "stripe", params.MaxSpendUSD,
		[]string{"/v1/payment_intents"}, "10m")
	if err != nil {
		return "", fmt.Errorf("vault key: %w", err)
	}

	vk, _ := vault.FromContext(ctx)
	defer vault.Revoke(context.Background(), vk.ID)

	// Use vk.Token to call proxy.keybrake.com/stripe/...
	return createPaymentIntent(ctx, vk.Token, params.Amount)
}

For Temporal workflows that span multiple activities, don't share a single vault key across the workflow — issue one per activity. This matches Temporal's activity retry model: if an activity retries, the revoked key from the previous attempt doesn't block the retry.

Get early access

Related questions

Further reading

• Temporal AI agent API key — vault keys in Temporal workflow and activity patterns, including retry-safe scoping and the workflow-vs-activity key granularity decision.
• AI agent API key lifecycle — issuance, enforcement, expiration, and revocation phases mapped to Go's context cancellation and goroutine lifecycle.
• AI agent idempotency — idempotency key patterns for Go HTTP handlers calling Stripe, including Temporal saga compensation patterns.
• AI agent error handling — handling spend cap 429s, retry logic with exponential backoff, and error propagation patterns for Go agent backends.
• Multi-tenant isolation — per-tenant vault key issuance patterns for Go SaaS backends serving multiple organizations.