feat(bedrock): Ochre runtime MVP — Converse over self-host + Anthropic-direct (#520)

* feat(bedrock): Ochre runtime MVP — Converse over self-host + Anthropic-direct

Add the AWS-Bedrock-compatible gateway services (bedrock, bedrock-runtime)
following the EKS REST-JSON pattern: ListFoundationModels / GetFoundationModel
over a static connectivity-tiered catalog, and non-streaming Converse routed
per modelId to one of two provider backends — self-hosted vLLM (OpenAI wire)
and Anthropic-direct (Anthropic Messages wire).

Per-account provider API keys are a master-key-encrypted JetStream KV secret
with an optional platform default; the advertised catalog gates external
models on a resolvable credential. Self-host endpoints are pinned via config
(OCHRE_VLLM_ENDPOINTS) behind an EndpointResolver seam.

Adds gateway/bedrock/ (catalog, provider router, Anthropic + vLLM adapters,
credential store), awserrors REST-JSON codes, and the bedrock E2E suite
(list/get light; Converse gated on live backends). Daemon/NATS control plane
and dynamic endpoint lifecycle are deferred to Phase 3.

* test(bedrock): cover dispatch, router, credentials KV + endpoint config

Add unit coverage for the Ochre Phase 1 diff: Bedrock_Request /
BedrockRuntime_Request end-to-end dispatch, provider router, JSON
response/error writers, upstream-status mapping, the credential-store KV
round-trip, the Anthropic bound-provider path, and parseBedrockEndpoints.
Lifts PR diff coverage above the 60% gate.
This commit is contained in:
Josh Siviour
2026-07-13 21:40:34 +10:00
committed by GitHub
parent 1e1d2f3c6d
commit cee73671ff
28 changed files with 2383 additions and 41 deletions
+8
View File
@@ -300,6 +300,14 @@ suites:
# spinifex_daemon. docker/crane/skopeo data-plane subtests skip when the
# client binary or registry-host resolution is absent.
e2e-bedrock:
path: tests/e2e/bedrock
covers: [awsgw]
# Ochre (AWS-Bedrock) control plane: static catalog list/get served
# straight from awsgw (gateway/bedrock.go, gateway/bedrock/) — no NATS
# hop to spinifex_daemon. Converse (Anthropic-direct / self-host vLLM)
# subtests are gated on backend availability and skip by default.
e2e-ecs:
path: tests/e2e/ecs
covers: [awsgw, spinifex_daemon, vpcd, viperblockd, predastore]
+20
View File
@@ -499,6 +499,17 @@ var (
ErrorELBv2IncompatibleProtocols = "IncompatibleProtocols"
ErrorELBv2CertificateNotFound = "CertificateNotFound"
ErrorELBv2SSLPolicyNotFound = "SSLPolicyNotFound"
// Bedrock/bedrock-runtime error codes
ErrorValidationException = "ValidationException"
// ErrorResourceNotFoundException aliases ErrorEKSResourceNotFound: both
// services use the same wire code, so it is not a second ErrorLookup entry.
ErrorResourceNotFoundException = ErrorEKSResourceNotFound
ErrorAccessDeniedException = "AccessDeniedException"
ErrorThrottlingException = "ThrottlingException"
ErrorModelNotReadyException = "ModelNotReadyException"
ErrorServiceUnavailableException = "ServiceUnavailableException"
ErrorModelErrorException = "ModelErrorException"
)
// ValidErrorCode returns the error code if it exists in ErrorLookup,
@@ -1021,4 +1032,13 @@ var ErrorLookup = map[string]ErrorMessage{
ErrorELBv2IncompatibleProtocols: {HTTPCode: 400, Message: "The listener protocol is incompatible with the target group protocol."},
ErrorELBv2CertificateNotFound: {HTTPCode: 400, Message: "One or more specified certificates do not exist."},
ErrorELBv2SSLPolicyNotFound: {HTTPCode: 400, Message: "The specified SSL policy does not exist."},
// Bedrock/bedrock-runtime error codes. ResourceNotFoundException reuses the
// EKS entry (ErrorEKSResourceNotFound) — same wire code, same HTTP status.
ErrorValidationException: {HTTPCode: 400, Message: "The input fails to satisfy the constraints specified by the model or service."},
ErrorAccessDeniedException: {HTTPCode: 403, Message: "You do not have sufficient access to perform this action."},
ErrorThrottlingException: {HTTPCode: 429, Message: "The request was denied due to request throttling."},
ErrorModelNotReadyException: {HTTPCode: 429, Message: "The model specified in the request is not ready to serve inference requests."},
ErrorServiceUnavailableException: {HTTPCode: 503, Message: "The service isn't currently available. Try again later."},
ErrorModelErrorException: {HTTPCode: 424, Message: "The request failed because of an error while running the model."},
}
+1 -1
View File
@@ -13,7 +13,7 @@ func TestErrorLookup_Structure(t *testing.T) {
t.Fatalf("ErrorLookup unexpectedly small: %d entries", len(ErrorLookup))
}
validHTTP := map[int]bool{400: true, 403: true, 404: true, 409: true, 412: true, 413: true, 500: true, 501: true, 503: true}
validHTTP := map[int]bool{400: true, 403: true, 404: true, 409: true, 412: true, 413: true, 424: true, 429: true, 500: true, 501: true, 503: true}
for code, msg := range ErrorLookup {
if !validHTTP[msg.HTTPCode] {
t.Errorf("%s has invalid HTTPCode %d", code, msg.HTTPCode)
+128
View File
@@ -0,0 +1,128 @@
package gateway
import (
"context"
"errors"
"io"
"log/slog"
"net/http"
"net/url"
"regexp"
"github.com/aws/aws-sdk-go/service/bedrock"
"github.com/mulgadc/spinifex/spinifex/awserrors"
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
)
// bedrockRoute maps one HTTP method + path regex to an AWS action and handler.
type bedrockRoute struct {
method string
pattern *regexp.Regexp
action string
handler bedrockRouteHandler
}
// bedrockRouteHandler invokes a per-action bedrock (control-plane) gateway
// function. params holds the regex capture groups, PathUnescape'd. resolver
// is gw.bedrockResolver(): the configured credential store, or a no-op
// fallback.
type bedrockRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver) (any, error)
// bedrockRoutes is the dispatch table. More-specific paths must precede
// less-specific ones with the same prefix so the regex matcher picks the
// deeper route first.
var bedrockRoutes = []bedrockRoute{
{"GET", regexp.MustCompile(`^/foundation-models$`), "ListFoundationModels",
func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver) (any, error) {
return gateway_bedrock.ListFoundationModels(ctx, acct, resolver, new(bedrock.ListFoundationModelsInput))
}},
{"GET", regexp.MustCompile(`^/foundation-models/([^/]+)$`), "GetFoundationModel",
func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver) (any, error) {
return gateway_bedrock.GetFoundationModel(ctx, acct, p[0])
}},
}
// lookupBedrockAction matches method+path against bedrockRoutes, returning the
// action, path params, and handler, or ("", nil, nil, false) on no match.
// path must be r.URL.EscapedPath(): captured params are PathUnescape'd before
// returning, mirroring lookupEKSAction.
func lookupBedrockAction(method, path string) (string, []string, bedrockRouteHandler, bool) {
for _, route := range bedrockRoutes {
if route.method != method {
continue
}
m := route.pattern.FindStringSubmatch(path)
if m == nil {
continue
}
var params []string
if len(m) > 1 {
params = make([]string, 0, len(m)-1)
for _, raw := range m[1:] {
decoded, err := url.PathUnescape(raw)
if err != nil {
slog.Debug("bedrock: bad percent-encoding in path param", "param", raw, "err", err)
decoded = raw
}
params = append(params, decoded)
}
}
return route.action, params, route.handler, true
}
return "", nil, nil, false
}
// Bedrock_Request dispatches bedrock (control-plane) REST-JSON requests:
// resolves method+path to an action, reads the body, calls the handler, and
// serialises the output as JSON.
func (gw *GatewayConfig) Bedrock_Request(w http.ResponseWriter, r *http.Request) error {
action, params, handler, ok := lookupBedrockAction(r.Method, r.URL.EscapedPath())
if !ok {
slog.DebugContext(r.Context(), "bedrock: no route for request", "method", r.Method, "path", r.URL.Path)
return errors.New(awserrors.ErrorInvalidAction)
}
if err := gw.checkPolicy(r, "bedrock", action); err != nil {
return err
}
if gw.NATSConn == nil {
return errors.New(awserrors.ErrorServerInternal)
}
accountID, _ := r.Context().Value(ctxAccountID).(string)
if accountID == "" {
slog.ErrorContext(r.Context(), "Bedrock_Request: no account ID in auth context")
return errors.New(awserrors.ErrorServerInternal)
}
body, err := io.ReadAll(r.Body)
if err != nil {
slog.ErrorContext(r.Context(), "Bedrock_Request: failed to read body", "err", err)
return errors.New(awserrors.ErrorInvalidParameterValue)
}
output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver())
if err != nil {
return err
}
gateway_bedrock.WriteJSONResponse(w, output)
return nil
}
// bedrockResolver returns gw.BedrockCredentials as a CredentialResolver, or
// the no-op fallback when no credential store is configured.
func (gw *GatewayConfig) bedrockResolver() gateway_bedrock.CredentialResolver {
if gw.BedrockCredentials != nil {
return gw.BedrockCredentials
}
return gateway_bedrock.NoopCredentialResolver
}
// bedrockEndpointResolver returns an EndpointResolver over the configured
// pinned self-host endpoints (gw.BedrockEndpoints). A nil/empty map resolves
// nothing, so self-host models return ModelNotReady until configured.
func (gw *GatewayConfig) bedrockEndpointResolver() gateway_bedrock.EndpointResolver {
return gateway_bedrock.NewStaticEndpointResolver(gw.BedrockEndpoints)
}
+251
View File
@@ -0,0 +1,251 @@
package gateway_bedrock
import (
"bytes"
"context"
"encoding/json"
"errors"
"io"
"log/slog"
"net/http"
"regexp"
"strings"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
const (
anthropicDefaultBaseURL = "https://api.anthropic.com"
anthropicMessagesPath = "/v1/messages"
anthropicAPIVersion = "2023-06-01"
anthropicDefaultMaxTokens = 4096
)
// anthropicModelIDPattern strips a "anthropic." vendor prefix and a trailing
// "-v<N>:<M>" Bedrock version suffix, e.g.
// "anthropic.claude-3-5-sonnet-20240620-v1:0" -> "claude-3-5-sonnet-20240620".
var anthropicModelIDPattern = regexp.MustCompile(`^anthropic\.(.+)-v\d+:\d+$`)
// anthropicModelID maps a Bedrock modelId to its Anthropic model name. IDs
// that don't match the expected shape pass through unchanged.
func anthropicModelID(bedrockModelID string) string {
if m := anthropicModelIDPattern.FindStringSubmatch(bedrockModelID); m != nil {
return m[1]
}
return bedrockModelID
}
// anthropicContentBlock is a single Messages API content block. Phase 1 is
// text-only, so only "text" blocks are produced or read.
type anthropicContentBlock struct {
Type string `json:"type"`
Text string `json:"text"`
}
type anthropicMessage struct {
Role string `json:"role"`
Content []anthropicContentBlock `json:"content"`
}
type anthropicRequest struct {
Model string `json:"model"`
MaxTokens int64 `json:"max_tokens"`
System string `json:"system,omitempty"`
Messages []anthropicMessage `json:"messages"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
StopSequences []string `json:"stop_sequences,omitempty"`
}
type anthropicUsage struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
}
type anthropicResponse struct {
Content []anthropicContentBlock `json:"content"`
StopReason string `json:"stop_reason"`
Usage anthropicUsage `json:"usage"`
}
// anthropicStopReasons maps Anthropic stop_reason values to Bedrock's
// StopReason enum. Unrecognised values (and tool/guardrail reasons Phase 1
// doesn't model) fall back to end_turn.
var anthropicStopReasons = map[string]string{
"end_turn": bedrockruntime.StopReasonEndTurn,
"max_tokens": bedrockruntime.StopReasonMaxTokens,
"stop_sequence": bedrockruntime.StopReasonStopSequence,
"tool_use": bedrockruntime.StopReasonToolUse,
}
func mapAnthropicStopReason(reason string) string {
if mapped, ok := anthropicStopReasons[reason]; ok {
return mapped
}
return bedrockruntime.StopReasonEndTurn
}
// anthropicProvider calls the Anthropic Messages API directly. httpClient
// and baseURL are injectable for tests.
type anthropicProvider struct {
httpClient *http.Client
baseURL string
}
func newAnthropicProviderClient() *anthropicProvider {
return &anthropicProvider{
httpClient: &http.Client{Timeout: providerHTTPTimeout},
baseURL: anthropicDefaultBaseURL,
}
}
// anthropicBoundProvider adapts anthropicProvider to Provider by baking in a
// resolved per-account (or platform-default) API key.
type anthropicBoundProvider struct {
inner *anthropicProvider
apiKey string
}
var _ Provider = (*anthropicBoundProvider)(nil)
func (b *anthropicBoundProvider) Converse(ctx context.Context, modelID string, input *bedrockruntime.ConverseInput) (*bedrockruntime.ConverseOutput, error) {
return b.inner.Converse(ctx, modelID, input, b.apiKey)
}
// newAnthropicProvider returns a Provider that calls the Anthropic Messages
// API with apiKey.
func newAnthropicProvider(apiKey string) Provider {
return &anthropicBoundProvider{inner: newAnthropicProviderClient(), apiKey: apiKey}
}
// Converse translates input to an Anthropic Messages request, calls the API
// with apiKey, and translates the response back to a Bedrock ConverseOutput.
func (p *anthropicProvider) Converse(ctx context.Context, modelID string, input *bedrockruntime.ConverseInput, apiKey string) (*bedrockruntime.ConverseOutput, error) {
reqBody, err := json.Marshal(buildAnthropicRequest(modelID, input))
if err != nil {
slog.Error("anthropic: failed to marshal request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+anthropicMessagesPath, bytes.NewReader(reqBody))
if err != nil {
slog.Error("anthropic: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq.Header.Set("x-api-key", apiKey)
httpReq.Header.Set("anthropic-version", anthropicAPIVersion)
httpReq.Header.Set("content-type", "application/json")
start := time.Now()
resp, err := p.httpClient.Do(httpReq)
latency := time.Since(start)
if err != nil {
slog.Error("anthropic: request failed", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
slog.Error("anthropic: failed to read response body", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
slog.Error("anthropic: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
return nil, errors.New(mapUpstreamStatus(resp.StatusCode))
}
var ar anthropicResponse
if err := json.Unmarshal(respBody, &ar); err != nil {
slog.Error("anthropic: failed to parse response", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorModelErrorException)
}
return anthropicToConverseOutput(ar, latency), nil
}
// buildAnthropicRequest maps a Bedrock ConverseInput to an Anthropic Messages
// request. Non-text content blocks are dropped (Phase 1 is text-only).
func buildAnthropicRequest(modelID string, input *bedrockruntime.ConverseInput) anthropicRequest {
maxTokens := int64(anthropicDefaultMaxTokens)
var temperature, topP *float64
var stopSequences []string
if input.InferenceConfig != nil {
if input.InferenceConfig.MaxTokens != nil {
maxTokens = *input.InferenceConfig.MaxTokens
}
temperature = input.InferenceConfig.Temperature
topP = input.InferenceConfig.TopP
stopSequences = aws.StringValueSlice(input.InferenceConfig.StopSequences)
}
var systemParts []string
for _, s := range input.System {
if s != nil && s.Text != nil {
systemParts = append(systemParts, *s.Text)
}
}
messages := make([]anthropicMessage, 0, len(input.Messages))
for _, m := range input.Messages {
if m == nil {
continue
}
var blocks []anthropicContentBlock
for _, c := range m.Content {
// TODO: Phase 1 is text-only; non-text content blocks are dropped.
if c == nil || c.Text == nil {
continue
}
blocks = append(blocks, anthropicContentBlock{Type: "text", Text: *c.Text})
}
messages = append(messages, anthropicMessage{Role: aws.StringValue(m.Role), Content: blocks})
}
return anthropicRequest{
Model: anthropicModelID(modelID),
MaxTokens: maxTokens,
System: strings.Join(systemParts, "\n"),
Messages: messages,
Temperature: temperature,
TopP: topP,
StopSequences: stopSequences,
}
}
// anthropicToConverseOutput maps an Anthropic Messages response to a Bedrock
// ConverseOutput, measuring latency from the caller-supplied wall-clock delta.
func anthropicToConverseOutput(ar anthropicResponse, latency time.Duration) *bedrockruntime.ConverseOutput {
var content []*bedrockruntime.ContentBlock
for _, block := range ar.Content {
if block.Type != "text" {
continue
}
content = append(content, &bedrockruntime.ContentBlock{Text: aws.String(block.Text)})
}
inputTokens := ar.Usage.InputTokens
outputTokens := ar.Usage.OutputTokens
return &bedrockruntime.ConverseOutput{
Output: &bedrockruntime.ConverseOutput_{
Message: &bedrockruntime.Message{
Role: aws.String(bedrockruntime.ConversationRoleAssistant),
Content: content,
},
},
StopReason: aws.String(mapAnthropicStopReason(ar.StopReason)),
Usage: &bedrockruntime.TokenUsage{
InputTokens: aws.Int64(inputTokens),
OutputTokens: aws.Int64(outputTokens),
TotalTokens: aws.Int64(inputTokens + outputTokens),
},
Metrics: &bedrockruntime.ConverseMetrics{
LatencyMs: aws.Int64(latency.Milliseconds()),
},
}
}
+142
View File
@@ -0,0 +1,142 @@
package gateway_bedrock
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAnthropicProvider_Converse_MapsResponse(t *testing.T) {
var capturedBody anthropicRequest
var capturedAPIKey string
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
capturedAPIKey = r.Header.Get("x-api-key")
require.NoError(t, json.NewDecoder(r.Body).Decode(&capturedBody))
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{
"content": [{"type": "text", "text": "Hello there"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 12, "output_tokens": 5}
}`))
}))
defer ts.Close()
p := &anthropicProvider{httpClient: ts.Client(), baseURL: ts.URL}
input := &bedrockruntime.ConverseInput{
System: []*bedrockruntime.SystemContentBlock{{Text: aws.String("be nice")}},
Messages: []*bedrockruntime.Message{
{
Role: aws.String(bedrockruntime.ConversationRoleUser),
Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}},
},
},
InferenceConfig: &bedrockruntime.InferenceConfiguration{
MaxTokens: aws.Int64(256),
},
}
out, err := p.Converse(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input, "sk-test-key")
require.NoError(t, err)
assert.Equal(t, "sk-test-key", capturedAPIKey)
assert.Equal(t, "claude-3-5-sonnet-20240620", capturedBody.Model)
assert.Equal(t, int64(256), capturedBody.MaxTokens)
assert.Equal(t, "be nice", capturedBody.System)
require.Len(t, capturedBody.Messages, 1)
assert.Equal(t, "user", capturedBody.Messages[0].Role)
require.NotNil(t, out.Output.Message)
assert.Equal(t, "assistant", *out.Output.Message.Role)
require.Len(t, out.Output.Message.Content, 1)
assert.Equal(t, "Hello there", *out.Output.Message.Content[0].Text)
assert.Equal(t, int64(12), *out.Usage.InputTokens)
assert.Equal(t, int64(5), *out.Usage.OutputTokens)
assert.Equal(t, int64(17), *out.Usage.TotalTokens)
assert.Equal(t, bedrockruntime.StopReasonEndTurn, *out.StopReason)
require.NotNil(t, out.Metrics)
assert.GreaterOrEqual(t, *out.Metrics.LatencyMs, int64(0))
}
func TestAnthropicProvider_Converse_401ReturnsAccessDenied(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error": {"message": "invalid x-api-key"}}`))
}))
defer ts.Close()
p := &anthropicProvider{httpClient: ts.Client(), baseURL: ts.URL}
input := &bedrockruntime.ConverseInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}},
},
}
_, err := p.Converse(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input, "bad-key")
require.Error(t, err)
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
}
// TestAnthropicBoundProvider_Converse exercises newAnthropicProvider,
// newAnthropicProviderClient, and anthropicBoundProvider.Converse (the
// Provider used by Router for the "provider:anthropic" tier). The real
// Anthropic base URL is hardcoded by newAnthropicProvider, so — same package
// as production code — the test reaches into the unexported inner client to
// redirect it at an httptest stub instead of the network.
func TestAnthropicBoundProvider_Converse(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "sk-bound-test", r.Header.Get("x-api-key"))
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{
"content": [{"type": "text", "text": "bound response"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 3, "output_tokens": 2}
}`))
}))
defer ts.Close()
p := newAnthropicProvider("sk-bound-test")
bound, ok := p.(*anthropicBoundProvider)
require.True(t, ok)
require.NotNil(t, bound.inner)
bound.inner.httpClient = ts.Client()
bound.inner.baseURL = ts.URL
input := &bedrockruntime.ConverseInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}},
},
}
out, err := p.Converse(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input)
require.NoError(t, err)
require.NotNil(t, out.Output.Message)
assert.Equal(t, "bound response", *out.Output.Message.Content[0].Text)
}
func TestAnthropicModelID(t *testing.T) {
cases := []struct {
name, in, want string
}{
{"sonnet with version suffix", "anthropic.claude-3-5-sonnet-20240620-v1:0", "claude-3-5-sonnet-20240620"},
{"haiku with version suffix", "anthropic.claude-3-haiku-20240307-v1:0", "claude-3-haiku-20240307"},
{"non-matching id passes through", "meta.llama3-70b-instruct-v1:0", "meta.llama3-70b-instruct-v1:0"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, anthropicModelID(tc.in))
})
}
}
+154
View File
@@ -0,0 +1,154 @@
package gateway_bedrock
import (
"context"
"errors"
"fmt"
"strings"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrock"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
// Tier flags on catalogEntry.Provider. "self-host" models run on Spinifex GPU
// compute; "provider:<vendor>" models are served by calling the vendor's own
// API and require a resolvable per-account (or platform-default) credential.
const (
tierSelfHost = "self-host"
providerPrefix = "provider:"
vendorAnthropic = "anthropic"
modelARNFormat = "arn:aws:bedrock:*::foundation-model/%s"
)
// catalogEntry is one static catalog record. Model IDs mirror AWS exactly so
// existing SDK code and configs are drop-in.
type catalogEntry struct {
ModelID string
ModelName string
ProviderName string
Provider string // "self-host" or "provider:<vendor>"
InputModalities []string
OutputModalities []string
ResponseStreamingSupported bool
InferenceTypesSupported []string
CustomizationsSupported []string
}
// catalog is the static Phase-1 model set: one self-hosted open model and one
// Anthropic-direct model. Later phases extend this list.
var catalog = []catalogEntry{
{
ModelID: "meta.llama3-70b-instruct-v1:0",
ModelName: "Llama 3 70B Instruct",
ProviderName: "Meta",
Provider: tierSelfHost,
InputModalities: []string{"TEXT"},
OutputModalities: []string{"TEXT"},
ResponseStreamingSupported: false,
InferenceTypesSupported: []string{"ON_DEMAND"},
},
{
ModelID: "anthropic.claude-3-5-sonnet-20240620-v1:0",
ModelName: "Claude 3.5 Sonnet",
ProviderName: "Anthropic",
Provider: providerPrefix + vendorAnthropic,
InputModalities: []string{"TEXT", "IMAGE"},
OutputModalities: []string{"TEXT"},
ResponseStreamingSupported: false,
InferenceTypesSupported: []string{"ON_DEMAND"},
},
}
// CredentialResolver resolves accountID's usable provider credential for
// vendor: a per-account key, else an optional platform default. key is only
// meaningful when ok is true.
type CredentialResolver interface {
Resolve(ctx context.Context, accountID, vendor string) (key string, ok bool, err error)
}
// tieredCatalog returns the catalog entries advertised to accountID: self-host
// entries always, provider entries only when resolver finds a usable credential.
func tieredCatalog(ctx context.Context, accountID string, resolver CredentialResolver) []catalogEntry {
var out []catalogEntry
for _, entry := range catalog {
if entry.Provider == tierSelfHost {
out = append(out, entry)
continue
}
vendor, ok := strings.CutPrefix(entry.Provider, providerPrefix)
if !ok {
continue
}
_, resolvable, err := resolver.Resolve(ctx, accountID, vendor)
if err != nil || !resolvable {
continue
}
out = append(out, entry)
}
return out
}
func (e catalogEntry) toSummary() *bedrock.FoundationModelSummary {
return &bedrock.FoundationModelSummary{
ModelArn: aws.String(modelARN(e.ModelID)),
ModelId: aws.String(e.ModelID),
ModelName: aws.String(e.ModelName),
ProviderName: aws.String(e.ProviderName),
InputModalities: aws.StringSlice(e.InputModalities),
OutputModalities: aws.StringSlice(e.OutputModalities),
ResponseStreamingSupported: aws.Bool(e.ResponseStreamingSupported),
InferenceTypesSupported: aws.StringSlice(e.InferenceTypesSupported),
CustomizationsSupported: aws.StringSlice(e.CustomizationsSupported),
}
}
func (e catalogEntry) toDetails() *bedrock.FoundationModelDetails {
return &bedrock.FoundationModelDetails{
ModelArn: aws.String(modelARN(e.ModelID)),
ModelId: aws.String(e.ModelID),
ModelName: aws.String(e.ModelName),
ProviderName: aws.String(e.ProviderName),
InputModalities: aws.StringSlice(e.InputModalities),
OutputModalities: aws.StringSlice(e.OutputModalities),
ResponseStreamingSupported: aws.Bool(e.ResponseStreamingSupported),
InferenceTypesSupported: aws.StringSlice(e.InferenceTypesSupported),
CustomizationsSupported: aws.StringSlice(e.CustomizationsSupported),
}
}
func modelARN(modelID string) string {
return fmt.Sprintf(modelARNFormat, modelID)
}
// lookupCatalogEntry finds a catalog entry by exact modelId, ignoring tier
// gating (used by the runtime router, which returns its own error class).
func lookupCatalogEntry(modelID string) (catalogEntry, bool) {
for _, entry := range catalog {
if entry.ModelID == modelID {
return entry, true
}
}
return catalogEntry{}, false
}
// ListFoundationModels returns the deployment-tiered catalog: self-host
// entries always, provider entries only where a credential resolves.
func ListFoundationModels(ctx context.Context, accountID string, resolver CredentialResolver, _ *bedrock.ListFoundationModelsInput) (*bedrock.ListFoundationModelsOutput, error) {
entries := tieredCatalog(ctx, accountID, resolver)
summaries := make([]*bedrock.FoundationModelSummary, 0, len(entries))
for _, entry := range entries {
summaries = append(summaries, entry.toSummary())
}
return &bedrock.ListFoundationModelsOutput{ModelSummaries: summaries}, nil
}
// GetFoundationModel looks up a single model by exact modelId, independent of
// tiering. Unknown models return ResourceNotFoundException.
func GetFoundationModel(_ context.Context, _ string, modelID string) (*bedrock.GetFoundationModelOutput, error) {
entry, ok := lookupCatalogEntry(modelID)
if !ok {
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
return &bedrock.GetFoundationModelOutput{ModelDetails: entry.toDetails()}, nil
}
+62
View File
@@ -0,0 +1,62 @@
package gateway_bedrock
import (
"context"
"testing"
"github.com/aws/aws-sdk-go/service/bedrock"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// stubResolver reports a resolvable credential (with a fixed key) only for
// vendors in ok.
type stubResolver struct {
ok map[string]bool
}
func (s stubResolver) Resolve(_ context.Context, _, vendor string) (string, bool, error) {
if !s.ok[vendor] {
return "", false, nil
}
return "stub-key", true, nil
}
func modelIDs(out *bedrock.ListFoundationModelsOutput) []string {
ids := make([]string, 0, len(out.ModelSummaries))
for _, m := range out.ModelSummaries {
ids = append(ids, *m.ModelId)
}
return ids
}
func TestListFoundationModels_SelfHostAlwaysIncluded(t *testing.T) {
out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{})
require.NoError(t, err)
assert.Contains(t, modelIDs(out), "meta.llama3-70b-instruct-v1:0")
}
func TestListFoundationModels_ProviderIncludedWhenResolvable(t *testing.T) {
out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{"anthropic": true}}, &bedrock.ListFoundationModelsInput{})
require.NoError(t, err)
assert.Contains(t, modelIDs(out), "anthropic.claude-3-5-sonnet-20240620-v1:0")
}
func TestListFoundationModels_ProviderExcludedWhenUnresolvable(t *testing.T) {
out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{})
require.NoError(t, err)
assert.NotContains(t, modelIDs(out), "anthropic.claude-3-5-sonnet-20240620-v1:0")
}
func TestGetFoundationModel_KnownModel(t *testing.T) {
out, err := GetFoundationModel(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0")
require.NoError(t, err)
require.NotNil(t, out.ModelDetails)
assert.Equal(t, "meta.llama3-70b-instruct-v1:0", *out.ModelDetails.ModelId)
}
func TestGetFoundationModel_UnknownModelReturnsNotFound(t *testing.T) {
_, err := GetFoundationModel(context.Background(), "000000000001", "does-not-exist")
require.Error(t, err)
assert.Equal(t, "ResourceNotFoundException", err.Error())
}
+130
View File
@@ -0,0 +1,130 @@
package gateway_bedrock
import (
"context"
"fmt"
"sync"
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
"github.com/nats-io/nats.go"
)
// bedrockCredentialsBucket is the cluster-replicated KV bucket holding
// per-account provider API keys, encrypted at rest with the IAM master key.
const bedrockCredentialsBucket = "bedrock-credentials" //nolint:gosec // G101: KV bucket name, not a credential
// bedrockCredentialsHistory keeps one revision; rotation overwrites in place.
const bedrockCredentialsHistory = 1
// credentialKey returns the KV key for accountID's vendor credential.
func credentialKey(accountID, vendor string) string {
return fmt.Sprintf("providers/%s/%s", accountID, vendor)
}
// CredentialStore resolves per-account provider API keys from the
// bedrock-credentials JetStream KV bucket, falling back to a platform-wide
// default when an account has none. Keys are encrypted at rest.
type CredentialStore struct {
js nats.JetStreamContext
masterKey []byte
replicas int
platformDefaults map[string]string
mu sync.Mutex
kv nats.KeyValue
}
var _ CredentialResolver = (*CredentialStore)(nil)
// NewCredentialStore constructs a CredentialStore. platformDefaults maps
// vendor to a platform-wide API key used when an account has no key of its
// own; pass an empty map to disable platform defaults entirely.
func NewCredentialStore(js nats.JetStreamContext, masterKey []byte, replicas int, platformDefaults map[string]string) *CredentialStore {
return &CredentialStore{
js: js,
masterKey: masterKey,
replicas: replicas,
platformDefaults: platformDefaults,
}
}
// bucket lazily opens (or creates) the bedrock-credentials KV bucket,
// caching the handle, mirroring gateway_ecrauth.openSigningBucket.
func (s *CredentialStore) bucket() (nats.KeyValue, error) {
s.mu.Lock()
defer s.mu.Unlock()
if s.kv != nil {
return s.kv, nil
}
kv, err := s.js.CreateKeyValue(&nats.KeyValueConfig{
Bucket: bedrockCredentialsBucket,
History: bedrockCredentialsHistory,
Replicas: s.replicas,
})
if err != nil {
if kv, err = s.js.KeyValue(bedrockCredentialsBucket); err != nil {
return nil, fmt.Errorf("open %s bucket: %w", bedrockCredentialsBucket, err)
}
}
s.kv = kv
return kv, nil
}
// Resolve returns accountID's vendor API key: a per-account key if one is
// stored, else the platform default, else ("", false, nil). A nil js is
// tolerated so a store built with only platformDefaults (e.g. in tests) can
// still serve the default-only path without touching JetStream.
func (s *CredentialStore) Resolve(_ context.Context, accountID, vendor string) (string, bool, error) {
if s.js != nil {
kv, err := s.bucket()
if err != nil {
return "", false, err
}
entry, err := kv.Get(credentialKey(accountID, vendor))
switch {
case err == nil:
plaintext, err := handlers_iam.DecryptSecret(string(entry.Value()), s.masterKey)
if err != nil {
return "", false, fmt.Errorf("decrypt credential for %s/%s: %w", accountID, vendor, err)
}
return plaintext, true, nil
case err != nats.ErrKeyNotFound:
return "", false, fmt.Errorf("kv get credential for %s/%s: %w", accountID, vendor, err)
}
}
if key, ok := s.platformDefaults[vendor]; ok && key != "" {
return key, true, nil
}
return "", false, nil
}
// PutCredential encrypts and stores key as accountID's vendor credential.
func (s *CredentialStore) PutCredential(_ context.Context, accountID, vendor, key string) error {
kv, err := s.bucket()
if err != nil {
return err
}
ciphertext, err := handlers_iam.EncryptSecret(key, s.masterKey)
if err != nil {
return fmt.Errorf("encrypt credential for %s/%s: %w", accountID, vendor, err)
}
if _, err := kv.Put(credentialKey(accountID, vendor), []byte(ciphertext)); err != nil {
return fmt.Errorf("kv put credential for %s/%s: %w", accountID, vendor, err)
}
return nil
}
// noopCredentialResolver is the zero-value fallback used when no
// CredentialStore is configured (e.g. unit tests of unrelated routes): it
// resolves no vendors.
type noopCredentialResolver struct{}
var _ CredentialResolver = (*noopCredentialResolver)(nil)
func (noopCredentialResolver) Resolve(_ context.Context, _, _ string) (string, bool, error) {
return "", false, nil
}
// NoopCredentialResolver resolves no vendors. Used as the fallback wherever
// no CredentialStore is configured (nil GatewayConfig.BedrockCredentials).
var NoopCredentialResolver CredentialResolver = noopCredentialResolver{}
@@ -0,0 +1,116 @@
package gateway_bedrock
import (
"bytes"
"context"
"testing"
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
"github.com/mulgadc/spinifex/spinifex/testutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// fixedMasterKey returns a deterministic 32-byte AES-256 key for tests.
func fixedMasterKey() []byte {
return bytes.Repeat([]byte{0x42}, 32)
}
// TestCredentialStore_Resolve_NilJS covers the platform-default and
// vendor-miss paths of a store built with a nil JetStream context (as a
// live NATS server is out of scope here). Resolve must consult
// platformDefaults without dereferencing js.
func TestCredentialStore_Resolve_NilJS(t *testing.T) {
store := NewCredentialStore(nil, fixedMasterKey(), 1, map[string]string{"anthropic": "sk-platform-default"})
cases := []struct {
name string
vendor string
wantKey string
wantOK bool
}{
{"platform default hit", "anthropic", "sk-platform-default", true},
{"vendor miss", "openai", "", false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
key, ok, err := store.Resolve(context.Background(), "000000000001", tc.vendor)
require.NoError(t, err)
assert.Equal(t, tc.wantOK, ok)
assert.Equal(t, tc.wantKey, key)
})
}
}
func TestCredentialStore_Resolve_NoPlatformDefaults(t *testing.T) {
store := NewCredentialStore(nil, fixedMasterKey(), 1, nil)
key, ok, err := store.Resolve(context.Background(), "000000000001", "anthropic")
require.NoError(t, err)
assert.False(t, ok)
assert.Empty(t, key)
}
func TestNoopCredentialResolver_ResolvesNothing(t *testing.T) {
key, ok, err := NoopCredentialResolver.Resolve(context.Background(), "000000000001", "anthropic")
require.NoError(t, err)
assert.False(t, ok)
assert.Empty(t, key)
}
// TestCredentialStore_PutAndResolve_KV exercises the real JetStream KV path:
// bucket (lazy create), credentialKey, PutCredential (encrypt+Put), and the
// KV-hit branch of Resolve (Get+decrypt).
func TestCredentialStore_PutAndResolve_KV(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewCredentialStore(js, fixedMasterKey(), 1, nil)
ctx := context.Background()
require.NoError(t, store.PutCredential(ctx, "000000000001", "anthropic", "sk-test"))
key, ok, err := store.Resolve(ctx, "000000000001", "anthropic")
require.NoError(t, err)
assert.True(t, ok)
assert.Equal(t, "sk-test", key)
}
// TestCredentialStore_Resolve_KVMissFallsBackToPlatformDefault covers a KV
// miss (nats.ErrKeyNotFound) that falls through to the platform default.
func TestCredentialStore_Resolve_KVMissFallsBackToPlatformDefault(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewCredentialStore(js, fixedMasterKey(), 1, map[string]string{"anthropic": "sk-platform-default"})
key, ok, err := store.Resolve(context.Background(), "999999999999", "anthropic")
require.NoError(t, err)
assert.True(t, ok)
assert.Equal(t, "sk-platform-default", key)
}
// TestCredentialStore_Resolve_KVMissNoPlatformDefault covers a KV miss on an
// unknown account with no platform defaults configured: ("", false, nil).
func TestCredentialStore_Resolve_KVMissNoPlatformDefault(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewCredentialStore(js, fixedMasterKey(), 1, nil)
key, ok, err := store.Resolve(context.Background(), "999999999999", "anthropic")
require.NoError(t, err)
assert.False(t, ok)
assert.Empty(t, key)
}
func TestCredentialKey(t *testing.T) {
assert.Equal(t, "providers/000000000001/anthropic", credentialKey("000000000001", "anthropic"))
}
func TestEncryptDecryptSecret_RoundTrip(t *testing.T) {
key := fixedMasterKey()
plaintext := "sk-ant-test-secret"
ciphertext, err := handlers_iam.EncryptSecret(plaintext, key)
require.NoError(t, err)
assert.NotEqual(t, plaintext, ciphertext)
decrypted, err := handlers_iam.DecryptSecret(ciphertext, key)
require.NoError(t, err)
assert.Equal(t, plaintext, decrypted)
}
+71
View File
@@ -0,0 +1,71 @@
// Package gateway_bedrock is the HTTP-side glue between awsgw and the Ochre
// (bedrock / bedrock-runtime) handlers. Both speak AWS REST-JSON 1.1; error
// envelopes are JSON and Content-Type is application/x-amz-json-1.1 throughout.
package gateway_bedrock
import (
"encoding/json"
"fmt"
"log/slog"
"net/http"
"strings"
"github.com/aws/aws-sdk-go/private/protocol/json/jsonutil"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
// JSONContentType is the AWS REST-JSON 1.1 content type bedrock clients expect.
const JSONContentType = "application/x-amz-json-1.1"
// jsonError is the AWS REST-JSON error envelope. SDKs key off __type for
// awserr.Code() and message for awserr.Message().
type jsonError struct {
Type string `json:"__type"`
Message string `json:"message"`
}
// generateJSONError marshals the AWS REST-JSON error envelope. The suffix
// "Exception" is appended idempotently — codes that already end in it are
// left unchanged to avoid the double "ExceptionException" the SDK rejects.
func generateJSONError(code, message string) []byte {
if !strings.HasSuffix(code, "Exception") {
code += "Exception"
}
body, err := json.Marshal(jsonError{Type: code, Message: message})
if err != nil {
slog.Error("Failed to marshal bedrock error JSON", "code", code, "err", err)
return fmt.Appendf(nil, `{"__type":"InternalErrorException","message":%q}`, message)
}
return body
}
// WriteJSONResponse serializes obj as AWS REST-JSON and writes a 200 response.
// Uses jsonutil.BuildJSON (the SDK's own restjson marshaler) to honour
// locationName tags that encoding/json does not understand.
func WriteJSONResponse(w http.ResponseWriter, obj any) {
body, err := jsonutil.BuildJSON(obj)
if err != nil {
slog.Error("Failed to marshal bedrock response JSON", "err", err)
WriteJSONError(w, awserrors.ErrorInternalError, "failed to marshal response", http.StatusInternalServerError)
return
}
w.Header().Set("Content-Type", JSONContentType)
w.WriteHeader(http.StatusOK)
if _, err := w.Write(body); err != nil {
slog.Error("Failed to write bedrock response", "err", err)
}
}
// WriteJSONError emits the AWS REST-JSON error envelope with the given code,
// message, and HTTP status.
func WriteJSONError(w http.ResponseWriter, code, message string, httpStatus int) {
body := generateJSONError(code, message)
w.Header().Set("Content-Type", JSONContentType)
if httpStatus == 0 {
httpStatus = http.StatusInternalServerError
}
w.WriteHeader(httpStatus)
if _, err := w.Write(body); err != nil {
slog.Error("Failed to write bedrock error response", "err", err)
}
}
+75
View File
@@ -0,0 +1,75 @@
package gateway_bedrock
import (
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestWriteJSONResponse_WritesBodyAndContentType(t *testing.T) {
type payload struct {
Foo string `json:"foo"`
}
w := httptest.NewRecorder()
WriteJSONResponse(w, payload{Foo: "bar"})
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, JSONContentType, w.Header().Get("Content-Type"))
var out payload
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &out))
assert.Equal(t, "bar", out.Foo)
}
func TestWriteJSONError_WritesEnvelopeAndStatus(t *testing.T) {
cases := []struct {
name string
code string
httpStatus int
wantStatus int
wantType string
}{
{"explicit status", awserrors.ErrorValidationException, http.StatusBadRequest, http.StatusBadRequest, "ValidationException"},
{"zero status defaults to 500", awserrors.ErrorInternalError, 0, http.StatusInternalServerError, "InternalErrorException"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
w := httptest.NewRecorder()
WriteJSONError(w, tc.code, "boom", tc.httpStatus)
assert.Equal(t, tc.wantStatus, w.Code)
assert.Equal(t, JSONContentType, w.Header().Get("Content-Type"))
var env jsonError
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &env))
assert.Equal(t, tc.wantType, env.Type)
assert.Equal(t, "boom", env.Message)
})
}
}
func TestGenerateJSONError_AppendsExceptionSuffixIdempotently(t *testing.T) {
cases := []struct {
name string
code string
wantType string
}{
{"bare code gets suffix", "Validation", "ValidationException"},
{"already-suffixed code unchanged", "ValidationException", "ValidationException"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
body := generateJSONError(tc.code, "msg")
var env jsonError
require.NoError(t, json.Unmarshal(body, &env))
assert.Equal(t, tc.wantType, env.Type)
assert.Equal(t, "msg", env.Message)
})
}
}
+24
View File
@@ -0,0 +1,24 @@
package gateway_bedrock
import (
"net/http"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
// mapUpstreamStatus maps an upstream provider HTTP status to a bedrock
// awserrors code, shared by the Anthropic and vLLM adapters.
func mapUpstreamStatus(status int) string {
switch {
case status == http.StatusUnauthorized || status == http.StatusForbidden:
return awserrors.ErrorAccessDeniedException
case status == http.StatusTooManyRequests:
return awserrors.ErrorThrottlingException
case status == http.StatusBadRequest:
return awserrors.ErrorValidationException
case status >= 500:
return awserrors.ErrorServiceUnavailableException
default:
return awserrors.ErrorModelErrorException
}
}
@@ -0,0 +1,31 @@
package gateway_bedrock
import (
"net/http"
"testing"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/stretchr/testify/assert"
)
func TestMapUpstreamStatus(t *testing.T) {
cases := []struct {
name string
status int
want string
}{
{"401 unauthorized -> access denied", http.StatusUnauthorized, awserrors.ErrorAccessDeniedException},
{"403 forbidden -> access denied", http.StatusForbidden, awserrors.ErrorAccessDeniedException},
{"429 too many requests -> throttling", http.StatusTooManyRequests, awserrors.ErrorThrottlingException},
{"400 bad request -> validation", http.StatusBadRequest, awserrors.ErrorValidationException},
{"500 internal server error -> service unavailable", http.StatusInternalServerError, awserrors.ErrorServiceUnavailableException},
{"503 service unavailable -> service unavailable", http.StatusServiceUnavailable, awserrors.ErrorServiceUnavailableException},
{"418 teapot -> model error", http.StatusTeapot, awserrors.ErrorModelErrorException},
{"other unmapped status -> model error", http.StatusConflict, awserrors.ErrorModelErrorException},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, mapUpstreamStatus(tc.status))
})
}
}
+86
View File
@@ -0,0 +1,86 @@
package gateway_bedrock
import (
"context"
"errors"
"strings"
"time"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
// providerHTTPTimeout bounds outbound calls to both provider backends. It is
// long because Phase 1 is synchronous (no streaming) and matches the
// platform's tolerance for a single Converse call.
const providerHTTPTimeout = 15 * time.Minute
// Provider translates a Converse request into a backend's native wire format
// and back. vllmProvider serves self-hosted models over OpenAI chat
// completions; anthropicProvider (via newAnthropicProvider) serves Claude
// over the Anthropic Messages API with a per-call API key baked in.
type Provider interface {
Converse(ctx context.Context, modelID string, input *bedrockruntime.ConverseInput) (*bedrockruntime.ConverseOutput, error)
}
// Router resolves a modelId to its catalog entry and dispatches to the
// matching Provider, resolving self-host endpoints and provider credentials
// as needed.
type Router struct {
resolver CredentialResolver
endpointResolver EndpointResolver
}
// NewRouter constructs a Router. A nil resolver or endpointResolver falls
// back to a resolver/resolver that finds nothing, so a Router is always safe
// to use even before the real stores are wired in.
func NewRouter(resolver CredentialResolver, endpointResolver EndpointResolver) *Router {
if resolver == nil {
resolver = NoopCredentialResolver
}
if endpointResolver == nil {
endpointResolver = NewStaticEndpointResolver(nil)
}
return &Router{resolver: resolver, endpointResolver: endpointResolver}
}
// Converse routes modelID to its provider via the catalog. Unknown modelIds
// and unresolvable vendors return ResourceNotFoundException; a vendor with no
// resolvable credential returns AccessDeniedException.
func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput) (*bedrockruntime.ConverseOutput, error) {
entry, ok := lookupCatalogEntry(modelID)
if !ok {
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
var p Provider
switch {
case entry.Provider == tierSelfHost:
p = newVLLMProvider(rt.endpointResolver)
case strings.HasPrefix(entry.Provider, providerPrefix):
switch strings.TrimPrefix(entry.Provider, providerPrefix) {
case vendorAnthropic:
key, ok, err := rt.resolver.Resolve(ctx, accountID, vendorAnthropic)
if err != nil {
return nil, err
}
if !ok {
return nil, errors.New(awserrors.ErrorAccessDeniedException)
}
p = newAnthropicProvider(key)
default:
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
default:
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
return p.Converse(ctx, modelID, input)
}
// Converse is the bedrock-runtime Converse entry point used by the gateway
// route table. resolver and endpointResolver may be nil; NewRouter supplies
// no-op fallbacks.
func Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput, resolver CredentialResolver, endpointResolver EndpointResolver) (*bedrockruntime.ConverseOutput, error) {
return NewRouter(resolver, endpointResolver).Converse(ctx, accountID, modelID, input)
}
+104
View File
@@ -0,0 +1,104 @@
package gateway_bedrock
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// stubCredentialResolver returns a fixed (key, ok, err) for every Resolve
// call, letting tests drive the Router's anthropic branch without a real
// CredentialStore.
type stubCredentialResolver struct {
key string
ok bool
err error
}
func (s stubCredentialResolver) Resolve(_ context.Context, _, _ string) (string, bool, error) {
return s.key, s.ok, s.err
}
func converseInput() *bedrockruntime.ConverseInput {
return &bedrockruntime.ConverseInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
}
}
func TestRouter_Converse_SelfHostSuccess(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{
"choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1}
}`))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
out, err := rt.Converse(context.Background(), "000000000001", modelID, converseInput())
require.NoError(t, err)
require.NotNil(t, out.Output.Message)
assert.Equal(t, "hi", *out.Output.Message.Content[0].Text)
}
func TestRouter_Converse_UnknownModelReturnsResourceNotFound(t *testing.T) {
rt := NewRouter(nil, nil)
_, err := rt.Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput())
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
func TestRouter_Converse_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) {
rt := NewRouter(stubCredentialResolver{ok: false}, nil)
_, err := rt.Converse(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", converseInput())
require.Error(t, err)
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
}
func TestConverse_PackageEntryPoint_SelfHostSuccess(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{
"choices": [{"message": {"role": "assistant", "content": "via package Converse"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1}
}`))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
out, err := Converse(context.Background(), "000000000001", modelID, converseInput(), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
require.NoError(t, err)
require.NotNil(t, out.Output.Message)
assert.Equal(t, "via package Converse", *out.Output.Message.Content[0].Text)
}
func TestConverse_PackageEntryPoint_UnknownModel(t *testing.T) {
_, err := Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput(), nil, nil)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
func TestNewRouter_NilArgumentsFallBackToNoops(t *testing.T) {
rt := NewRouter(nil, nil)
require.NotNil(t, rt)
// Self-host model with no endpoint resolver configured resolves nothing,
// so it must report ModelNotReady rather than panic on a nil resolver.
_, err := rt.Converse(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0", converseInput())
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
}
+224
View File
@@ -0,0 +1,224 @@
package gateway_bedrock
import (
"bytes"
"context"
"encoding/json"
"errors"
"io"
"log/slog"
"net/http"
"strings"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
const vllmChatCompletionsPath = "/v1/chat/completions"
// EndpointResolver resolves a self-hosted model's OpenAI-compatible base URL.
// A later task backs this with the daemon's endpoint registry.
type EndpointResolver interface {
Endpoint(ctx context.Context, modelID string) (baseURL string, ok bool, err error)
}
// staticEndpointResolver is a fixed modelID->baseURL map, for tests and
// config-driven single-endpoint setups.
type staticEndpointResolver map[string]string
var _ EndpointResolver = staticEndpointResolver(nil)
// NewStaticEndpointResolver returns an EndpointResolver backed by a fixed
// modelID->baseURL map. A nil or empty map resolves nothing.
func NewStaticEndpointResolver(endpoints map[string]string) EndpointResolver {
return staticEndpointResolver(endpoints)
}
func (s staticEndpointResolver) Endpoint(_ context.Context, modelID string) (string, bool, error) {
baseURL, ok := s[modelID]
return baseURL, ok, nil
}
type vllmMessage struct {
Role string `json:"role"`
Content string `json:"content"`
}
type vllmChatRequest struct {
Model string `json:"model"`
Messages []vllmMessage `json:"messages"`
MaxTokens *int64 `json:"max_tokens,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
Stop []string `json:"stop,omitempty"`
}
type vllmChoice struct {
Message vllmMessage `json:"message"`
FinishReason string `json:"finish_reason"`
}
type vllmUsage struct {
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
}
type vllmChatResponse struct {
Choices []vllmChoice `json:"choices"`
Usage vllmUsage `json:"usage"`
}
// vllmFinishReasons maps OpenAI finish_reason values to Bedrock's StopReason
// enum. Unrecognised values fall back to end_turn.
var vllmFinishReasons = map[string]string{
"stop": bedrockruntime.StopReasonEndTurn,
"length": bedrockruntime.StopReasonMaxTokens,
}
func mapVLLMFinishReason(reason string) string {
if mapped, ok := vllmFinishReasons[reason]; ok {
return mapped
}
return bedrockruntime.StopReasonEndTurn
}
// vllmProvider serves self-hosted open-weight models over the OpenAI
// chat-completions wire. httpClient and endpointResolver are injectable for
// tests.
type vllmProvider struct {
endpointResolver EndpointResolver
httpClient *http.Client
}
var _ Provider = (*vllmProvider)(nil)
func newVLLMProvider(endpointResolver EndpointResolver) *vllmProvider {
return &vllmProvider{
endpointResolver: endpointResolver,
httpClient: &http.Client{Timeout: providerHTTPTimeout},
}
}
// Converse resolves modelID's endpoint, translates input to an OpenAI
// chat-completions request, calls the endpoint, and translates the response
// back to a Bedrock ConverseOutput.
func (p *vllmProvider) Converse(ctx context.Context, modelID string, input *bedrockruntime.ConverseInput) (*bedrockruntime.ConverseOutput, error) {
baseURL, ok, err := p.endpointResolver.Endpoint(ctx, modelID)
if err != nil {
slog.Error("vllm: endpoint resolution failed", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
if !ok {
return nil, errors.New(awserrors.ErrorModelNotReadyException)
}
reqBody, err := json.Marshal(buildVLLMRequest(modelID, input))
if err != nil {
slog.Error("vllm: failed to marshal request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+vllmChatCompletionsPath, bytes.NewReader(reqBody))
if err != nil {
slog.Error("vllm: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq.Header.Set("content-type", "application/json")
start := time.Now()
resp, err := p.httpClient.Do(httpReq)
latency := time.Since(start)
if err != nil {
slog.Error("vllm: request failed", "model", modelID, "endpoint", baseURL, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
slog.Error("vllm: failed to read response body", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
slog.Error("vllm: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
return nil, errors.New(mapUpstreamStatus(resp.StatusCode))
}
var cr vllmChatResponse
if err := json.Unmarshal(respBody, &cr); err != nil {
slog.Error("vllm: failed to parse response", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorModelErrorException)
}
return vllmToConverseOutput(cr, latency)
}
// buildVLLMRequest maps a Bedrock ConverseInput to an OpenAI chat-completions
// request. Non-text content blocks are dropped (Phase 1 is text-only).
func buildVLLMRequest(modelID string, input *bedrockruntime.ConverseInput) vllmChatRequest {
var messages []vllmMessage
var systemParts []string
for _, s := range input.System {
if s != nil && s.Text != nil {
systemParts = append(systemParts, *s.Text)
}
}
if len(systemParts) > 0 {
messages = append(messages, vllmMessage{Role: "system", Content: strings.Join(systemParts, "\n")})
}
for _, m := range input.Messages {
if m == nil {
continue
}
var textParts []string
for _, c := range m.Content {
// TODO: Phase 1 is text-only; non-text content blocks are dropped.
if c == nil || c.Text == nil {
continue
}
textParts = append(textParts, *c.Text)
}
messages = append(messages, vllmMessage{Role: aws.StringValue(m.Role), Content: strings.Join(textParts, "\n")})
}
req := vllmChatRequest{Model: modelID, Messages: messages}
if input.InferenceConfig != nil {
req.MaxTokens = input.InferenceConfig.MaxTokens
req.Temperature = input.InferenceConfig.Temperature
req.TopP = input.InferenceConfig.TopP
req.Stop = aws.StringValueSlice(input.InferenceConfig.StopSequences)
}
return req
}
// vllmToConverseOutput maps an OpenAI chat-completions response to a Bedrock
// ConverseOutput, measuring latency from the caller-supplied wall-clock delta.
func vllmToConverseOutput(cr vllmChatResponse, latency time.Duration) (*bedrockruntime.ConverseOutput, error) {
if len(cr.Choices) == 0 {
return nil, errors.New(awserrors.ErrorModelErrorException)
}
choice := cr.Choices[0]
inputTokens := cr.Usage.PromptTokens
outputTokens := cr.Usage.CompletionTokens
return &bedrockruntime.ConverseOutput{
Output: &bedrockruntime.ConverseOutput_{
Message: &bedrockruntime.Message{
Role: aws.String(bedrockruntime.ConversationRoleAssistant),
Content: []*bedrockruntime.ContentBlock{{Text: aws.String(choice.Message.Content)}},
},
},
StopReason: aws.String(mapVLLMFinishReason(choice.FinishReason)),
Usage: &bedrockruntime.TokenUsage{
InputTokens: aws.Int64(inputTokens),
OutputTokens: aws.Int64(outputTokens),
TotalTokens: aws.Int64(inputTokens + outputTokens),
},
Metrics: &bedrockruntime.ConverseMetrics{LatencyMs: aws.Int64(latency.Milliseconds())},
}, nil
}
+73
View File
@@ -0,0 +1,73 @@
package gateway_bedrock
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestVLLMProvider_Converse_MapsResponse(t *testing.T) {
var captured vllmChatRequest
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.NoError(t, json.NewDecoder(r.Body).Decode(&captured))
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{
"choices": [{"message": {"role": "assistant", "content": "hi there"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 8, "completion_tokens": 3}
}`))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
p.httpClient = ts.Client()
input := &bedrockruntime.ConverseInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
InferenceConfig: &bedrockruntime.InferenceConfiguration{MaxTokens: aws.Int64(128)},
}
out, err := p.Converse(context.Background(), modelID, input)
require.NoError(t, err)
assert.Equal(t, modelID, captured.Model)
require.NotNil(t, captured.MaxTokens)
assert.Equal(t, int64(128), *captured.MaxTokens)
require.Len(t, captured.Messages, 1)
assert.Equal(t, "user", captured.Messages[0].Role)
assert.Equal(t, "hello", captured.Messages[0].Content)
require.NotNil(t, out.Output.Message)
require.Len(t, out.Output.Message.Content, 1)
assert.Equal(t, "hi there", *out.Output.Message.Content[0].Text)
assert.Equal(t, int64(8), *out.Usage.InputTokens)
assert.Equal(t, int64(3), *out.Usage.OutputTokens)
assert.Equal(t, int64(11), *out.Usage.TotalTokens)
assert.Equal(t, bedrockruntime.StopReasonEndTurn, *out.StopReason)
}
func TestVLLMProvider_Converse_UnresolvedEndpointReturnsModelNotReady(t *testing.T) {
p := newVLLMProvider(NewStaticEndpointResolver(nil))
input := &bedrockruntime.ConverseInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
}
_, err := p.Converse(context.Background(), "meta.llama3-70b-instruct-v1:0", input)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
}
+166
View File
@@ -0,0 +1,166 @@
package gateway
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/aws/aws-sdk-go/service/bedrock"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/testutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
const bedrockTestAccount = "123456789012"
const bedrockTestLlamaModelID = "meta.llama3-70b-instruct-v1:0"
// newVLLMStub stands up an httptest server that answers the OpenAI
// chat-completions wire, mirroring gateway/bedrock/vllm_test.go's stub.
func newVLLMStub(t *testing.T) *httptest.Server {
t.Helper()
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{
"choices": [{"message": {"role": "assistant", "content": "hi from vllm"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 4, "completion_tokens": 2}
}`))
}))
t.Cleanup(ts.Close)
return ts
}
// newBedrockRequestGateway builds a GatewayConfig with a real NATS connection
// (satisfying Bedrock_Request/BedrockRuntime_Request's nil-check only — no
// NATS subject handling is required for these routes), no IAMService (so
// checkPolicy is a no-op), and BedrockEndpoints pinned at the given vLLM stub.
func newBedrockRequestGateway(t *testing.T, vllmURL string) *GatewayConfig {
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
return &GatewayConfig{
NATSConn: nc,
DisableLogging: true,
BedrockEndpoints: map[string]string{
bedrockTestLlamaModelID: vllmURL,
},
}
}
func bedrockRequestWithAccount(method, path, body string) *http.Request {
req := httptest.NewRequest(method, path, strings.NewReader(body))
ctx := context.WithValue(req.Context(), ctxAccountID, bedrockTestAccount)
return req.WithContext(ctx)
}
func TestBedrockRequest_ListFoundationModels(t *testing.T) {
ts := newVLLMStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
req := bedrockRequestWithAccount(http.MethodGet, "/foundation-models", "")
w := httptest.NewRecorder()
require.NoError(t, gw.Bedrock_Request(w, req))
require.Equal(t, http.StatusOK, w.Code)
var out bedrock.ListFoundationModelsOutput
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &out))
var ids []string
for _, m := range out.ModelSummaries {
ids = append(ids, *m.ModelId)
}
assert.Contains(t, ids, bedrockTestLlamaModelID)
}
func TestBedrockRequest_GetFoundationModel(t *testing.T) {
ts := newVLLMStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
req := bedrockRequestWithAccount(http.MethodGet, "/foundation-models/"+bedrockTestLlamaModelID, "")
w := httptest.NewRecorder()
require.NoError(t, gw.Bedrock_Request(w, req))
require.Equal(t, http.StatusOK, w.Code)
var out bedrock.GetFoundationModelOutput
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &out))
require.NotNil(t, out.ModelDetails)
assert.Equal(t, bedrockTestLlamaModelID, *out.ModelDetails.ModelId)
}
func TestBedrockRequest_GetFoundationModel_NotFound(t *testing.T) {
ts := newVLLMStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
req := bedrockRequestWithAccount(http.MethodGet, "/foundation-models/does.not-exist-v1:0", "")
w := httptest.NewRecorder()
err := gw.Bedrock_Request(w, req)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
func TestBedrockRequest_UnknownRouteReturnsInvalidAction(t *testing.T) {
ts := newVLLMStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
req := bedrockRequestWithAccount(http.MethodDelete, "/foundation-models", "")
w := httptest.NewRecorder()
err := gw.Bedrock_Request(w, req)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorInvalidAction, err.Error())
}
func TestBedrockRequest_MissingAccountIDReturnsServerInternal(t *testing.T) {
ts := newVLLMStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
req := httptest.NewRequest(http.MethodGet, "/foundation-models", nil)
w := httptest.NewRecorder()
err := gw.Bedrock_Request(w, req)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorServerInternal, err.Error())
}
func TestBedrockRuntimeRequest_Converse(t *testing.T) {
ts := newVLLMStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
body := `{"messages":[{"role":"user","content":[{"text":"hello"}]}]}`
req := bedrockRequestWithAccount(http.MethodPost, "/model/"+bedrockTestLlamaModelID+"/converse", body)
w := httptest.NewRecorder()
require.NoError(t, gw.BedrockRuntime_Request(w, req))
require.Equal(t, http.StatusOK, w.Code)
var out bedrockruntime.ConverseOutput
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &out))
require.NotNil(t, out.Output)
require.NotNil(t, out.Output.Message)
assert.Equal(t, bedrockruntime.ConversationRoleAssistant, *out.Output.Message.Role)
require.Len(t, out.Output.Message.Content, 1)
assert.Equal(t, "hi from vllm", *out.Output.Message.Content[0].Text)
}
func TestBedrockRuntimeRequest_MalformedBodyReturnsValidationException(t *testing.T) {
ts := newVLLMStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
req := bedrockRequestWithAccount(http.MethodPost, "/model/"+bedrockTestLlamaModelID+"/converse", "{not-json")
w := httptest.NewRecorder()
err := gw.BedrockRuntime_Request(w, req)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorValidationException, err.Error())
}
func TestBedrockRuntimeRequest_MissingAccountIDReturnsServerInternal(t *testing.T) {
ts := newVLLMStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
req := httptest.NewRequest(http.MethodPost, "/model/"+bedrockTestLlamaModelID+"/converse", strings.NewReader("{}"))
w := httptest.NewRecorder()
err := gw.BedrockRuntime_Request(w, req)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorServerInternal, err.Error())
}
+48
View File
@@ -0,0 +1,48 @@
package gateway
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestLookupBedrockAction_ResolvesKnownRoutes(t *testing.T) {
cases := []struct {
method, path string
wantAction string
wantParams []string
}{
{"GET", "/foundation-models", "ListFoundationModels", nil},
{"GET", "/foundation-models/anthropic.claude-3-5-sonnet-20240620-v1:0", "GetFoundationModel", []string{"anthropic.claude-3-5-sonnet-20240620-v1:0"}},
}
for _, tc := range cases {
t.Run(tc.method+"_"+tc.path, func(t *testing.T) {
action, params, handler, ok := lookupBedrockAction(tc.method, tc.path)
require.True(t, ok, "expected route to match for %s %s", tc.method, tc.path)
require.NotNil(t, handler)
assert.Equal(t, tc.wantAction, action)
assert.Equal(t, tc.wantParams, params)
})
}
}
func TestLookupBedrockAction_UnknownReturnsFalse(t *testing.T) {
_, _, handler, ok := lookupBedrockAction("DELETE", "/foundation-models")
assert.False(t, ok)
assert.Nil(t, handler)
}
func TestLookupBedrockRuntimeAction_ResolvesConverse(t *testing.T) {
action, params, handler, ok := lookupBedrockRuntimeAction("POST", "/model/meta.llama3-70b-instruct-v1:0/converse")
require.True(t, ok)
require.NotNil(t, handler)
assert.Equal(t, "Converse", action)
assert.Equal(t, []string{"meta.llama3-70b-instruct-v1:0"}, params)
}
func TestLookupBedrockRuntimeAction_UnknownReturnsFalse(t *testing.T) {
_, _, handler, ok := lookupBedrockRuntimeAction("GET", "/model/foo/converse")
assert.False(t, ok)
assert.Nil(t, handler)
}
+113
View File
@@ -0,0 +1,113 @@
package gateway
import (
"context"
"encoding/json"
"errors"
"io"
"log/slog"
"net/http"
"net/url"
"regexp"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
)
// bedrockRuntimeRoute maps one HTTP method + path regex to an AWS action and handler.
type bedrockRuntimeRoute struct {
method string
pattern *regexp.Regexp
action string
handler bedrockRuntimeRouteHandler
}
// bedrockRuntimeRouteHandler invokes a per-action bedrock-runtime (data-plane)
// gateway function. params holds the regex capture groups, PathUnescape'd.
// resolver is gw.bedrockResolver() (credential store or no-op); endpoints is
// gw.bedrockEndpointResolver() over the configured pinned self-host endpoints.
type bedrockRuntimeRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver) (any, error)
// bedrockRuntimeRoutes is the dispatch table.
var bedrockRuntimeRoutes = []bedrockRuntimeRoute{
{"POST", regexp.MustCompile(`^/model/([^/]+)/converse$`), "Converse",
func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver) (any, error) {
input := new(bedrockruntime.ConverseInput)
if len(b) > 0 {
if err := json.Unmarshal(b, input); err != nil {
return nil, errors.New(awserrors.ErrorValidationException)
}
}
return gateway_bedrock.Converse(ctx, acct, p[0], input, resolver, endpoints)
}},
}
// lookupBedrockRuntimeAction matches method+path against bedrockRuntimeRoutes,
// returning the action, path params, and handler, or ("", nil, nil, false) on
// no match. path must be r.URL.EscapedPath(): captured params are
// PathUnescape'd before returning, mirroring lookupEKSAction.
func lookupBedrockRuntimeAction(method, path string) (string, []string, bedrockRuntimeRouteHandler, bool) {
for _, route := range bedrockRuntimeRoutes {
if route.method != method {
continue
}
m := route.pattern.FindStringSubmatch(path)
if m == nil {
continue
}
var params []string
if len(m) > 1 {
params = make([]string, 0, len(m)-1)
for _, raw := range m[1:] {
decoded, err := url.PathUnescape(raw)
if err != nil {
slog.Debug("bedrock-runtime: bad percent-encoding in path param", "param", raw, "err", err)
decoded = raw
}
params = append(params, decoded)
}
}
return route.action, params, route.handler, true
}
return "", nil, nil, false
}
// BedrockRuntime_Request dispatches bedrock-runtime (data-plane) REST-JSON
// requests: resolves method+path to an action, reads the body, calls the
// handler, and serialises the output as JSON.
func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.Request) error {
action, params, handler, ok := lookupBedrockRuntimeAction(r.Method, r.URL.EscapedPath())
if !ok {
slog.DebugContext(r.Context(), "bedrock-runtime: no route for request", "method", r.Method, "path", r.URL.Path)
return errors.New(awserrors.ErrorInvalidAction)
}
if err := gw.checkPolicy(r, "bedrock-runtime", action); err != nil {
return err
}
if gw.NATSConn == nil {
return errors.New(awserrors.ErrorServerInternal)
}
accountID, _ := r.Context().Value(ctxAccountID).(string)
if accountID == "" {
slog.ErrorContext(r.Context(), "BedrockRuntime_Request: no account ID in auth context")
return errors.New(awserrors.ErrorServerInternal)
}
body, err := io.ReadAll(r.Body)
if err != nil {
slog.ErrorContext(r.Context(), "BedrockRuntime_Request: failed to read body", "err", err)
return errors.New(awserrors.ErrorInvalidParameterValue)
}
output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockEndpointResolver())
if err != nil {
return err
}
gateway_bedrock.WriteJSONResponse(w, output)
return nil
}
+18 -2
View File
@@ -20,6 +20,7 @@ import (
"github.com/mulgadc/predastore/pkg/iampolicy"
"github.com/mulgadc/predastore/ratelimit"
"github.com/mulgadc/spinifex/spinifex/awserrors"
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
gateway_ecr "github.com/mulgadc/spinifex/spinifex/gateway/ecr"
gateway_ecrauth "github.com/mulgadc/spinifex/spinifex/gateway/ecrauth"
"github.com/mulgadc/spinifex/spinifex/gateway/policy"
@@ -110,6 +111,14 @@ type GatewayConfig struct {
// in unit tests of unrelated routes).
ECRTokenIssuer *gateway_ecrauth.Issuer
ECRTokenVerifier *gateway_ecrauth.Verifier
// BedrockCredentials resolves per-account provider API keys for bedrock
// routes. Nil falls back to no external providers (self-host models only).
BedrockCredentials *gateway_bedrock.CredentialStore
// BedrockEndpoints maps a self-hosted modelId to its OpenAI-compatible base
// URL. Phase 1 endpoints are pinned/always-resident, so this is static
// config; a later phase backs it with the daemon's dynamic endpoint
// registry. Nil/empty means no self-hosted models are reachable.
BedrockEndpoints map[string]string
}
var supportedServices = map[string]bool{
@@ -123,6 +132,8 @@ var supportedServices = map[string]bool{
"acm": true,
"tagging": true,
"spinifex": true,
"bedrock": true,
"bedrock-runtime": true,
}
// EC2ErrorResponse is the EC2 query-API error envelope.
@@ -340,6 +351,10 @@ func (gw *GatewayConfig) Request(w http.ResponseWriter, r *http.Request) {
err = gw.ELBv2_Request(w, r)
case "eks":
err = gw.EKS_Request(w, r)
case "bedrock":
err = gw.Bedrock_Request(w, r)
case "bedrock-runtime":
err = gw.BedrockRuntime_Request(w, r)
case "ecs":
err = gw.ECS_Request(w, r)
case "ecr":
@@ -501,8 +516,9 @@ func (gw *GatewayConfig) ErrorHandler(w http.ResponseWriter, r *http.Request, er
errorMsg.HTTPCode = 500
}
// EKS, ECR, ACM, ECS, and tagging use AWS JSON 1.1; query/XML services fall through.
if svc == "eks" || svc == "ecr" || svc == "acm" || svc == "ecs" || svc == "tagging" {
// EKS, ECR, ACM, ECS, tagging, and bedrock/bedrock-runtime use AWS JSON 1.1;
// query/XML services fall through.
if svc == "eks" || svc == "ecr" || svc == "acm" || svc == "ecs" || svc == "tagging" || svc == "bedrock" || svc == "bedrock-runtime" {
body := GenerateEKSErrorResponse(err.Error(), errorMsg.Message, requestId)
slog.Debug("Generated JSON error response", "service", svc, "error", err.Error(), "json", string(body), "requestId", requestId)
w.Header().Set("Content-Type", eksJSONContentType)
+59 -17
View File
@@ -18,6 +18,7 @@ import (
"github.com/mulgadc/spinifex/spinifex/admin"
"github.com/mulgadc/spinifex/spinifex/config"
"github.com/mulgadc/spinifex/spinifex/gateway"
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
gateway_ecr "github.com/mulgadc/spinifex/spinifex/gateway/ecr"
gateway_ecrauth "github.com/mulgadc/spinifex/spinifex/gateway/ecrauth"
"github.com/mulgadc/spinifex/spinifex/handlers/ecr"
@@ -283,24 +284,40 @@ func launchService(config *config.ClusterConfig) error {
registryHost = nodeConfig.AdvertiseIP
}
// Bedrock provider credentials: per-account keys live in the
// bedrock-credentials KV bucket; OCHRE_ANTHROPIC_API_KEY seeds an optional
// platform-wide default so accounts work before onboarding their own key.
bedrockPlatformDefaults := map[string]string{}
if key := os.Getenv("OCHRE_ANTHROPIC_API_KEY"); key != "" {
bedrockPlatformDefaults["anthropic"] = key
}
bedrockCredentials := gateway_bedrock.NewCredentialStore(js, masterKey, len(config.Nodes), bedrockPlatformDefaults)
// Bedrock self-host endpoints: Phase 1 models are pinned, so their
// OpenAI-compatible base URLs come from static config. OCHRE_VLLM_ENDPOINTS
// is a comma-separated list of modelId=baseURL pairs.
bedrockEndpoints := parseBedrockEndpoints(os.Getenv("OCHRE_VLLM_ENDPOINTS"))
gw := gateway.GatewayConfig{
Debug: nodeConfig.AWSGW.Debug,
DisableLogging: false,
NATSConn: natsConn,
Config: nodeConfig.AWSGW.Config,
ExpectedNodes: len(config.Nodes),
Region: nodeConfig.Region,
InternalSuffix: config.AWS.InternalSuffix,
RegistryPort: registryPort,
RegistryHost: registryHost,
AZ: nodeConfig.AZ,
IAMService: iamService,
STSService: stsService,
Version: version,
Commit: commit,
ECRRegistry: ecrRegistry,
ECRTokenIssuer: gateway_ecrauth.NewIssuer(signingKey, ecrAudience),
ECRTokenVerifier: gateway_ecrauth.NewVerifier(verifyKeys, ecrAudience),
Debug: nodeConfig.AWSGW.Debug,
DisableLogging: false,
NATSConn: natsConn,
Config: nodeConfig.AWSGW.Config,
ExpectedNodes: len(config.Nodes),
Region: nodeConfig.Region,
InternalSuffix: config.AWS.InternalSuffix,
RegistryPort: registryPort,
RegistryHost: registryHost,
AZ: nodeConfig.AZ,
IAMService: iamService,
STSService: stsService,
Version: version,
Commit: commit,
ECRRegistry: ecrRegistry,
ECRTokenIssuer: gateway_ecrauth.NewIssuer(signingKey, ecrAudience),
ECRTokenVerifier: gateway_ecrauth.NewVerifier(verifyKeys, ecrAudience),
BedrockCredentials: bedrockCredentials,
BedrockEndpoints: bedrockEndpoints,
}
// Rotate the ECR signing key on a 30-day cadence, retaining the previous keys
@@ -448,3 +465,28 @@ func findBootstrapFile(baseDir string) string {
func isConcreteRegistryHost(host string) bool {
return host != "" && host != "0.0.0.0" && host != "::"
}
// parseBedrockEndpoints parses a comma-separated list of modelId=baseURL pairs
// into a map for the bedrock self-host endpoint resolver. Malformed or empty
// entries are skipped. Returns nil for empty input.
func parseBedrockEndpoints(raw string) map[string]string {
raw = strings.TrimSpace(raw)
if raw == "" {
return nil
}
endpoints := map[string]string{}
for pair := range strings.SplitSeq(raw, ",") {
modelID, baseURL, ok := strings.Cut(pair, "=")
modelID = strings.TrimSpace(modelID)
baseURL = strings.TrimSpace(baseURL)
if !ok || modelID == "" || baseURL == "" {
slog.Warn("awsgw: skipping malformed OCHRE_VLLM_ENDPOINTS entry", "entry", pair)
continue
}
endpoints[modelID] = baseURL
}
if len(endpoints) == 0 {
return nil
}
return endpoints
}
@@ -0,0 +1,60 @@
package awsgw
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestParseBedrockEndpoints(t *testing.T) {
cases := []struct {
name string
raw string
want map[string]string
}{
{"empty input returns nil", "", nil},
{"whitespace-only input returns nil", " ", nil},
{
"single pair",
"meta.llama3-70b-instruct-v1:0=http://vllm.internal:8000",
map[string]string{"meta.llama3-70b-instruct-v1:0": "http://vllm.internal:8000"},
},
{
"multiple pairs",
"meta.llama3-70b-instruct-v1:0=http://vllm1:8000,mistral.mixtral-8x7b-v1:0=http://vllm2:8000",
map[string]string{
"meta.llama3-70b-instruct-v1:0": "http://vllm1:8000",
"mistral.mixtral-8x7b-v1:0": "http://vllm2:8000",
},
},
{
"whitespace around pairs and entries trimmed",
" meta.llama3-70b-instruct-v1:0 = http://vllm.internal:8000 , mistral.mixtral-8x7b-v1:0=http://vllm2:8000 ",
map[string]string{
"meta.llama3-70b-instruct-v1:0": "http://vllm.internal:8000",
"mistral.mixtral-8x7b-v1:0": "http://vllm2:8000",
},
},
{
"malformed entry with no equals sign is skipped",
"no-equals,meta.llama3-70b-instruct-v1:0=http://vllm.internal:8000",
map[string]string{"meta.llama3-70b-instruct-v1:0": "http://vllm.internal:8000"},
},
{
"malformed entry with empty modelId is skipped",
"=http://vllm.internal:8000,meta.llama3-70b-instruct-v1:0=http://vllm.internal:8000",
map[string]string{"meta.llama3-70b-instruct-v1:0": "http://vllm.internal:8000"},
},
{
"malformed entry with empty baseURL is skipped",
"meta.llama3-70b-instruct-v1:0=,mistral.mixtral-8x7b-v1:0=http://vllm2:8000",
map[string]string{"mistral.mixtral-8x7b-v1:0": "http://vllm2:8000"},
},
{"all entries malformed returns nil", "no-equals,=u,m=", nil},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, parseBedrockEndpoints(tc.raw))
})
}
}
+4 -3
View File
@@ -1,4 +1,4 @@
PKG ?= ./cert/... ./lb/... ./eks/... ./ecs/... ./iam/... ./ecr/... ./multinode/... ./baremetal/... ./reboot/... ./tofu/...
PKG ?= ./cert/... ./lb/... ./eks/... ./ecs/... ./iam/... ./ecr/... ./bedrock/... ./multinode/... ./baremetal/... ./reboot/... ./tofu/...
ARTIFACT_DIR ?= /tmp/spinifex-e2e-artifacts
TIMEOUT ?= 30m
MODE ?= single
@@ -11,9 +11,9 @@ SUITE_TAG ?= $(notdir $(patsubst %/...,%,$(firstword $(PKG))))
# named suite into BIN_DIR/<suite>.test so CI can ship binaries (not source +
# go toolchain) to the remote VM.
BIN_DIR ?= $(CURDIR)/_bin
SUITES_BUILD ?= cert lb eks ecs single iam ecr gpu natuplink
SUITES_BUILD ?= cert lb eks ecs single iam ecr bedrock gpu natuplink
.PHONY: e2e e2e-cert e2e-lb e2e-eks e2e-ecs e2e-single e2e-iam e2e-ecr e2e-multinode e2e-bm e2e-reboot e2e-tofu e2e-gpu e2e-natuplink e2e-all clean e2e-build
.PHONY: e2e e2e-cert e2e-lb e2e-eks e2e-ecs e2e-single e2e-iam e2e-ecr e2e-bedrock e2e-multinode e2e-bm e2e-reboot e2e-tofu e2e-gpu e2e-natuplink e2e-all clean e2e-build
ifeq ($(GOTESTSUM),1)
E2E_CMD = gotestsum \
@@ -39,6 +39,7 @@ e2e-ecs: ; $(MAKE) e2e PKG=./ecs/...
e2e-single: ; $(MAKE) e2e PKG=./single/...
e2e-iam: ; $(MAKE) e2e PKG=./iam/...
e2e-ecr: ; $(MAKE) e2e PKG=./ecr/...
e2e-bedrock: ; $(MAKE) e2e PKG=./bedrock/...
e2e-multinode: ; $(MAKE) e2e PKG=./multinode/... MODE=multinode
e2e-bm: ; $(MAKE) e2e PKG=./baremetal/... MODE=baremetal
e2e-reboot: ; $(MAKE) e2e PKG=./reboot/...
+130
View File
@@ -0,0 +1,130 @@
//go:build e2e
package bedrock
import (
"os"
"testing"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrock"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/tests/e2e/harness"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// Catalog model IDs mirror the static Phase-1 catalog in gateway_bedrock.
const (
modelSelfHostLlama = "meta.llama3-70b-instruct-v1:0"
modelAnthropic = "anthropic.claude-3-5-sonnet-20240620-v1:0"
modelBogus = "does.not-exist-v1:0"
)
// TestListFoundationModels confirms the self-host model is always advertised;
// the Anthropic model is only asserted on when present, since its advertisement
// depends on cluster credential configuration.
func TestListFoundationModels(t *testing.T) {
f := requireBedrockFixture(t)
out, err := f.AWS.Bedrock.ListFoundationModels(&bedrock.ListFoundationModelsInput{})
require.NoError(t, err, "list-foundation-models")
ids := make(map[string]*bedrock.FoundationModelSummary, len(out.ModelSummaries))
for _, m := range out.ModelSummaries {
ids[aws.StringValue(m.ModelId)] = m
}
assert.Contains(t, ids, modelSelfHostLlama, "self-host model must always be advertised")
if anthropic, ok := ids[modelAnthropic]; ok {
assert.Equal(t, "Anthropic", aws.StringValue(anthropic.ProviderName),
"anthropic model advertised with unexpected provider name")
}
}
// TestGetFoundationModel round-trips a known model and asserts
// ResourceNotFoundException for an unknown one.
func TestGetFoundationModel(t *testing.T) {
f := requireBedrockFixture(t)
t.Run("known model", func(t *testing.T) {
out, err := f.AWS.Bedrock.GetFoundationModel(&bedrock.GetFoundationModelInput{
ModelIdentifier: aws.String(modelSelfHostLlama),
})
require.NoError(t, err, "get-foundation-model %s", modelSelfHostLlama)
require.NotNil(t, out.ModelDetails, "empty ModelDetails")
assert.Equal(t, modelSelfHostLlama, aws.StringValue(out.ModelDetails.ModelId))
})
t.Run("unknown model", func(t *testing.T) {
harness.ExpectError(t, "ResourceNotFoundException", func() error {
_, e := f.AWS.Bedrock.GetFoundationModel(&bedrock.GetFoundationModelInput{
ModelIdentifier: aws.String(modelBogus),
})
return e
})
})
}
// TestConverseSelfHost exercises the real vLLM Converse path. Gated behind
// OCHRE_E2E_SELFHOST=1 — it needs a GPU-backed vLLM endpoint configured on the
// gateway via OCHRE_VLLM_ENDPOINTS at boot.
func TestConverseSelfHost(t *testing.T) {
if os.Getenv("OCHRE_E2E_SELFHOST") != "1" {
t.Skip("OCHRE_E2E_SELFHOST!=1; skipping real self-host Converse")
}
f := requireBedrockFixture(t)
out, err := f.AWS.BedrockRuntime.Converse(&bedrockruntime.ConverseInput{
ModelId: aws.String(modelSelfHostLlama),
Messages: []*bedrockruntime.Message{{
Role: aws.String(bedrockruntime.ConversationRoleUser),
Content: []*bedrockruntime.ContentBlock{{Text: aws.String("Say hello in one word.")}},
}},
InferenceConfig: &bedrockruntime.InferenceConfiguration{MaxTokens: aws.Int64(64)},
})
require.NoError(t, err, "converse self-host")
require.NotNil(t, out.Output, "nil output")
require.NotNil(t, out.Output.Message, "nil output message")
assert.NotEmpty(t, converseText(out.Output.Message), "empty assistant text")
require.NotNil(t, out.Usage, "nil usage")
assert.Greater(t, aws.Int64Value(out.Usage.OutputTokens), int64(0), "OutputTokens must be > 0")
}
// TestConverseAnthropic exercises the real Anthropic-direct Converse path.
// Gated behind OCHRE_E2E_ANTHROPIC=1 — it needs a live Anthropic API key
// configured on the cluster (OCHRE_ANTHROPIC_API_KEY or a per-account KV entry).
func TestConverseAnthropic(t *testing.T) {
if os.Getenv("OCHRE_E2E_ANTHROPIC") != "1" {
t.Skip("OCHRE_E2E_ANTHROPIC!=1; skipping real Anthropic Converse")
}
f := requireBedrockFixture(t)
out, err := f.AWS.BedrockRuntime.Converse(&bedrockruntime.ConverseInput{
ModelId: aws.String(modelAnthropic),
Messages: []*bedrockruntime.Message{{
Role: aws.String(bedrockruntime.ConversationRoleUser),
Content: []*bedrockruntime.ContentBlock{{Text: aws.String("Say hello in one word.")}},
}},
InferenceConfig: &bedrockruntime.InferenceConfiguration{MaxTokens: aws.Int64(64)},
})
require.NoError(t, err, "converse anthropic")
require.NotNil(t, out.Output, "nil output")
require.NotNil(t, out.Output.Message, "nil output message")
assert.NotEmpty(t, converseText(out.Output.Message), "empty assistant text")
require.NotNil(t, out.Usage, "nil usage")
assert.Greater(t, aws.Int64Value(out.Usage.InputTokens), int64(0), "InputTokens must be > 0")
assert.Greater(t, aws.Int64Value(out.Usage.OutputTokens), int64(0), "OutputTokens must be > 0")
}
// converseText concatenates the text content blocks of a Converse output message.
func converseText(msg *bedrockruntime.Message) string {
var text string
for _, c := range msg.Content {
if c != nil && c.Text != nil {
text += *c.Text
}
}
return text
}
+61
View File
@@ -0,0 +1,61 @@
//go:build e2e
// Package bedrock is the Ochre (AWS-Bedrock) E2E suite. Hermetic API-shape and
// provider-translation parity (Anthropic Messages / vLLM OpenAI wire formats)
// is covered by the gateway_bedrock unit tests; this suite exercises
// ListFoundationModels/GetFoundationModel against any running cluster and
// gates real Converse calls on backend availability (self-host GPU or a live
// Anthropic key), so CI parity holds without requiring either to be present.
package bedrock
import (
"os"
"sync"
"testing"
"github.com/mulgadc/spinifex/tests/e2e/harness"
)
// Package-scoped singleton fixture, built lazily by requireBedrockFixture and
// torn down by TestMain after the run.
var (
pkgFixOnce sync.Once
pkgFix *Fixture
pkgFixErr error
)
// Fixture carries per-process state shared across every Test* in this package.
type Fixture struct {
Env *harness.Env
AWS *harness.AWSClient
Account string
}
func TestMain(m *testing.M) {
os.Exit(m.Run())
}
// requireBedrockFixture returns the package-scoped Fixture, building it on
// first call. Skips the calling test if SPINIFEX_E2E is unset.
func requireBedrockFixture(t *testing.T) *Fixture {
t.Helper()
pkgFixOnce.Do(func() {
if os.Getenv("SPINIFEX_E2E") == "" {
return
}
env := harness.LoadEnv(t)
awsCli := harness.NewAWSClient(t, env)
pkgFix = &Fixture{
Env: env,
AWS: awsCli,
Account: harness.IAMAccountID(t, awsCli),
}
})
if pkgFixErr != nil {
t.Fatalf("bedrock fixture init failed: %v", pkgFixErr)
}
if pkgFix == nil {
t.Skip("bedrock fixture unavailable (SPINIFEX_E2E unset)")
}
return pkgFix
}
+24 -18
View File
@@ -14,6 +14,8 @@ import (
"github.com/aws/aws-sdk-go/aws/credentials"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/acm"
"github.com/aws/aws-sdk-go/service/bedrock"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
"github.com/aws/aws-sdk-go/service/ecr"
@@ -30,15 +32,17 @@ type AWSClient struct {
// EC2 is wrapped so RunInstances auto-serialises against cluster capacity
// (see capacityRetryEC2). EC2Conf is the underlying SDK client shared by
// both; kept for call sites that read or mutate its Config.
EC2 ec2iface.EC2API
EC2Conf *ec2.EC2
ELBv2 *elbv2.ELBV2
IAM *iam.IAM
STS *sts.STS
EKS *eks.EKS
ACM *acm.ACM
ECR *ecr.ECR
ECS *ecs.ECS
EC2 ec2iface.EC2API
EC2Conf *ec2.EC2
ELBv2 *elbv2.ELBV2
IAM *iam.IAM
STS *sts.STS
EKS *eks.EKS
ACM *acm.ACM
ECR *ecr.ECR
ECS *ecs.ECS
Bedrock *bedrock.Bedrock
BedrockRuntime *bedrockruntime.BedrockRuntime
}
// NewAWSClient builds clients pointed at the spinifex gateway using the
@@ -128,15 +132,17 @@ func newAWSClient(t *testing.T, env *Env, accessKey, secretKey, sessionToken str
rawEC2 := ec2.New(sess)
return &AWSClient{
EC2: &capacityRetryEC2{EC2API: rawEC2},
EC2Conf: rawEC2,
ELBv2: elbv2.New(sess),
IAM: iam.New(sess),
STS: sts.New(sess),
EKS: eks.New(sess),
ACM: acm.New(sess),
ECR: ecr.New(sess),
ECS: ecs.New(sess),
EC2: &capacityRetryEC2{EC2API: rawEC2},
EC2Conf: rawEC2,
ELBv2: elbv2.New(sess),
IAM: iam.New(sess),
STS: sts.New(sess),
EKS: eks.New(sess),
ACM: acm.New(sess),
ECR: ecr.New(sess),
ECS: ecs.New(sess),
Bedrock: bedrock.New(sess),
BedrockRuntime: bedrockruntime.New(sess),
}
}