feat(backend): add kiro account support

2026-04-30 14:02:05 +08:00
parent 9d801595c9
commit 05bc424c9a
60 changed files with 11916 additions and 38 deletions
@@ -0,0 +1,208 @@
+//go:build unit
+
+package service
+
+import (
+	"context"
+	"net/http"
+	"strings"
+	"testing"
+
+	kiropkg "github.com/Wei-Shaw/sub2api/internal/pkg/kiro"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBuildKiroAccountKeyIgnoresAccessToken(t *testing.T) {
+	accountA := &Account{
+		ID: 99,
+		Credentials: map[string]any{
+			"access_token": "token-a",
+		},
+	}
+	accountB := &Account{
+		ID: 99,
+		Credentials: map[string]any{
+			"access_token": "token-b",
+		},
+	}
+
+	require.Equal(t, buildKiroAccountKey(accountA), buildKiroAccountKey(accountB))
+}
+
+func TestBuildKiroMachineIDPrefersExplicitCredential(t *testing.T) {
+	account := &Account{
+		ID:       101,
+		Platform: PlatformKiro,
+		Type:     AccountTypeOAuth,
+		Credentials: map[string]any{
+			"machineId":     "2582956e-cc88-4669-b546-07adbffcb894",
+			"refresh_token": "refresh-token",
+		},
+	}
+
+	require.Equal(t, "2582956ecc884669b54607adbffcb8942582956ecc884669b54607adbffcb894", buildKiroMachineID(account))
+}
+
+func TestBuildKiroMachineIDDerivesFromRefreshToken(t *testing.T) {
+	account := &Account{
+		ID:       102,
+		Platform: PlatformKiro,
+		Type:     AccountTypeOAuth,
+		Credentials: map[string]any{
+			"refresh_token": "refresh-token",
+		},
+	}
+
+	require.Equal(t, kiropkg.BuildMachineID("refresh-token", "", "account:102"), buildKiroMachineID(account))
+}
+
+func TestBuildKiroMachineIDDerivesFromAPIKeyAccount(t *testing.T) {
+	account := &Account{
+		ID:       103,
+		Platform: PlatformKiro,
+		Type:     AccountTypeAPIKey,
+		Credentials: map[string]any{
+			"kiroApiKey": "kiro-api-key",
+		},
+	}
+
+	require.Equal(t, kiropkg.BuildMachineID("", "kiro-api-key", "account:103"), buildKiroMachineID(account))
+}
+
+func TestNewKiroJSONRequestAddsConditionalHeaders(t *testing.T) {
+	account := &Account{
+		Credentials: map[string]any{
+			"auth_method": "external_idp",
+			"provider":    "Internal",
+			"profile_arn": "arn:aws:codewhisperer:us-east-1:123456789012:profile/HEADER",
+		},
+	}
+
+	req, err := newKiroJSONRequest(
+		context.Background(),
+		"https://q.us-east-1.amazonaws.com/generateAssistantResponse",
+		[]byte(`{"ok":true}`),
+		"access-token",
+		"account-key",
+		buildKiroMachineID(account),
+		"",
+		account,
+	)
+	require.NoError(t, err)
+	require.Equal(t, "EXTERNAL_IDP", req.Header.Get("TokenType"))
+	require.Equal(t, "true", req.Header.Get("redirect-for-internal"))
+	require.Equal(t, "arn:aws:codewhisperer:us-east-1:123456789012:profile/HEADER", req.Header.Get("x-amzn-kiro-profile-arn"))
+	require.Equal(t, "vibe", req.Header.Get("x-amzn-kiro-agent-mode"))
+	require.Equal(t, "true", req.Header.Get("x-amzn-codewhisperer-optout"))
+	require.Contains(t, req.Header.Get("User-Agent"), "aws-sdk-js/1.0.34")
+	require.Contains(t, req.Header.Get("User-Agent"), "md/nodejs#22.22.0")
+	require.Contains(t, req.Header.Get("User-Agent"), buildKiroMachineID(account))
+	require.Contains(t, req.Header.Get("X-Amz-User-Agent"), buildKiroMachineID(account))
+	require.True(t, strings.Contains(req.Header.Get("User-Agent"), "api/codewhispererstreaming#1.0.34"))
+	require.Empty(t, req.Header.Get("Anthropic-Beta"))
+}
+
+func TestIsKiroInvalidModelIDBodyRecognizesKnownForms(t *testing.T) {
+	tests := []string{
+		`{"message":"Invalid model ID. Please select a different model to continue.","reason":"INVALID_MODEL_ID"}`,
+		`{"message":"Invalid model. Please select a different model to continue."}`,
+		`API Error: 400 {"error":{"message":"Invalid model. Please select a different model to continue.","type":"upstream_error"}}`,
+	}
+
+	for _, body := range tests {
+		require.True(t, isKiroInvalidModelIDBody([]byte(body)), body)
+	}
+}
+
+func TestBuildKiroPayloadForAccountPropagatesThinkingHeaders(t *testing.T) {
+	account := &Account{
+		ID:       7,
+		Platform: PlatformKiro,
+		Type:     AccountTypeOAuth,
+		Credentials: map[string]any{
+			"profile_arn": "arn:aws:codewhisperer:us-east-1:123456789012:profile/test",
+		},
+	}
+	body := []byte(`{
+		"model":"claude-sonnet-4-6",
+		"messages":[{"role":"user","content":"hello"}]
+	}`)
+	headers := http.Header{}
+	headers.Set("Anthropic-Beta", "interleaved-thinking-2025-05-14")
+
+	payload, err := buildKiroPayloadForAccount(
+		context.Background(),
+		account,
+		body,
+		"claude-sonnet-4.6",
+		"kiro-access-token",
+		"claude-sonnet-4-6",
+		headers,
+	)
+	require.NoError(t, err)
+	require.NotContains(t, string(payload), "CHUNKED WRITE PROTOCOL")
+	require.Contains(t, string(payload), "\\u003cthinking_mode\\u003eenabled\\u003c/thinking_mode\\u003e")
+}
+
+func TestKiroAPIRegionPrefersAPIRegionOverProfileARN(t *testing.T) {
+	account := &Account{
+		Credentials: map[string]any{
+			"api_region":  "eu-west-1",
+			"profile_arn": "arn:aws:codewhisperer:us-west-2:123456789012:profile/test",
+			"region":      "ap-northeast-1",
+		},
+	}
+
+	require.Equal(t, "eu-west-1", kiroAPIRegion(account))
+}
+
+func TestKiroAPIRegionIgnoresProfileARNRegionFallback(t *testing.T) {
+	account := &Account{
+		Credentials: map[string]any{
+			"profile_arn": "arn:aws:codewhisperer:us-west-2:123456789012:profile/test",
+		},
+	}
+
+	require.Equal(t, kiroDefaultRegion, kiroAPIRegion(account))
+}
+
+func TestKiroAPIRegionIgnoresOIDCRegionFallback(t *testing.T) {
+	account := &Account{
+		Credentials: map[string]any{
+			"region": "ap-northeast-2",
+		},
+	}
+
+	require.Equal(t, kiroDefaultRegion, kiroAPIRegion(account))
+}
+
+func TestBuildKiroEndpointsUsesOnlyAmazonQEndpoint(t *testing.T) {
+	account := &Account{
+		Credentials: map[string]any{
+			"api_region":         "us-west-2",
+			"preferred_endpoint": "cw",
+		},
+	}
+
+	endpoints := buildKiroEndpoints(account)
+	require.Len(t, endpoints, 1)
+	require.Equal(t, "AmazonQ", endpoints[0].Name)
+	require.Equal(t, "q.us-west-2.amazonaws.com/generateAssistantResponse", endpoints[0].URL[8:])
+	require.Empty(t, endpoints[0].AmzTarget)
+}
+
+func TestBuildKiroEndpointsIgnoresPreferredEndpoint(t *testing.T) {
+	for _, preferred := range []string{"codewhisperer", "cw", "unknown"} {
+		account := &Account{
+			Credentials: map[string]any{
+				"api_region":         "us-west-2",
+				"preferred_endpoint": preferred,
+			},
+		}
+
+		endpoints := buildKiroEndpoints(account)
+		require.Len(t, endpoints, 1)
+		require.Equal(t, "AmazonQ", endpoints[0].Name)
+		require.Equal(t, "q.us-west-2.amazonaws.com/generateAssistantResponse", endpoints[0].URL[8:])
+	}
+}