selfheal-lab-apps/ai/internal/anthropic/federation_test.go

package anthropic

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai"
)

// Federation: the identity token from the (fake) metadata server is exchanged once at /v1/oauth/token with the
// jwt-bearer grant and the four ids; /v1/messages gets "Authorization: Bearer" and no x-api-key; the cached token is
// reused; a 401 triggers exactly one re-exchange; an expired cache re-exchanges before the next call.
func TestFederationExchangeAndReuse(t *testing.T) {
	var idCalls, exchanges, msgCalls int32
	var fail401 atomic.Bool
	md := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Metadata-Flavor") != "Google" || !strings.Contains(r.URL.RawQuery, "audience=https://api.anthropic.com") || !strings.Contains(r.URL.RawQuery, "format=full") {
			t.Errorf("bad metadata request: %s %v", r.URL.String(), r.Header)
		}
		atomic.AddInt32(&idCalls, 1)
		w.Write([]byte("eyJhbGciOiJSUzI1NiJ9.eyJpc3MiOiJodHRwczovL2FjY291bnRzLmdvb2dsZS5jb20ifQ.sig"))
	}))
	defer md.Close()
	api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/v1/oauth/token":
			var b map[string]string
			_ = json.NewDecoder(r.Body).Decode(&b)
			if b["grant_type"] != "urn:ietf:params:oauth:grant-type:jwt-bearer" || !strings.HasPrefix(b["assertion"], "eyJ") || b["federation_rule_id"] != "fdrl_1" || b["organization_id"] != "org-1" || b["service_account_id"] != "svac_1" || b["workspace_id"] != "wrkspc_1" {
				t.Errorf("bad exchange body: %v", b)
			}
			n := atomic.AddInt32(&exchanges, 1)
			w.Write([]byte(`{"access_token":"sk-ant-oat01-tok` + string(rune('0'+n)) + `","token_type":"Bearer","expires_in":600}`))
		case "/v1/messages":
			atomic.AddInt32(&msgCalls, 1)
			if r.Header.Get("x-api-key") != "" || !strings.HasPrefix(r.Header.Get("Authorization"), "Bearer sk-ant-oat01-tok") {
				t.Errorf("bad auth headers: %v", r.Header)
			}
			if fail401.Swap(false) {
				w.WriteHeader(401)
				w.Write([]byte(`{"type":"error","error":{"type":"authentication_error","message":"Authentication failed"}}`))
				return
			}
			w.Write([]byte(`{"id":"msg_f","type":"message","model":"claude-sonnet-5-5","stop_reason":"end_turn","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":3,"output_tokens":1}}`))
		default:
			t.Errorf("unexpected path %s", r.URL.Path)
		}
	}))
	defer api.Close()
	c := New(api.URL, func() (string, error) { return "", nil }) // no key file: federation must carry the call
	fed := &Federation{RuleID: "fdrl_1", OrgID: "org-1", ServiceAccountID: "svac_1", WorkspaceID: "wrkspc_1", IdentityTokenURL: md.URL + "/identity?audience=https://api.anthropic.com&format=full"}
	c.SetFederation(fed)
	if c.AuthMode() != "federation" {
		t.Fatalf("auth mode %s", c.AuthMode())
	}
	req := openai.Request{Model: "claude-sonnet-5-5", Input: []openai.Item{{Type: "message", Role: "user", Content: "hi"}}}
	for i := 0; i < 2; i++ {
		if _, err := c.CreateResponse(context.Background(), req); err != nil {
			t.Fatal(err)
		}
	}
	if atomic.LoadInt32(&exchanges) != 1 || atomic.LoadInt32(&idCalls) != 1 || atomic.LoadInt32(&msgCalls) != 2 {
		t.Fatalf("expected one exchange for two calls, got exchanges=%d id=%d msgs=%d", exchanges, idCalls, msgCalls)
	}
	// 401 -> one re-exchange and a successful retry
	fail401.Store(true)
	if _, err := c.CreateResponse(context.Background(), req); err != nil {
		t.Fatal(err)
	}
	if atomic.LoadInt32(&exchanges) != 2 || atomic.LoadInt32(&msgCalls) != 4 {
		t.Fatalf("401 must re-exchange once: exchanges=%d msgs=%d", exchanges, msgCalls)
	}
	// cached token close to expiry -> advisory refresh before the next call
	fed.mu.Lock()
	fed.expiry = time.Now().Add(60 * time.Second)
	fed.mu.Unlock()
	if _, err := c.CreateResponse(context.Background(), req); err != nil {
		t.Fatal(err)
	}
	if atomic.LoadInt32(&exchanges) != 3 {
		t.Fatalf("near-expiry token must be refreshed: exchanges=%d", exchanges)
	}
	// probe reports a successful exchange
	p := fed.Probe(context.Background(), api.URL)
	if p["ok"] != true || p["expires_in_s"].(int) < 500 || !strings.HasPrefix(p["token_prefix"].(string), "sk-ant-oat01") {
		t.Fatalf("probe: %v", p)
	}
}

func TestFederationFailuresAndPrecedence(t *testing.T) {
	// incomplete ids are ignored: the key path stays active
	c := New("http://127.0.0.1:1", func() (string, error) { return "k", nil })
	c.SetFederation(&Federation{RuleID: "fdrl_1"})
	if c.AuthMode() != "api_key" || c.Federation() != nil {
		t.Fatalf("incomplete federation must not activate: %s", c.AuthMode())
	}
	// exchange denied (opaque 401) -> clear error, no message call
	api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/v1/oauth/token" {
			w.WriteHeader(401)
			w.Write([]byte(`{"type":"error","error":{"type":"authentication_error","message":"Authentication failed"}}`))
			return
		}
		t.Errorf("messages must not be called when the exchange fails")
	}))
	defer api.Close()
	md := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Write([]byte("a.b.c")) }))
	defer md.Close()
	c2 := New(api.URL, func() (string, error) { return "", nil })
	fed := &Federation{RuleID: "fdrl_1", OrgID: "org-1", ServiceAccountID: "svac_1", IdentityTokenURL: md.URL}
	c2.SetFederation(fed)
	_, err := c2.CreateResponse(context.Background(), openai.Request{Model: "claude-haiku-5-5", Input: []openai.Item{{Type: "message", Role: "user", Content: "x"}}})
	if err == nil || !strings.Contains(err.Error(), "token exchange: status 401") {
		t.Fatalf("expected exchange failure, got %v", err)
	}
	if p := fed.Probe(context.Background(), api.URL); p["ok"] != false || !strings.Contains(p["error"].(string), "status 401") {
		t.Fatalf("probe must report the denial: %v", p)
	}
	// metadata server unreachable
	fed2 := &Federation{RuleID: "fdrl_1", OrgID: "org-1", ServiceAccountID: "svac_1", IdentityTokenURL: "http://127.0.0.1:1/identity", HTTP: &http.Client{Timeout: 2 * time.Second}}
	if _, err := fed2.AccessToken(context.Background(), api.URL); err == nil || !strings.Contains(err.Error(), "identity token") {
		t.Fatalf("expected identity token error, got %v", err)
	}
}