From 8a493481ba2aae595de23bce766c31d32997d7df Mon Sep 17 00:00:00 2001 From: michaeljguarino Date: Thu, 21 May 2026 13:20:12 -0400 Subject: [PATCH] Support token exchange for bedrock This has the following needs: 1. Baseurl support - handled in reqllm but not bifrost 2. Custom tuning of stream enablement 3. Potentially more --- assets/src/generated/graphql.ts | 12 ++ go/client/models_gen.go | 12 ++ .../api/v1alpha1/deploymentsettings_types.go | 25 +++ .../api/v1alpha1/zz_generated.deepcopy.go | 15 ++ ...loyments.plural.sh_deploymentsettings.yaml | 49 +++++ .../api/v1alpha1/agentruntime_types.go | 29 ++- .../controller/agentrun_controller.go | 10 +- .../pkg/agentrun-harness/agentrun/v1/types.go | 18 +- .../agentrun-harness/tool/claude/claude.go | 2 + .../tool/claude/claude_streaming.go | 24 +++ .../pkg/agentrun-harness/tool/codex/codex.go | 15 +- .../tool/codex/codex_templates.go | 7 +- .../tool/codex/codex_types.go | 9 +- go/nexus/go.mod | 2 + go/nexus/go.sum | 4 +- go/nexus/internal/proto/console.pb.go | 67 ++++-- go/nexus/internal/router/account.go | 12 +- go/nexus/internal/router/account_keys.go | 37 +++- go/nexus/internal/router/account_keys_test.go | 137 ++++++++++++ .../internal/router/anthropic_messages.go | 1 + .../router/bedrock_allowed_requests.go | 23 +++ .../router/bedrock_allowed_requests_test.go | 27 +++ go/nexus/internal/router/oauth_token.go | 25 +++ go/nexus/internal/router/openai_chat.go | 5 +- go/nexus/internal/router/openai_responses.go | 5 +- go/nexus/internal/router/provider_base_url.go | 5 + go/nexus/internal/router/streaming_control.go | 53 +++++ .../internal/router/streaming_control_test.go | 45 ++++ .../internal/proto/console.pb.go | 195 ++++++++++++++---- lib/console/ai/provider/bedrock.ex | 71 ++++--- lib/console/graphql/deployments/settings.ex | 6 + lib/console/grpc/server.ex | 9 +- lib/console/schema/deployment_settings.ex | 7 +- lib/grpc/console.pb.ex | 3 + proto/console.proto | 3 + schema/schema.graphql | 18 ++ test/console/ai/provider/bedrock_test.exs | 142 +++++++++++++ 37 files changed, 1007 insertions(+), 122 deletions(-) create mode 100644 go/deployment-operator/pkg/agentrun-harness/tool/claude/claude_streaming.go create mode 100644 go/nexus/internal/router/bedrock_allowed_requests.go create mode 100644 go/nexus/internal/router/bedrock_allowed_requests_test.go create mode 100644 go/nexus/internal/router/oauth_token.go create mode 100644 go/nexus/internal/router/streaming_control.go create mode 100644 go/nexus/internal/router/streaming_control_test.go create mode 100644 test/console/ai/provider/bedrock_test.exs diff --git a/assets/src/generated/graphql.ts b/assets/src/generated/graphql.ts index 86b261b3cd..6ff6cbd050 100644 --- a/assets/src/generated/graphql.ts +++ b/assets/src/generated/graphql.ts @@ -1504,16 +1504,22 @@ export type BedrockAiAttributes = { awsAccessKeyId?: InputMaybe; /** the aws secret access key to use (DEPRECATED) */ awsSecretAccessKey?: InputMaybe; + /** the base url to use when querying a Bedrock-compatible API */ + baseUrl?: InputMaybe; /** Deprecated for most configurations: prefer regional-prefixed inference profile IDs in modelId or proxyModels (aliases are inferred automatically). Still needed for explicit client model name overrides, application inference profile resource IDs (profile suffix only, not full ARN), or when alias mapping cannot be inferred. Maps client-facing model ID to inference profile ID. Example: {"anthropic.claude-3-5-sonnet-20241022-v2:0": "us.anthropic.claude-3-5-sonnet-20241022-v2:0"} */ deployments?: InputMaybe; /** Bedrock model or inference profile for embeddings. Same ID formats as modelId. */ embeddingModel?: InputMaybe; + /** whether to enable streaming responses */ + enableStream?: InputMaybe; /** AWS Bedrock model or inference profile identifier. Use a foundation model ID (e.g. anthropic.claude-3-5-sonnet-20241022-v2:0) or a regional inference profile ID with three dot-separated segments (e.g. us.anthropic.claude-3-5-sonnet-20241022-v2:0, global.anthropic.claude-haiku-4-5-20251001-v1:0). Nexus registers the bare model ID for routing and auto-maps 3-part profile IDs to Bifrost aliases. */ modelId?: InputMaybe; /** Additional Bedrock model or inference profile IDs exposed through the Nexus OpenAI-compatible proxy beyond modelId, toolModelId, and embeddingModel. Same ID formats as modelId. */ proxyModels?: InputMaybe>>; /** the aws region the model is hosted in */ region?: InputMaybe; + /** OAuth2 client credentials against a token endpoint to obtain access tokens */ + tokenExchange?: InputMaybe; /** Bedrock model or inference profile for tool calls. Same ID formats as modelId. */ toolModelId?: InputMaybe; }; @@ -1523,16 +1529,22 @@ export type BedrockAiSettings = { __typename?: 'BedrockAiSettings'; /** the openai bedrock aws access key id to use (DEPRECATED) */ accessKeyId?: Maybe; + /** the base url to use when querying a Bedrock-compatible API */ + baseUrl?: Maybe; /** Deprecated for most configurations: prefer regional-prefixed inference profile IDs in modelId or proxyModels (aliases are inferred automatically). Still needed for explicit client model name overrides, application inference profile resource IDs (profile suffix only, not full ARN), or when alias mapping cannot be inferred. Maps client-facing model ID to inference profile ID. Example: {"anthropic.claude-3-5-sonnet-20241022-v2:0": "us.anthropic.claude-3-5-sonnet-20241022-v2:0"} */ deployments?: Maybe; /** Bedrock model or inference profile for embeddings. Same ID formats as modelId. */ embeddingModel?: Maybe; + /** whether streaming responses are enabled */ + enableStream?: Maybe; /** AWS Bedrock model or inference profile identifier. Use a foundation model ID (e.g. anthropic.claude-3-5-sonnet-20241022-v2:0) or a regional inference profile ID with three dot-separated segments (e.g. us.anthropic.claude-3-5-sonnet-20241022-v2:0, global.anthropic.claude-haiku-4-5-20251001-v1:0). Nexus registers the bare model ID for routing and auto-maps 3-part profile IDs to Bifrost aliases. Omit for Plural defaults. */ modelId?: Maybe; /** Additional Bedrock model or inference profile IDs exposed through the Nexus OpenAI-compatible proxy beyond modelId, toolModelId, and embeddingModel. Same ID formats as modelId. */ proxyModels?: Maybe>>; /** the aws region the model is hosted in */ region?: Maybe; + /** OAuth2 client credentials configured for token endpoint exchange */ + tokenExchange?: Maybe; /** Bedrock model or inference profile for tool calls. Same ID formats as modelId. */ toolModelId?: Maybe; }; diff --git a/go/client/models_gen.go b/go/client/models_gen.go index 5dc56aa603..5d13f5d8b2 100644 --- a/go/client/models_gen.go +++ b/go/client/models_gen.go @@ -1201,6 +1201,8 @@ type BedrockAiAttributes struct { ModelID *string `json:"modelId,omitempty"` // Bedrock model or inference profile for tool calls. Same ID formats as modelId. ToolModelID *string `json:"toolModelId,omitempty"` + // the base url to use when querying a Bedrock-compatible API + BaseURL *string `json:"baseUrl,omitempty"` // the openai bedrock access token to use AccessToken *string `json:"accessToken,omitempty"` // the aws region the model is hosted in @@ -1215,6 +1217,10 @@ type BedrockAiAttributes struct { ProxyModels []*string `json:"proxyModels,omitempty"` // Deprecated for most configurations: prefer regional-prefixed inference profile IDs in modelId or proxyModels (aliases are inferred automatically). Still needed for explicit client model name overrides, application inference profile resource IDs (profile suffix only, not full ARN), or when alias mapping cannot be inferred. Maps client-facing model ID to inference profile ID. Example: {"anthropic.claude-3-5-sonnet-20241022-v2:0": "us.anthropic.claude-3-5-sonnet-20241022-v2:0"} Deployments *string `json:"deployments,omitempty"` + // whether to enable streaming responses + EnableStream *bool `json:"enableStream,omitempty"` + // OAuth2 client credentials against a token endpoint to obtain access tokens + TokenExchange *OpenaiTokenExchangeAttributes `json:"tokenExchange,omitempty"` } // Settings for usage of AWS Bedrock for LLMs @@ -1223,6 +1229,8 @@ type BedrockAiSettings struct { ModelID *string `json:"modelId,omitempty"` // Bedrock model or inference profile for tool calls. Same ID formats as modelId. ToolModelID *string `json:"toolModelId,omitempty"` + // the base url to use when querying a Bedrock-compatible API + BaseURL *string `json:"baseUrl,omitempty"` // the openai bedrock aws access key id to use (DEPRECATED) AccessKeyID *string `json:"accessKeyId,omitempty"` // the aws region the model is hosted in @@ -1233,6 +1241,10 @@ type BedrockAiSettings struct { ProxyModels []*string `json:"proxyModels,omitempty"` // Deprecated for most configurations: prefer regional-prefixed inference profile IDs in modelId or proxyModels (aliases are inferred automatically). Still needed for explicit client model name overrides, application inference profile resource IDs (profile suffix only, not full ARN), or when alias mapping cannot be inferred. Maps client-facing model ID to inference profile ID. Example: {"anthropic.claude-3-5-sonnet-20241022-v2:0": "us.anthropic.claude-3-5-sonnet-20241022-v2:0"} Deployments map[string]any `json:"deployments,omitempty"` + // whether streaming responses are enabled + EnableStream *bool `json:"enableStream,omitempty"` + // OAuth2 client credentials configured for token endpoint exchange + TokenExchange *OpenaiTokenExchange `json:"tokenExchange,omitempty"` } type BindingAttributes struct { diff --git a/go/controller/api/v1alpha1/deploymentsettings_types.go b/go/controller/api/v1alpha1/deploymentsettings_types.go index e7d401a420..5c3cbef04e 100644 --- a/go/controller/api/v1alpha1/deploymentsettings_types.go +++ b/go/controller/api/v1alpha1/deploymentsettings_types.go @@ -749,6 +749,7 @@ func (in *AISettings) Attributes(ctx context.Context, c client.Client, namespace attr.Bedrock = &console.BedrockAiAttributes{ ModelID: in.Bedrock.ModelID, ToolModelID: in.Bedrock.ToolModelId, + BaseURL: in.Bedrock.BaseUrl, AccessToken: secret, Region: lo.ToPtr(in.Bedrock.Region), EmbeddingModel: in.Bedrock.EmbeddingModel, @@ -756,6 +757,14 @@ func (in *AISettings) Attributes(ctx context.Context, c client.Client, namespace AWSSecretAccessKey: secretKey, AWSAccessKeyID: in.Bedrock.AwsAccessKeyID, Deployments: deployments, + EnableStream: in.Bedrock.EnableStream, + } + if in.Bedrock.TokenExchange != nil { + tokenExchange, err := in.Bedrock.TokenExchange.Attributes(ctx, c, namespace) + if err != nil { + return nil, err + } + attr.Bedrock.TokenExchange = tokenExchange } } @@ -1047,6 +1056,22 @@ type BedrockSettings struct { // +kubebuilder:validation:Optional ToolModelId *string `json:"toolModelId,omitempty"` + // BaseUrl is the base url to use when querying a Bedrock-compatible API. + // + // +kubebuilder:validation:Optional + BaseUrl *string `json:"baseUrl,omitempty"` + + // EnableStream controls whether streaming responses are enabled. Set to false explicitly to disable, otherwise on. + // + // +kubebuilder:default=true + // +kubebuilder:validation:Optional + EnableStream *bool `json:"enableStream,omitempty"` + + // TokenExchange configures OAuth2 client credentials against a token endpoint to obtain access tokens. + // + // +kubebuilder:validation:Optional + TokenExchange *OAuth2TokenExchange `json:"tokenExchange,omitempty"` + // EmbeddingModel is the Bedrock model or inference profile for embeddings. Same ID formats as modelId. // // +kubebuilder:validation:Optional diff --git a/go/controller/api/v1alpha1/zz_generated.deepcopy.go b/go/controller/api/v1alpha1/zz_generated.deepcopy.go index bfa9f2670d..522052ce3c 100644 --- a/go/controller/api/v1alpha1/zz_generated.deepcopy.go +++ b/go/controller/api/v1alpha1/zz_generated.deepcopy.go @@ -444,6 +444,21 @@ func (in *BedrockSettings) DeepCopyInto(out *BedrockSettings) { *out = new(string) **out = **in } + if in.BaseUrl != nil { + in, out := &in.BaseUrl, &out.BaseUrl + *out = new(string) + **out = **in + } + if in.EnableStream != nil { + in, out := &in.EnableStream, &out.EnableStream + *out = new(bool) + **out = **in + } + if in.TokenExchange != nil { + in, out := &in.TokenExchange, &out.TokenExchange + *out = new(OAuth2TokenExchange) + (*in).DeepCopyInto(*out) + } if in.EmbeddingModel != nil { in, out := &in.EmbeddingModel, &out.EmbeddingModel *out = new(string) diff --git a/go/controller/config/crd/bases/deployments.plural.sh_deploymentsettings.yaml b/go/controller/config/crd/bases/deployments.plural.sh_deploymentsettings.yaml index 5b50d337fc..a6b245509e 100644 --- a/go/controller/config/crd/bases/deployments.plural.sh_deploymentsettings.yaml +++ b/go/controller/config/crd/bases/deployments.plural.sh_deploymentsettings.yaml @@ -235,6 +235,10 @@ spec: - key type: object x-kubernetes-map-type: atomic + baseUrl: + description: BaseUrl is the base url to use when querying + a Bedrock-compatible API. + type: string deployments: additionalProperties: type: string @@ -250,6 +254,11 @@ spec: description: EmbeddingModel is the Bedrock model or inference profile for embeddings. Same ID formats as modelId. type: string + enableStream: + default: true + description: EnableStream controls whether streaming responses + are enabled. + type: boolean modelId: description: |- ModelID is the primary AWS Bedrock model or inference profile identifier. @@ -267,6 +276,46 @@ spec: description: Region is the AWS region the model is hosted in type: string + tokenExchange: + description: TokenExchange configures OAuth2 client credentials + against a token endpoint to obtain access tokens. + properties: + clientId: + description: ClientID is the OAuth2 client identifier. + type: string + clientSecretSecretRef: + description: ClientSecretSecretRef is a reference to a + Kubernetes secret key holding the OAuth2 client secret. + properties: + key: + description: The key of the secret to select from. Must + be a valid secret key. + type: string + name: + default: "" + description: |- + Name of the referent. + This field is effectively required, but due to backwards compatibility is + allowed to be empty. Instances of this type with an empty value here are + almost certainly wrong. + More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names + type: string + optional: + description: Specify whether the Secret or its key + must be defined + type: boolean + required: + - key + type: object + x-kubernetes-map-type: atomic + enabled: + description: Enabled turns token exchange on for obtaining + access tokens via the configured token endpoint. + type: boolean + tokenUrl: + description: TokenURL is the OAuth2 token endpoint URL. + type: string + type: object tokenSecretRef: description: |- TokenSecretRef is a reference to the local secret holding the token to access diff --git a/go/deployment-operator/api/v1alpha1/agentruntime_types.go b/go/deployment-operator/api/v1alpha1/agentruntime_types.go index 4e76927a80..8be15937e8 100644 --- a/go/deployment-operator/api/v1alpha1/agentruntime_types.go +++ b/go/deployment-operator/api/v1alpha1/agentruntime_types.go @@ -315,6 +315,13 @@ type CodexConfigRaw struct { // Timeout bounds a single codex run invocation. // +kubebuilder:validation:Optional Timeout *metav1.Duration `json:"timeout,omitempty"` + + // DisableStream disables model provider API streaming when set to true explicitly. + // When true with aiProxy, the harness sets X-Plural-Enable-Stream: false on the + // plural model provider in config.toml. Codex has no direct env var; non-proxy + // mode may not apply. + // +kubebuilder:validation:Optional + DisableStream *bool `json:"disableStream,omitempty"` } type CodexConfig struct { @@ -333,6 +340,10 @@ type CodexConfig struct { // Timeout bounds a single codex run invocation. // +kubebuilder:validation:Optional Timeout *metav1.Duration `json:"timeout,omitempty"` + + // DisableStream disables model provider API streaming when set to true explicitly. + // +kubebuilder:validation:Optional + DisableStream *bool `json:"disableStream,omitempty"` } func (in *CodexConfig) ToCodexConfigRaw(secretGetter func(corev1.SecretKeySelector) (*corev1.Secret, error)) (*CodexConfigRaw, error) { @@ -341,9 +352,10 @@ func (in *CodexConfig) ToCodexConfigRaw(secretGetter func(corev1.SecretKeySelect } result := &CodexConfigRaw{ - Model: in.Model, - Endpoint: in.Endpoint, - Timeout: in.Timeout, + Model: in.Model, + Endpoint: in.Endpoint, + Timeout: in.Timeout, + DisableStream: in.DisableStream, } if !secretKeySelectorSet(in.ApiKeySecretRef) { @@ -393,6 +405,12 @@ type ClaudeConfig struct { // for a command before it is terminated. // +kubebuilder:validation:Optional BashMaxTimeout *metav1.Duration `json:"bashMaxTimeout,omitempty"` + + // DisableStream disables model provider API streaming when set to true explicitly. + // The harness sets CLAUDE_CODE_EXTRA_BODY to {"stream":false} in settings.local.json. + // With aiProxy, ANTHROPIC_CUSTOM_HEADERS also sends X-Plural-Enable-Stream: false. + // +kubebuilder:validation:Optional + DisableStream *bool `json:"disableStream,omitempty"` } // ClaudeConfigRaw contains configuration for the Claude CLI runtime. @@ -426,6 +444,10 @@ type ClaudeConfigRaw struct { // for a command before it is terminated. // +kubebuilder:validation:Optional BashMaxTimeout *metav1.Duration `json:"bashMaxTimeout,omitempty"` + + // DisableStream disables model provider API streaming when set to true explicitly. + // +kubebuilder:validation:Optional + DisableStream *bool `json:"disableStream,omitempty"` } func (in *ClaudeConfig) ToClaudeConfigRaw(secretGetter func(corev1.SecretKeySelector) (*corev1.Secret, error)) (*ClaudeConfigRaw, error) { @@ -440,6 +462,7 @@ func (in *ClaudeConfig) ToClaudeConfigRaw(secretGetter func(corev1.SecretKeySele Timeout: in.Timeout, BashTimeout: in.BashTimeout, BashMaxTimeout: in.BashMaxTimeout, + DisableStream: in.DisableStream, } if !secretKeySelectorSet(in.ApiKeySecretRef) { diff --git a/go/deployment-operator/internal/controller/agentrun_controller.go b/go/deployment-operator/internal/controller/agentrun_controller.go index 5e4e94844a..7e17b70a22 100644 --- a/go/deployment-operator/internal/controller/agentrun_controller.go +++ b/go/deployment-operator/internal/controller/agentrun_controller.go @@ -50,6 +50,7 @@ const ( EnvClaudeBashDefaultTimeout = "PLRL_CLAUDE_BASH_DEFAULT_TIMEOUT" EnvClaudeBashMaxTimeout = "PLRL_CLAUDE_BASH_MAX_TIMEOUT" EnvClaudeEndpoint = "PLRL_CLAUDE_ENDPOINT" + EnvClaudeDisableStream = "PLRL_CLAUDE_DISABLE_STREAM" EnvGeminiModel = "PLRL_GEMINI_MODEL" EnvGeminiAPIKey = "PLRL_GEMINI_API_KEY" @@ -58,7 +59,8 @@ const ( EnvCodexModel = "PLRL_CODEX_MODEL" EnvCodexAPIKey = "PLRL_CODEX_API_KEY" - EnvCodexEndpoint = "PLRL_CODEX_ENDPOINT" + EnvCodexEndpoint = "PLRL_CODEX_ENDPOINT" + EnvCodexDisableStream = "PLRL_CODEX_DISABLE_STREAM" EnvDindEnabled = "PLRL_DIND_ENABLED" EnvBrowserEnabled = "PLRL_BROWSER_ENABLED" @@ -523,6 +525,9 @@ func (r *AgentRunReconciler) getSecretData(run *v1alpha1.AgentRun, config *v1alp if config.Claude.Endpoint != nil { result[EnvClaudeEndpoint] = lo.FromPtr(config.Claude.Endpoint) } + if config.Claude.DisableStream != nil && *config.Claude.DisableStream { + result[EnvClaudeDisableStream] = "true" + } } if runtimeType == console.AgentRuntimeTypeGemini { @@ -555,6 +560,9 @@ func (r *AgentRunReconciler) getSecretData(run *v1alpha1.AgentRun, config *v1alp if config.Codex.Endpoint != nil { result[EnvCodexEndpoint] = lo.FromPtr(config.Codex.Endpoint) } + if config.Codex.DisableStream != nil && *config.Codex.DisableStream { + result[EnvCodexDisableStream] = "true" + } } return result diff --git a/go/deployment-operator/pkg/agentrun-harness/agentrun/v1/types.go b/go/deployment-operator/pkg/agentrun-harness/agentrun/v1/types.go index c4fd8aa951..dc6da0991f 100644 --- a/go/deployment-operator/pkg/agentrun-harness/agentrun/v1/types.go +++ b/go/deployment-operator/pkg/agentrun-harness/agentrun/v1/types.go @@ -92,6 +92,7 @@ type ClaudeConfig struct { BashTimeout time.Duration `json:"bashTimeout"` BashMaxTimeout time.Duration `json:"bashMaxTimeout"` Endpoint *string `json:"endpoint,omitempty"` + DisableStream bool `json:"disableStream,omitempty"` } type GeminiConfig struct { @@ -103,10 +104,11 @@ type GeminiConfig struct { } type CodexConfig struct { - ApiKey string `json:"apiKey"` - Model string `json:"model,omitempty"` - Timeout time.Duration `json:"timeout"` - Endpoint *string `json:"endpoint,omitempty"` + ApiKey string `json:"apiKey"` + Model string `json:"model,omitempty"` + Timeout time.Duration `json:"timeout"` + Endpoint *string `json:"endpoint,omitempty"` + DisableStream bool `json:"disableStream,omitempty"` } // FromAgentRunFragment converts Console API fragment to harness type @@ -175,6 +177,7 @@ func (ar *AgentRun) fromEnv(runtime *console.AgentRuntimeFragment) *AgentRuntime Timeout: helpers.GetPluralEnvDuration(controller.EnvExecTimeout, defaultTimeout), BashTimeout: helpers.GetPluralEnvDuration(controller.EnvClaudeBashDefaultTimeout, defaultBashTimeout), BashMaxTimeout: helpers.GetPluralEnvDuration(controller.EnvClaudeBashMaxTimeout, defaultBashMaxTimeout), + DisableStream: helpers.GetPluralEnvBool(controller.EnvClaudeDisableStream, false), } if endpoint := helpers.GetPluralEnv(controller.EnvClaudeEndpoint, ""); endpoint != "" { config.Claude.Endpoint = &endpoint @@ -200,9 +203,10 @@ func (ar *AgentRun) fromEnv(runtime *console.AgentRuntimeFragment) *AgentRuntime } case console.AgentRuntimeTypeCodex: config.Codex = &CodexConfig{ - ApiKey: helpers.GetPluralEnv(controller.EnvCodexAPIKey, ""), - Model: helpers.GetPluralEnv(controller.EnvCodexModel, ""), - Timeout: helpers.GetPluralEnvDuration(controller.EnvExecTimeout, defaultTimeout), + ApiKey: helpers.GetPluralEnv(controller.EnvCodexAPIKey, ""), + Model: helpers.GetPluralEnv(controller.EnvCodexModel, ""), + Timeout: helpers.GetPluralEnvDuration(controller.EnvExecTimeout, defaultTimeout), + DisableStream: helpers.GetPluralEnvBool(controller.EnvCodexDisableStream, false), } if endpoint := helpers.GetPluralEnv(controller.EnvCodexEndpoint, ""); endpoint != "" { config.Codex.Endpoint = &endpoint diff --git a/go/deployment-operator/pkg/agentrun-harness/tool/claude/claude.go b/go/deployment-operator/pkg/agentrun-harness/tool/claude/claude.go index 7bcdca5669..6494fc78f1 100644 --- a/go/deployment-operator/pkg/agentrun-harness/tool/claude/claude.go +++ b/go/deployment-operator/pkg/agentrun-harness/tool/claude/claude.go @@ -286,6 +286,7 @@ func (in *Claude) ConfigureBabysitRun() error { settings.WithEnv("BASH_DEFAULT_TIMEOUT_MS", defaultTimeout) settings.WithEnv("BASH_MAX_TIMEOUT_MS", maxTimeout) klog.V(log.LogLevelInfo).InfoS("claude timeouts configured", "default_timeout", defaultTimeout, "max_timeout", maxTimeout) + in.applyProviderStreamingSettings(settings) return settings.WriteToFile(filepath.Join(in.configPath(), "settings.local.json")) } @@ -374,6 +375,7 @@ func (in *Claude) Configure(consoleURL, consoleToken, _ string) error { settings.WithEnv("BASH_DEFAULT_TIMEOUT_MS", defaultTimeout) settings.WithEnv("BASH_MAX_TIMEOUT_MS", maxTimeout) klog.V(log.LogLevelInfo).InfoS("claude timeouts configured", "default_timeout", defaultTimeout, "max_timeout", maxTimeout) + in.applyProviderStreamingSettings(settings) return settings.WriteToFile(filepath.Join(in.configPath(), "settings.local.json")) } diff --git a/go/deployment-operator/pkg/agentrun-harness/tool/claude/claude_streaming.go b/go/deployment-operator/pkg/agentrun-harness/tool/claude/claude_streaming.go new file mode 100644 index 0000000000..7f837d4363 --- /dev/null +++ b/go/deployment-operator/pkg/agentrun-harness/tool/claude/claude_streaming.go @@ -0,0 +1,24 @@ +package claude + +// PluralDisableStreamHeader is sent to the Console AI proxy when provider streaming is disabled. +const PluralDisableStreamHeader = "X-Plural-Enable-Stream: false" + +const claudeExtraBodyDisableStream = `{"stream":false}` + +func (in *Claude) providerStreamingDisabled() bool { + if in.Config.Run.Runtime == nil || in.Config.Run.Runtime.Config == nil || in.Config.Run.Runtime.Config.Claude == nil { + return false + } + return in.Config.Run.Runtime.Config.Claude.DisableStream +} + +func (in *Claude) applyProviderStreamingSettings(settings *SettingsBuilder) { + if !in.providerStreamingDisabled() { + return + } + + settings.WithEnv("CLAUDE_CODE_EXTRA_BODY", claudeExtraBodyDisableStream) + if in.Config.Run.IsProxyEnabled() { + settings.WithEnv("ANTHROPIC_CUSTOM_HEADERS", PluralDisableStreamHeader) + } +} diff --git a/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex.go b/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex.go index d33abf1df2..99b73f4ffe 100644 --- a/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex.go +++ b/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex.go @@ -57,6 +57,13 @@ func (in *Codex) ensure() error { return nil } +func (in *Codex) providerStreamingDisabled() bool { + if in.Config.Run.Runtime == nil || in.Config.Run.Runtime.Config == nil || in.Config.Run.Runtime.Config.Codex == nil { + return false + } + return in.Config.Run.Runtime.Config.Codex.DisableStream +} + func (in *Codex) Run(ctx context.Context, options ...exec.Option) { go in.start(ctx, options...) } @@ -99,11 +106,15 @@ func (in *Codex) Configure(consoleURL, consoleToken, deployToken string) error { modelProvider := "" if in.proxy { modelProvider = "plural" - providers = []ModelProviderInput{{ + provider := ModelProviderInput{ Name: "plural", BaseURL: fmt.Sprintf("%s/ext/ai/v1", consoleURL), EnvKey: consoleTokenEnv, - }} + } + if in.providerStreamingDisabled() { + provider.HttpHeaders = map[string]string{"X-Plural-Enable-Stream": "false"} + } + providers = []ModelProviderInput{provider} } else if in.Config.Run.Runtime.Config.Codex.Endpoint != nil { modelProvider = "custom" providers = []ModelProviderInput{{ diff --git a/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex_templates.go b/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex_templates.go index 7df61a9e50..3cc550cacc 100644 --- a/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex_templates.go +++ b/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex_templates.go @@ -34,9 +34,10 @@ func BuildCodexConfig(repositoryDir string, agents []AgentInput, mcps []MCPInput cfg.ModelProviders = make(map[string]*ModelProviderConfig, len(providers)) for _, p := range providers { cfg.ModelProviders[p.Name] = &ModelProviderConfig{ - Name: p.Name, - BaseURL: p.BaseURL, - EnvKey: p.EnvKey, + Name: p.Name, + BaseURL: p.BaseURL, + EnvKey: p.EnvKey, + HttpHeaders: p.HttpHeaders, } } } diff --git a/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex_types.go b/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex_types.go index 98d0f8ab08..ecf5e689d1 100644 --- a/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex_types.go +++ b/go/deployment-operator/pkg/agentrun-harness/tool/codex/codex_types.go @@ -204,13 +204,16 @@ type ModelProviderInput struct { BaseURL string // EnvKey is the name of the environment variable that holds the API key. EnvKey string + // HttpHeaders are sent on each request to this provider (for example X-Plural-Enable-Stream). + HttpHeaders map[string]string } // ModelProviderConfig is serialized into [model_providers.] in config.toml. type ModelProviderConfig struct { - Name string `toml:"name,omitempty"` - BaseURL string `toml:"base_url,omitempty"` - EnvKey string `toml:"env_key,omitempty"` + Name string `toml:"name,omitempty"` + BaseURL string `toml:"base_url,omitempty"` + EnvKey string `toml:"env_key,omitempty"` + HttpHeaders map[string]string `toml:"http_headers,omitempty"` } type ShellEnvPolicy struct { diff --git a/go/nexus/go.mod b/go/nexus/go.mod index c718477318..0ebc3716c5 100644 --- a/go/nexus/go.mod +++ b/go/nexus/go.mod @@ -93,3 +93,5 @@ require ( google.golang.org/genproto/googleapis/rpc v0.0.0-20260401024825-9d38bb4040a9 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) + +replace github.com/maximhq/bifrost/core => github.com/pluralsh/bifrost/core v0.0.0-20260521165326-8d3ebe79c07e diff --git a/go/nexus/go.sum b/go/nexus/go.sum index fdd4018965..d6647d5e98 100644 --- a/go/nexus/go.sum +++ b/go/nexus/go.sum @@ -123,8 +123,8 @@ github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/ github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= -github.com/maximhq/bifrost/core v1.5.10 h1:xC+i5KlrRE+OK5UDCuaI7s4lq7mwJaDR2+/H+vVdfE0= -github.com/maximhq/bifrost/core v1.5.10/go.mod h1:WX3sHkss9Lk2M+uzVKStd52sw7/bXDFx48Iv83LdjIQ= +github.com/pluralsh/bifrost/core v0.0.0-20260521165326-8d3ebe79c07e h1:VqJhHNH/ZPbZd1tP9SrQ1Ep3dMSIdMOgHmN/MAG62J8= +github.com/pluralsh/bifrost/core v0.0.0-20260521165326-8d3ebe79c07e/go.mod h1:bzcFxKjI2NCRtOeQOwA0KxKJkjsY/UgnXd/cmg7Q+4U= github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4= github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ= diff --git a/go/nexus/internal/proto/console.pb.go b/go/nexus/internal/proto/console.pb.go index ae4af9bfd9..79644c945f 100644 --- a/go/nexus/internal/proto/console.pb.go +++ b/go/nexus/internal/proto/console.pb.go @@ -565,6 +565,9 @@ type BedrockConfig struct { AwsSecretAccessKey *string `protobuf:"bytes,7,opt,name=awsSecretAccessKey,proto3,oneof" json:"awsSecretAccessKey,omitempty"` ProxyModels []string `protobuf:"bytes,8,rep,name=proxyModels,proto3" json:"proxyModels,omitempty"` Deployments map[string]string `protobuf:"bytes,9,rep,name=deployments,proto3" json:"deployments,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + BaseUrl *string `protobuf:"bytes,10,opt,name=baseUrl,proto3,oneof" json:"baseUrl,omitempty"` + TokenExchange *OpenAiTokenExchange `protobuf:"bytes,11,opt,name=tokenExchange,proto3,oneof" json:"tokenExchange,omitempty"` + EnableStream *bool `protobuf:"varint,12,opt,name=enableStream,proto3,oneof" json:"enableStream,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -662,6 +665,27 @@ func (x *BedrockConfig) GetDeployments() map[string]string { return nil } +func (x *BedrockConfig) GetBaseUrl() string { + if x != nil && x.BaseUrl != nil { + return *x.BaseUrl + } + return "" +} + +func (x *BedrockConfig) GetTokenExchange() *OpenAiTokenExchange { + if x != nil { + return x.TokenExchange + } + return nil +} + +func (x *BedrockConfig) GetEnableStream() bool { + if x != nil && x.EnableStream != nil { + return *x.EnableStream + } + return false +} + type AzureOpenAiConfig struct { state protoimpl.MessageState `protogen:"open.v1"` ApiVersion *string `protobuf:"bytes,1,opt,name=apiVersion,proto3,oneof" json:"apiVersion,omitempty"` @@ -1140,7 +1164,7 @@ const file_console_proto_rawDesc = "" + "_toolModelB\n" + "\n" + "\b_projectB\v\n" + - "\t_location\"\xcc\x04\n" + + "\t_location\"\x89\x06\n" + "\rBedrockConfig\x12\x1d\n" + "\amodelId\x18\x01 \x01(\tH\x00R\amodelId\x88\x01\x01\x12%\n" + "\vtoolModelId\x18\x02 \x01(\tH\x01R\vtoolModelId\x88\x01\x01\x12%\n" + @@ -1150,7 +1174,11 @@ const file_console_proto_rawDesc = "" + "\x0eawsAccessKeyId\x18\x06 \x01(\tH\x05R\x0eawsAccessKeyId\x88\x01\x01\x123\n" + "\x12awsSecretAccessKey\x18\a \x01(\tH\x06R\x12awsSecretAccessKey\x88\x01\x01\x12 \n" + "\vproxyModels\x18\b \x03(\tR\vproxyModels\x12F\n" + - "\vdeployments\x18\t \x03(\v2$.plrl.BedrockConfig.DeploymentsEntryR\vdeployments\x1a>\n" + + "\vdeployments\x18\t \x03(\v2$.plrl.BedrockConfig.DeploymentsEntryR\vdeployments\x12\x1d\n" + + "\abaseUrl\x18\n" + + " \x01(\tH\aR\abaseUrl\x88\x01\x01\x12D\n" + + "\rtokenExchange\x18\v \x01(\v2\x19.plrl.OpenAiTokenExchangeH\bR\rtokenExchange\x88\x01\x01\x12'\n" + + "\fenableStream\x18\f \x01(\bH\tR\fenableStream\x88\x01\x01\x1a>\n" + "\x10DeploymentsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01B\n" + @@ -1161,7 +1189,11 @@ const file_console_proto_rawDesc = "" + "\a_regionB\x13\n" + "\x11_embeddingModelIdB\x11\n" + "\x0f_awsAccessKeyIdB\x15\n" + - "\x13_awsSecretAccessKey\"\xf0\x03\n" + + "\x13_awsSecretAccessKeyB\n" + + "\n" + + "\b_baseUrlB\x10\n" + + "\x0e_tokenExchangeB\x0f\n" + + "\r_enableStream\"\xf0\x03\n" + "\x11AzureOpenAiConfig\x12#\n" + "\n" + "apiVersion\x18\x01 \x01(\tH\x00R\n" + @@ -1260,20 +1292,21 @@ var file_console_proto_depIdxs = []int32{ 3, // 5: plrl.OpenAiConfig.tokenExchange:type_name -> plrl.OpenAiTokenExchange 0, // 6: plrl.OpenAiConfig.method:type_name -> plrl.OpenAiMethod 15, // 7: plrl.BedrockConfig.deployments:type_name -> plrl.BedrockConfig.DeploymentsEntry - 16, // 8: plrl.AzureOpenAiConfig.deployments:type_name -> plrl.AzureOpenAiConfig.DeploymentsEntry - 12, // 9: plrl.PluralServer.MeterMetrics:input_type -> plrl.MeterMetricsRequest - 1, // 10: plrl.PluralServer.GetAiConfig:input_type -> plrl.AiConfigRequest - 14, // 11: plrl.PluralServer.GetObservabilityConfig:input_type -> plrl.ObservabilityConfigRequest - 9, // 12: plrl.PluralServer.ProxyAuthentication:input_type -> plrl.ProxyAuthenticationRequest - 13, // 13: plrl.PluralServer.MeterMetrics:output_type -> plrl.MeterMetricsResponse - 2, // 14: plrl.PluralServer.GetAiConfig:output_type -> plrl.AiConfig - 11, // 15: plrl.PluralServer.GetObservabilityConfig:output_type -> plrl.ObservabilityConfig - 10, // 16: plrl.PluralServer.ProxyAuthentication:output_type -> plrl.ProxyAuthenticationResponse - 13, // [13:17] is the sub-list for method output_type - 9, // [9:13] is the sub-list for method input_type - 9, // [9:9] is the sub-list for extension type_name - 9, // [9:9] is the sub-list for extension extendee - 0, // [0:9] is the sub-list for field type_name + 3, // 8: plrl.BedrockConfig.tokenExchange:type_name -> plrl.OpenAiTokenExchange + 16, // 9: plrl.AzureOpenAiConfig.deployments:type_name -> plrl.AzureOpenAiConfig.DeploymentsEntry + 12, // 10: plrl.PluralServer.MeterMetrics:input_type -> plrl.MeterMetricsRequest + 1, // 11: plrl.PluralServer.GetAiConfig:input_type -> plrl.AiConfigRequest + 14, // 12: plrl.PluralServer.GetObservabilityConfig:input_type -> plrl.ObservabilityConfigRequest + 9, // 13: plrl.PluralServer.ProxyAuthentication:input_type -> plrl.ProxyAuthenticationRequest + 13, // 14: plrl.PluralServer.MeterMetrics:output_type -> plrl.MeterMetricsResponse + 2, // 15: plrl.PluralServer.GetAiConfig:output_type -> plrl.AiConfig + 11, // 16: plrl.PluralServer.GetObservabilityConfig:output_type -> plrl.ObservabilityConfig + 10, // 17: plrl.PluralServer.ProxyAuthentication:output_type -> plrl.ProxyAuthenticationResponse + 14, // [14:18] is the sub-list for method output_type + 10, // [10:14] is the sub-list for method input_type + 10, // [10:10] is the sub-list for extension type_name + 10, // [10:10] is the sub-list for extension extendee + 0, // [0:10] is the sub-list for field type_name } func init() { file_console_proto_init() } diff --git a/go/nexus/internal/router/account.go b/go/nexus/internal/router/account.go index 8bd75ad85a..42ce6c8a96 100644 --- a/go/nexus/internal/router/account.go +++ b/go/nexus/internal/router/account.go @@ -153,7 +153,17 @@ func (in *Account) GetConfigForProvider(provider schemas.ModelProvider) (*schema // Vertex uses project/location + auth in keys; no base URL override required. case schemas.Bedrock: - // Bedrock uses AWS region + credentials in keys; no base URL override required. + if cfg := aiConfig.GetBedrock(); cfg != nil { + if cfg.GetBaseUrl() != "" { + config.NetworkConfig.BaseURL = bedrockNetworkBaseURL(cfg.GetBaseUrl()) + } + if allowed := bedrockAllowedRequests(cfg); allowed != nil { + config.CustomProviderConfig = &schemas.CustomProviderConfig{ + BaseProviderType: schemas.Bedrock, + AllowedRequests: allowed, + } + } + } case schemas.Azure: // Azure uses endpoint + token in keys; no base URL override required. diff --git a/go/nexus/internal/router/account_keys.go b/go/nexus/internal/router/account_keys.go index 4aeb4f7472..214b02582b 100644 --- a/go/nexus/internal/router/account_keys.go +++ b/go/nexus/internal/router/account_keys.go @@ -30,7 +30,7 @@ func (in *Account) GetKeysForProvider(ctx context.Context, provider schemas.Mode case schemas.Vertex: return in.handleVertexKeys(aiConfig.GetVertexAi()) case schemas.Bedrock: - return in.handleBedrockKeys(aiConfig.GetBedrock()) + return in.handleBedrockKeys(ctx, aiConfig.GetBedrock()) case schemas.Azure: return in.handleAzureKeys(aiConfig.GetAzure()) default: @@ -73,15 +73,10 @@ func (in *Account) handleOpenAIKeys(ctx context.Context, config *pb.OpenAiConfig // token (when tokenExchange is enabled in the Console gRPC schema), the static apiKey, or empty when // requests are authenticated by an upstream internal proxy (e.g. Plural Cloud) instead. func (in *Account) openAIAPIKey(ctx context.Context, config *pb.OpenAiConfig) (string, error) { - tx := config.GetTokenExchange() - if tx != nil && tx.GetEnabled() { - tokenURL := tx.GetTokenUrl() - clientID := tx.GetClientId() - clientSecret := tx.GetClientSecret() - if tokenURL == "" || clientID == "" || clientSecret == "" { - return "", fmt.Errorf("openai tokenExchange enabled but tokenUrl, clientId, and clientSecret must all be set") - } - return in.tokenCache.AccessToken(ctx, tokenURL, clientID, clientSecret) + if token, err := in.oauthAccessToken(ctx, config.GetTokenExchange(), "openai"); err != nil { + return "", err + } else if token != "" { + return token, nil } return config.GetApiKey(), nil } @@ -146,7 +141,7 @@ func (in *Account) handleVertexKeys(config *pb.VertexAiConfig) ([]schemas.Key, e }, nil } -func (in *Account) handleBedrockKeys(config *pb.BedrockConfig) ([]schemas.Key, error) { +func (in *Account) handleBedrockKeys(ctx context.Context, config *pb.BedrockConfig) ([]schemas.Key, error) { if config == nil { return nil, fmt.Errorf("bedrock not configured") } @@ -155,10 +150,19 @@ func (in *Account) handleBedrockKeys(config *pb.BedrockConfig) ([]schemas.Key, e zap.String("model", config.GetModelId()), zap.String("tool_model", config.GetToolModelId()), zap.String("embedding_model", config.GetEmbeddingModelId()), + zap.String("base_url", config.GetBaseUrl()), ) + apiKey, err := in.bedrockAPIKey(ctx, config) + if err != nil { + return nil, err + } + return []schemas.Key{ { + Value: schemas.EnvVar{ + Val: apiKey, + }, Models: in.toBedrockModels(config), Aliases: schemas.KeyAliases(in.toBedrockDeployments(config)), BedrockKeyConfig: &schemas.BedrockKeyConfig{ @@ -178,6 +182,17 @@ func (in *Account) handleBedrockKeys(config *pb.BedrockConfig) ([]schemas.Key, e }, nil } +// bedrockAPIKey returns the bearer credential for Bedrock-compatible APIs: OAuth2 token exchange, +// a static access token, or empty when AWS IAM credentials on the key are used instead. +func (in *Account) bedrockAPIKey(ctx context.Context, config *pb.BedrockConfig) (string, error) { + if token, err := in.oauthAccessToken(ctx, config.GetTokenExchange(), "bedrock"); err != nil { + return "", err + } else if token != "" { + return token, nil + } + return config.GetAccessToken(), nil +} + // toBedrockModels returns client-facing model IDs registered on the Bifrost key. // Configured values may be foundation model IDs (e.g. anthropic.claude-3-5-sonnet-20241022-v2:0) // or regional inference profile IDs with three dot-separated segments (e.g. diff --git a/go/nexus/internal/router/account_keys_test.go b/go/nexus/internal/router/account_keys_test.go index 0f6464ab2f..1e562b846a 100644 --- a/go/nexus/internal/router/account_keys_test.go +++ b/go/nexus/internal/router/account_keys_test.go @@ -117,3 +117,140 @@ func TestHandleOpenAIKeys_tokenExchangeEnabledIncomplete(t *testing.T) { _, err := acct.GetKeysForProvider(context.Background(), schemas.OpenAI) require.Error(t, err) } + +func TestHandleBedrockKeys_accessTokenStatic(t *testing.T) { + cfg := &pb.AiConfig{ + Enabled: true, + Bedrock: &pb.BedrockConfig{ + ModelId: lo.ToPtr("anthropic.claude-3-5-sonnet-20241022-v2:0"), + Region: lo.ToPtr("us-east-1"), + AccessToken: lo.ToPtr("static-bedrock-token"), + }, + } + acct := &Account{ + consoleClient: &mockConsoleClient{cfg: cfg}, + tokenCache: tokenexchange.NewCache(), + logger: zap.NewNop(), + } + + keys, err := acct.GetKeysForProvider(context.Background(), schemas.Bedrock) + require.NoError(t, err) + require.Len(t, keys, 1) + require.Equal(t, "static-bedrock-token", keys[0].Value.Val) +} + +func TestHandleBedrockKeys_noBearerUsesIAMOnly(t *testing.T) { + cfg := &pb.AiConfig{ + Enabled: true, + Bedrock: &pb.BedrockConfig{ + ModelId: lo.ToPtr("anthropic.claude-3-5-sonnet-20241022-v2:0"), + Region: lo.ToPtr("us-east-1"), + AwsAccessKeyId: lo.ToPtr("AKIA"), + }, + } + acct := &Account{ + consoleClient: &mockConsoleClient{cfg: cfg}, + tokenCache: tokenexchange.NewCache(), + logger: zap.NewNop(), + } + + keys, err := acct.GetKeysForProvider(context.Background(), schemas.Bedrock) + require.NoError(t, err) + require.Len(t, keys, 1) + require.Empty(t, keys[0].Value.Val) + require.Equal(t, "AKIA", keys[0].BedrockKeyConfig.AccessKey.Val) +} + +func TestHandleBedrockKeys_tokenExchangeUsesCache(t *testing.T) { + var tokenCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + tokenCalls.Add(1) + _ = r.ParseForm() + if r.FormValue("grant_type") != "client_credentials" { + http.Error(w, "bad grant", http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"bedrock-oauth-at","expires_in":3600}`)) + })) + t.Cleanup(srv.Close) + + tokenURL := srv.URL + "/oauth/token" + cfg := &pb.AiConfig{ + Enabled: true, + Bedrock: &pb.BedrockConfig{ + ModelId: lo.ToPtr("anthropic.claude-3-5-sonnet-20241022-v2:0"), + Region: lo.ToPtr("us-east-1"), + TokenExchange: &pb.OpenAiTokenExchange{ + Enabled: lo.ToPtr(true), + TokenUrl: lo.ToPtr(tokenURL), + ClientId: lo.ToPtr("id"), + ClientSecret: lo.ToPtr("secret"), + }, + }, + } + + acct := &Account{ + consoleClient: &mockConsoleClient{cfg: cfg}, + tokenCache: tokenexchange.NewCacheWithHTTPClient(srv.Client()), + logger: zap.NewNop(), + } + ctx := context.Background() + + keys1, err := acct.GetKeysForProvider(ctx, schemas.Bedrock) + require.NoError(t, err) + require.Len(t, keys1, 1) + require.Equal(t, "bedrock-oauth-at", keys1[0].Value.Val) + c1 := tokenCalls.Load() + require.GreaterOrEqual(t, c1, int32(1)) + + keys2, err := acct.GetKeysForProvider(ctx, schemas.Bedrock) + require.NoError(t, err) + require.Equal(t, "bedrock-oauth-at", keys2[0].Value.Val) + require.Equal(t, c1, tokenCalls.Load(), "second GetKeysForProvider should use cached OAuth token") +} + +func TestHandleBedrockKeys_tokenExchangeEnabledIncomplete(t *testing.T) { + cfg := &pb.AiConfig{ + Enabled: true, + Bedrock: &pb.BedrockConfig{ + ModelId: lo.ToPtr("anthropic.claude-3-5-sonnet-20241022-v2:0"), + Region: lo.ToPtr("us-east-1"), + TokenExchange: &pb.OpenAiTokenExchange{ + Enabled: lo.ToPtr(true), + ClientId: lo.ToPtr("only-id"), + }, + }, + } + acct := &Account{ + consoleClient: &mockConsoleClient{cfg: cfg}, + tokenCache: tokenexchange.NewCache(), + logger: zap.NewNop(), + } + _, err := acct.GetKeysForProvider(context.Background(), schemas.Bedrock) + require.Error(t, err) +} + +func TestGetConfigForProvider_bedrockBaseURLAndDisableStream(t *testing.T) { + cfg := &pb.AiConfig{ + Enabled: true, + Bedrock: &pb.BedrockConfig{ + ModelId: lo.ToPtr("anthropic.claude-3-5-sonnet-20241022-v2:0"), + Region: lo.ToPtr("us-east-1"), + BaseUrl: lo.ToPtr("https://bedrock-proxy.example.com/"), + EnableStream: lo.ToPtr(false), + }, + } + acct := &Account{ + consoleClient: &mockConsoleClient{cfg: cfg}, + tokenCache: tokenexchange.NewCache(), + logger: zap.NewNop(), + } + + providerCfg, err := acct.GetConfigForProvider(schemas.Bedrock) + require.NoError(t, err) + require.Equal(t, "https://bedrock-proxy.example.com", providerCfg.NetworkConfig.BaseURL) + require.NotNil(t, providerCfg.CustomProviderConfig) + require.False(t, providerCfg.CustomProviderConfig.AllowedRequests.ChatCompletionStream) + require.False(t, providerCfg.CustomProviderConfig.AllowedRequests.ResponsesStream) +} diff --git a/go/nexus/internal/router/anthropic_messages.go b/go/nexus/internal/router/anthropic_messages.go index ec66c74480..b4cba3fa8c 100644 --- a/go/nexus/internal/router/anthropic_messages.go +++ b/go/nexus/internal/router/anthropic_messages.go @@ -18,6 +18,7 @@ func (in *AnthropicRouter) newMessagesRoute() RouteConfig { RequestConverter: in.messagesRequestConverter, ResponsesResponseConverter: in.messagesResponseConverter, ErrorConverter: in.anthropicErrorConverter, + PreCallback: disableProviderStreamingPreCallback(), SendDoneMarker: new(false), StreamConfig: &StreamConfig{ ResponsesStreamResponseConverter: in.messagesStreamResponseConverter, diff --git a/go/nexus/internal/router/bedrock_allowed_requests.go b/go/nexus/internal/router/bedrock_allowed_requests.go new file mode 100644 index 0000000000..72e2e6c3ef --- /dev/null +++ b/go/nexus/internal/router/bedrock_allowed_requests.go @@ -0,0 +1,23 @@ +package router + +import ( + "github.com/maximhq/bifrost/core/schemas" + pb "github.com/pluralsh/console/go/nexus/internal/proto" +) + +// bedrockAllowedRequests configures Bifrost operation gating when streaming is disabled. +// Returns nil when all operations should remain allowed (default / enableStream unset or true). +func bedrockAllowedRequests(cfg *pb.BedrockConfig) *schemas.AllowedRequests { + if cfg == nil || cfg.EnableStream == nil || cfg.GetEnableStream() { + return nil + } + + return &schemas.AllowedRequests{ + ListModels: true, + ChatCompletion: true, + ChatCompletionStream: false, + Responses: true, + ResponsesStream: false, + Embedding: true, + } +} diff --git a/go/nexus/internal/router/bedrock_allowed_requests_test.go b/go/nexus/internal/router/bedrock_allowed_requests_test.go new file mode 100644 index 0000000000..f2e79e1099 --- /dev/null +++ b/go/nexus/internal/router/bedrock_allowed_requests_test.go @@ -0,0 +1,27 @@ +package router + +import ( + "testing" + + pb "github.com/pluralsh/console/go/nexus/internal/proto" + "github.com/samber/lo" + "github.com/stretchr/testify/require" +) + +func TestBedrockAllowedRequests(t *testing.T) { + require.Nil(t, bedrockAllowedRequests(nil)) + require.Nil(t, bedrockAllowedRequests(&pb.BedrockConfig{})) + require.Nil(t, bedrockAllowedRequests(&pb.BedrockConfig{EnableStream: lo.ToPtr(true)})) + + allowed := bedrockAllowedRequests(&pb.BedrockConfig{EnableStream: lo.ToPtr(false)}) + require.NotNil(t, allowed) + require.True(t, allowed.ChatCompletion) + require.False(t, allowed.ChatCompletionStream) + require.True(t, allowed.Responses) + require.False(t, allowed.ResponsesStream) +} + +func TestBedrockNetworkBaseURL(t *testing.T) { + require.Equal(t, "https://bedrock-proxy.example.com", bedrockNetworkBaseURL("https://bedrock-proxy.example.com/")) + require.Equal(t, "https://bedrock-proxy.example.com", bedrockNetworkBaseURL(" https://bedrock-proxy.example.com ")) +} diff --git a/go/nexus/internal/router/oauth_token.go b/go/nexus/internal/router/oauth_token.go new file mode 100644 index 0000000000..c40caaf22f --- /dev/null +++ b/go/nexus/internal/router/oauth_token.go @@ -0,0 +1,25 @@ +package router + +import ( + "context" + "fmt" + + pb "github.com/pluralsh/console/go/nexus/internal/proto" +) + +// oauthAccessToken returns a bearer credential from OAuth2 client-credentials exchange when +// enabled, otherwise an empty string. providerName is used in validation error messages. +func (in *Account) oauthAccessToken(ctx context.Context, tx *pb.OpenAiTokenExchange, providerName string) (string, error) { + if tx == nil || !tx.GetEnabled() { + return "", nil + } + + tokenURL := tx.GetTokenUrl() + clientID := tx.GetClientId() + clientSecret := tx.GetClientSecret() + if tokenURL == "" || clientID == "" || clientSecret == "" { + return "", fmt.Errorf("%s tokenExchange enabled but tokenUrl, clientId, and clientSecret must all be set", providerName) + } + + return in.tokenCache.AccessToken(ctx, tokenURL, clientID, clientSecret) +} diff --git a/go/nexus/internal/router/openai_chat.go b/go/nexus/internal/router/openai_chat.go index f0759a2da3..3f34596686 100644 --- a/go/nexus/internal/router/openai_chat.go +++ b/go/nexus/internal/router/openai_chat.go @@ -16,7 +16,10 @@ func (in *OpenAIRouter) newChatCompletionsRoute() RouteConfig { RequestConverter: in.chatCompletionsRequestConverter, ChatResponseConverter: in.chatCompletionsResponseConverter, ErrorConverter: in.errorConverter, - PreCallback: in.openAIRoutePreCallback(string(RouteChatCompletions)), + PreCallback: chainPreCallbacks( + in.openAIRoutePreCallback(string(RouteChatCompletions)), + disableProviderStreamingPreCallback(), + ), StreamConfig: &StreamConfig{ ChatStreamResponseConverter: in.chatCompletionsStreamConverter, ErrorConverter: in.errorConverter, diff --git a/go/nexus/internal/router/openai_responses.go b/go/nexus/internal/router/openai_responses.go index b369f6d9c6..a67c544a13 100644 --- a/go/nexus/internal/router/openai_responses.go +++ b/go/nexus/internal/router/openai_responses.go @@ -16,7 +16,10 @@ func (in *OpenAIRouter) newResponsesRoute() RouteConfig { RequestConverter: in.responsesRequestConverter, ResponsesResponseConverter: in.responsesResponseConverter, ErrorConverter: in.errorConverter, - PreCallback: in.openAIRoutePreCallback(string(RouteResponses)), + PreCallback: chainPreCallbacks( + in.openAIRoutePreCallback(string(RouteResponses)), + disableProviderStreamingPreCallback(), + ), StreamConfig: &StreamConfig{ ResponsesStreamResponseConverter: in.responsesStreamResponseConverter, ErrorConverter: in.errorConverter, diff --git a/go/nexus/internal/router/provider_base_url.go b/go/nexus/internal/router/provider_base_url.go index 236e1a5283..d6733c9699 100644 --- a/go/nexus/internal/router/provider_base_url.go +++ b/go/nexus/internal/router/provider_base_url.go @@ -13,3 +13,8 @@ func bifrostNetworkBaseURL(baseURL string) string { } return baseURL } + +// bedrockNetworkBaseURL normalizes custom Bedrock-compatible proxy base URLs for Bifrost. +func bedrockNetworkBaseURL(baseURL string) string { + return strings.TrimRight(strings.TrimSpace(baseURL), "/") +} diff --git a/go/nexus/internal/router/streaming_control.go b/go/nexus/internal/router/streaming_control.go new file mode 100644 index 0000000000..3569550319 --- /dev/null +++ b/go/nexus/internal/router/streaming_control.go @@ -0,0 +1,53 @@ +package router + +import ( + "net/http" + "strings" + + "github.com/maximhq/bifrost/core/providers/anthropic" + "github.com/maximhq/bifrost/core/providers/openai" + "github.com/maximhq/bifrost/core/schemas" +) + +// PluralEnableStreamHeader disables model-provider streaming when set to "false". +// Agent runtimes send this via ANTHROPIC_CUSTOM_HEADERS (Claude) or model provider +// http_headers in Codex config.toml when spec.config.*.disableStream is true. +const PluralEnableStreamHeader = "X-Plural-Enable-Stream" + +func providerStreamingDisabled(r *http.Request) bool { + v := strings.TrimSpace(r.Header.Get(PluralEnableStreamHeader)) + return strings.EqualFold(v, "false") || v == "0" +} + +func chainPreCallbacks(callbacks ...PreRequestCallback) PreRequestCallback { + return func(r *http.Request, ctx *schemas.BifrostContext, req interface{}) error { + for _, cb := range callbacks { + if cb == nil { + continue + } + if err := cb(r, ctx, req); err != nil { + return err + } + } + return nil + } +} + +func disableProviderStreamingPreCallback() PreRequestCallback { + return func(r *http.Request, _ *schemas.BifrostContext, req interface{}) error { + if !providerStreamingDisabled(r) { + return nil + } + + streamOff := false + switch t := req.(type) { + case *anthropic.AnthropicMessageRequest: + t.Stream = &streamOff + case *openai.OpenAIChatRequest: + t.Stream = &streamOff + case *openai.OpenAIResponsesRequest: + t.Stream = &streamOff + } + return nil + } +} diff --git a/go/nexus/internal/router/streaming_control_test.go b/go/nexus/internal/router/streaming_control_test.go new file mode 100644 index 0000000000..7c62451f9e --- /dev/null +++ b/go/nexus/internal/router/streaming_control_test.go @@ -0,0 +1,45 @@ +package router + +import ( + "net/http" + "testing" + + "github.com/maximhq/bifrost/core/providers/anthropic" + "github.com/maximhq/bifrost/core/providers/openai" + "github.com/stretchr/testify/require" +) + +func TestDisableProviderStreamingPreCallback(t *testing.T) { + cb := disableProviderStreamingPreCallback() + + t.Run("no header leaves stream unset", func(t *testing.T) { + req := &anthropic.AnthropicMessageRequest{} + stream := true + req.Stream = &stream + err := cb(&http.Request{Header: http.Header{}}, nil, req) + require.NoError(t, err) + require.True(t, *req.Stream) + }) + + t.Run("false header clears anthropic stream", func(t *testing.T) { + req := &anthropic.AnthropicMessageRequest{} + stream := true + req.Stream = &stream + r, _ := http.NewRequest(http.MethodPost, "http://example.com", nil) + r.Header.Set(PluralEnableStreamHeader, "false") + err := cb(r, nil, req) + require.NoError(t, err) + require.False(t, *req.Stream) + }) + + t.Run("false header clears openai responses stream", func(t *testing.T) { + req := &openai.OpenAIResponsesRequest{} + stream := true + req.Stream = &stream + r, _ := http.NewRequest(http.MethodPost, "http://example.com", nil) + r.Header.Set(PluralEnableStreamHeader, "false") + err := cb(r, nil, req) + require.NoError(t, err) + require.False(t, *req.Stream) + }) +} diff --git a/go/observability-proxy/internal/proto/console.pb.go b/go/observability-proxy/internal/proto/console.pb.go index 2bcd548c74..79644c945f 100644 --- a/go/observability-proxy/internal/proto/console.pb.go +++ b/go/observability-proxy/internal/proto/console.pb.go @@ -21,6 +21,59 @@ const ( _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) ) +// OpenAiMethod selects which OpenAI HTTP APIs the Nexus proxy exposes. +type OpenAiMethod int32 + +const ( + OpenAiMethod_OPEN_AI_METHOD_UNSPECIFIED OpenAiMethod = 0 + OpenAiMethod_CHAT OpenAiMethod = 1 + OpenAiMethod_RESPONSES OpenAiMethod = 2 + OpenAiMethod_AUTO OpenAiMethod = 3 +) + +// Enum value maps for OpenAiMethod. +var ( + OpenAiMethod_name = map[int32]string{ + 0: "OPEN_AI_METHOD_UNSPECIFIED", + 1: "CHAT", + 2: "RESPONSES", + 3: "AUTO", + } + OpenAiMethod_value = map[string]int32{ + "OPEN_AI_METHOD_UNSPECIFIED": 0, + "CHAT": 1, + "RESPONSES": 2, + "AUTO": 3, + } +) + +func (x OpenAiMethod) Enum() *OpenAiMethod { + p := new(OpenAiMethod) + *p = x + return p +} + +func (x OpenAiMethod) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (OpenAiMethod) Descriptor() protoreflect.EnumDescriptor { + return file_console_proto_enumTypes[0].Descriptor() +} + +func (OpenAiMethod) Type() protoreflect.EnumType { + return &file_console_proto_enumTypes[0] +} + +func (x OpenAiMethod) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use OpenAiMethod.Descriptor instead. +func (OpenAiMethod) EnumDescriptor() ([]byte, []int) { + return file_console_proto_rawDescGZIP(), []int{0} +} + type AiConfigRequest struct { state protoimpl.MessageState `protogen:"open.v1"` unknownFields protoimpl.UnknownFields @@ -218,6 +271,7 @@ type OpenAiConfig struct { BaseUrl *string `protobuf:"bytes,5,opt,name=baseUrl,proto3,oneof" json:"baseUrl,omitempty"` ProxyModels []string `protobuf:"bytes,6,rep,name=proxyModels,proto3" json:"proxyModels,omitempty"` TokenExchange *OpenAiTokenExchange `protobuf:"bytes,7,opt,name=tokenExchange,proto3,oneof" json:"tokenExchange,omitempty"` + Method *OpenAiMethod `protobuf:"varint,8,opt,name=method,proto3,enum=plrl.OpenAiMethod,oneof" json:"method,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -301,6 +355,13 @@ func (x *OpenAiConfig) GetTokenExchange() *OpenAiTokenExchange { return nil } +func (x *OpenAiConfig) GetMethod() OpenAiMethod { + if x != nil && x.Method != nil { + return *x.Method + } + return OpenAiMethod_OPEN_AI_METHOD_UNSPECIFIED +} + type AnthropicConfig struct { state protoimpl.MessageState `protogen:"open.v1"` BaseUrl *string `protobuf:"bytes,1,opt,name=baseUrl,proto3,oneof" json:"baseUrl,omitempty"` @@ -504,6 +565,9 @@ type BedrockConfig struct { AwsSecretAccessKey *string `protobuf:"bytes,7,opt,name=awsSecretAccessKey,proto3,oneof" json:"awsSecretAccessKey,omitempty"` ProxyModels []string `protobuf:"bytes,8,rep,name=proxyModels,proto3" json:"proxyModels,omitempty"` Deployments map[string]string `protobuf:"bytes,9,rep,name=deployments,proto3" json:"deployments,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + BaseUrl *string `protobuf:"bytes,10,opt,name=baseUrl,proto3,oneof" json:"baseUrl,omitempty"` + TokenExchange *OpenAiTokenExchange `protobuf:"bytes,11,opt,name=tokenExchange,proto3,oneof" json:"tokenExchange,omitempty"` + EnableStream *bool `protobuf:"varint,12,opt,name=enableStream,proto3,oneof" json:"enableStream,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -601,6 +665,27 @@ func (x *BedrockConfig) GetDeployments() map[string]string { return nil } +func (x *BedrockConfig) GetBaseUrl() string { + if x != nil && x.BaseUrl != nil { + return *x.BaseUrl + } + return "" +} + +func (x *BedrockConfig) GetTokenExchange() *OpenAiTokenExchange { + if x != nil { + return x.TokenExchange + } + return nil +} + +func (x *BedrockConfig) GetEnableStream() bool { + if x != nil && x.EnableStream != nil { + return *x.EnableStream + } + return false +} + type AzureOpenAiConfig struct { state protoimpl.MessageState `protogen:"open.v1"` ApiVersion *string `protobuf:"bytes,1,opt,name=apiVersion,proto3,oneof" json:"apiVersion,omitempty"` @@ -1027,7 +1112,7 @@ const file_console_proto_rawDesc = "" + "\b_enabledB\v\n" + "\t_tokenUrlB\v\n" + "\t_clientIdB\x0f\n" + - "\r_clientSecret\"\xf1\x02\n" + + "\r_clientSecret\"\xad\x03\n" + "\fOpenAiConfig\x12\x1b\n" + "\x06apiKey\x18\x01 \x01(\tH\x00R\x06apiKey\x88\x01\x01\x12\x19\n" + "\x05model\x18\x02 \x01(\tH\x01R\x05model\x88\x01\x01\x12+\n" + @@ -1035,7 +1120,8 @@ const file_console_proto_rawDesc = "" + "\ttoolModel\x18\x04 \x01(\tH\x03R\ttoolModel\x88\x01\x01\x12\x1d\n" + "\abaseUrl\x18\x05 \x01(\tH\x04R\abaseUrl\x88\x01\x01\x12 \n" + "\vproxyModels\x18\x06 \x03(\tR\vproxyModels\x12D\n" + - "\rtokenExchange\x18\a \x01(\v2\x19.plrl.OpenAiTokenExchangeH\x05R\rtokenExchange\x88\x01\x01B\t\n" + + "\rtokenExchange\x18\a \x01(\v2\x19.plrl.OpenAiTokenExchangeH\x05R\rtokenExchange\x88\x01\x01\x12/\n" + + "\x06method\x18\b \x01(\x0e2\x12.plrl.OpenAiMethodH\x06R\x06method\x88\x01\x01B\t\n" + "\a_apiKeyB\b\n" + "\x06_modelB\x11\n" + "\x0f_embeddingModelB\f\n" + @@ -1043,7 +1129,8 @@ const file_console_proto_rawDesc = "" + "_toolModelB\n" + "\n" + "\b_baseUrlB\x10\n" + - "\x0e_tokenExchange\"\x9c\x02\n" + + "\x0e_tokenExchangeB\t\n" + + "\a_method\"\x9c\x02\n" + "\x0fAnthropicConfig\x12\x1d\n" + "\abaseUrl\x18\x01 \x01(\tH\x00R\abaseUrl\x88\x01\x01\x12\x1b\n" + "\x06apiKey\x18\x02 \x01(\tH\x01R\x06apiKey\x88\x01\x01\x12\x19\n" + @@ -1077,7 +1164,7 @@ const file_console_proto_rawDesc = "" + "_toolModelB\n" + "\n" + "\b_projectB\v\n" + - "\t_location\"\xcc\x04\n" + + "\t_location\"\x89\x06\n" + "\rBedrockConfig\x12\x1d\n" + "\amodelId\x18\x01 \x01(\tH\x00R\amodelId\x88\x01\x01\x12%\n" + "\vtoolModelId\x18\x02 \x01(\tH\x01R\vtoolModelId\x88\x01\x01\x12%\n" + @@ -1087,7 +1174,11 @@ const file_console_proto_rawDesc = "" + "\x0eawsAccessKeyId\x18\x06 \x01(\tH\x05R\x0eawsAccessKeyId\x88\x01\x01\x123\n" + "\x12awsSecretAccessKey\x18\a \x01(\tH\x06R\x12awsSecretAccessKey\x88\x01\x01\x12 \n" + "\vproxyModels\x18\b \x03(\tR\vproxyModels\x12F\n" + - "\vdeployments\x18\t \x03(\v2$.plrl.BedrockConfig.DeploymentsEntryR\vdeployments\x1a>\n" + + "\vdeployments\x18\t \x03(\v2$.plrl.BedrockConfig.DeploymentsEntryR\vdeployments\x12\x1d\n" + + "\abaseUrl\x18\n" + + " \x01(\tH\aR\abaseUrl\x88\x01\x01\x12D\n" + + "\rtokenExchange\x18\v \x01(\v2\x19.plrl.OpenAiTokenExchangeH\bR\rtokenExchange\x88\x01\x01\x12'\n" + + "\fenableStream\x18\f \x01(\bH\tR\fenableStream\x88\x01\x01\x1a>\n" + "\x10DeploymentsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01B\n" + @@ -1098,7 +1189,11 @@ const file_console_proto_rawDesc = "" + "\a_regionB\x13\n" + "\x11_embeddingModelIdB\x11\n" + "\x0f_awsAccessKeyIdB\x15\n" + - "\x13_awsSecretAccessKey\"\xf0\x03\n" + + "\x13_awsSecretAccessKeyB\n" + + "\n" + + "\b_baseUrlB\x10\n" + + "\x0e_tokenExchangeB\x0f\n" + + "\r_enableStream\"\xf0\x03\n" + "\x11AzureOpenAiConfig\x12#\n" + "\n" + "apiVersion\x18\x01 \x01(\tH\x00R\n" + @@ -1143,7 +1238,12 @@ const file_console_proto_rawDesc = "" + "\x05bytes\x18\x01 \x01(\x03R\x05bytes\"0\n" + "\x14MeterMetricsResponse\x12\x18\n" + "\asuccess\x18\x01 \x01(\bR\asuccess\"\x1c\n" + - "\x1aObservabilityConfigRequest2\xbe\x02\n" + + "\x1aObservabilityConfigRequest*Q\n" + + "\fOpenAiMethod\x12\x1e\n" + + "\x1aOPEN_AI_METHOD_UNSPECIFIED\x10\x00\x12\b\n" + + "\x04CHAT\x10\x01\x12\r\n" + + "\tRESPONSES\x10\x02\x12\b\n" + + "\x04AUTO\x10\x032\xbe\x02\n" + "\fPluralServer\x12E\n" + "\fMeterMetrics\x12\x19.plrl.MeterMetricsRequest\x1a\x1a.plrl.MeterMetricsResponse\x124\n" + "\vGetAiConfig\x12\x15.plrl.AiConfigRequest\x1a\x0e.plrl.AiConfig\x12U\n" + @@ -1162,47 +1262,51 @@ func file_console_proto_rawDescGZIP() []byte { return file_console_proto_rawDescData } +var file_console_proto_enumTypes = make([]protoimpl.EnumInfo, 1) var file_console_proto_msgTypes = make([]protoimpl.MessageInfo, 16) var file_console_proto_goTypes = []any{ - (*AiConfigRequest)(nil), // 0: plrl.AiConfigRequest - (*AiConfig)(nil), // 1: plrl.AiConfig - (*OpenAiTokenExchange)(nil), // 2: plrl.OpenAiTokenExchange - (*OpenAiConfig)(nil), // 3: plrl.OpenAiConfig - (*AnthropicConfig)(nil), // 4: plrl.AnthropicConfig - (*VertexAiConfig)(nil), // 5: plrl.VertexAiConfig - (*BedrockConfig)(nil), // 6: plrl.BedrockConfig - (*AzureOpenAiConfig)(nil), // 7: plrl.AzureOpenAiConfig - (*ProxyAuthenticationRequest)(nil), // 8: plrl.ProxyAuthenticationRequest - (*ProxyAuthenticationResponse)(nil), // 9: plrl.ProxyAuthenticationResponse - (*ObservabilityConfig)(nil), // 10: plrl.ObservabilityConfig - (*MeterMetricsRequest)(nil), // 11: plrl.MeterMetricsRequest - (*MeterMetricsResponse)(nil), // 12: plrl.MeterMetricsResponse - (*ObservabilityConfigRequest)(nil), // 13: plrl.ObservabilityConfigRequest - nil, // 14: plrl.BedrockConfig.DeploymentsEntry - nil, // 15: plrl.AzureOpenAiConfig.DeploymentsEntry + (OpenAiMethod)(0), // 0: plrl.OpenAiMethod + (*AiConfigRequest)(nil), // 1: plrl.AiConfigRequest + (*AiConfig)(nil), // 2: plrl.AiConfig + (*OpenAiTokenExchange)(nil), // 3: plrl.OpenAiTokenExchange + (*OpenAiConfig)(nil), // 4: plrl.OpenAiConfig + (*AnthropicConfig)(nil), // 5: plrl.AnthropicConfig + (*VertexAiConfig)(nil), // 6: plrl.VertexAiConfig + (*BedrockConfig)(nil), // 7: plrl.BedrockConfig + (*AzureOpenAiConfig)(nil), // 8: plrl.AzureOpenAiConfig + (*ProxyAuthenticationRequest)(nil), // 9: plrl.ProxyAuthenticationRequest + (*ProxyAuthenticationResponse)(nil), // 10: plrl.ProxyAuthenticationResponse + (*ObservabilityConfig)(nil), // 11: plrl.ObservabilityConfig + (*MeterMetricsRequest)(nil), // 12: plrl.MeterMetricsRequest + (*MeterMetricsResponse)(nil), // 13: plrl.MeterMetricsResponse + (*ObservabilityConfigRequest)(nil), // 14: plrl.ObservabilityConfigRequest + nil, // 15: plrl.BedrockConfig.DeploymentsEntry + nil, // 16: plrl.AzureOpenAiConfig.DeploymentsEntry } var file_console_proto_depIdxs = []int32{ - 3, // 0: plrl.AiConfig.openai:type_name -> plrl.OpenAiConfig - 4, // 1: plrl.AiConfig.anthropic:type_name -> plrl.AnthropicConfig - 5, // 2: plrl.AiConfig.vertexAi:type_name -> plrl.VertexAiConfig - 6, // 3: plrl.AiConfig.bedrock:type_name -> plrl.BedrockConfig - 7, // 4: plrl.AiConfig.azure:type_name -> plrl.AzureOpenAiConfig - 2, // 5: plrl.OpenAiConfig.tokenExchange:type_name -> plrl.OpenAiTokenExchange - 14, // 6: plrl.BedrockConfig.deployments:type_name -> plrl.BedrockConfig.DeploymentsEntry - 15, // 7: plrl.AzureOpenAiConfig.deployments:type_name -> plrl.AzureOpenAiConfig.DeploymentsEntry - 11, // 8: plrl.PluralServer.MeterMetrics:input_type -> plrl.MeterMetricsRequest - 0, // 9: plrl.PluralServer.GetAiConfig:input_type -> plrl.AiConfigRequest - 13, // 10: plrl.PluralServer.GetObservabilityConfig:input_type -> plrl.ObservabilityConfigRequest - 8, // 11: plrl.PluralServer.ProxyAuthentication:input_type -> plrl.ProxyAuthenticationRequest - 12, // 12: plrl.PluralServer.MeterMetrics:output_type -> plrl.MeterMetricsResponse - 1, // 13: plrl.PluralServer.GetAiConfig:output_type -> plrl.AiConfig - 10, // 14: plrl.PluralServer.GetObservabilityConfig:output_type -> plrl.ObservabilityConfig - 9, // 15: plrl.PluralServer.ProxyAuthentication:output_type -> plrl.ProxyAuthenticationResponse - 12, // [12:16] is the sub-list for method output_type - 8, // [8:12] is the sub-list for method input_type - 8, // [8:8] is the sub-list for extension type_name - 8, // [8:8] is the sub-list for extension extendee - 0, // [0:8] is the sub-list for field type_name + 4, // 0: plrl.AiConfig.openai:type_name -> plrl.OpenAiConfig + 5, // 1: plrl.AiConfig.anthropic:type_name -> plrl.AnthropicConfig + 6, // 2: plrl.AiConfig.vertexAi:type_name -> plrl.VertexAiConfig + 7, // 3: plrl.AiConfig.bedrock:type_name -> plrl.BedrockConfig + 8, // 4: plrl.AiConfig.azure:type_name -> plrl.AzureOpenAiConfig + 3, // 5: plrl.OpenAiConfig.tokenExchange:type_name -> plrl.OpenAiTokenExchange + 0, // 6: plrl.OpenAiConfig.method:type_name -> plrl.OpenAiMethod + 15, // 7: plrl.BedrockConfig.deployments:type_name -> plrl.BedrockConfig.DeploymentsEntry + 3, // 8: plrl.BedrockConfig.tokenExchange:type_name -> plrl.OpenAiTokenExchange + 16, // 9: plrl.AzureOpenAiConfig.deployments:type_name -> plrl.AzureOpenAiConfig.DeploymentsEntry + 12, // 10: plrl.PluralServer.MeterMetrics:input_type -> plrl.MeterMetricsRequest + 1, // 11: plrl.PluralServer.GetAiConfig:input_type -> plrl.AiConfigRequest + 14, // 12: plrl.PluralServer.GetObservabilityConfig:input_type -> plrl.ObservabilityConfigRequest + 9, // 13: plrl.PluralServer.ProxyAuthentication:input_type -> plrl.ProxyAuthenticationRequest + 13, // 14: plrl.PluralServer.MeterMetrics:output_type -> plrl.MeterMetricsResponse + 2, // 15: plrl.PluralServer.GetAiConfig:output_type -> plrl.AiConfig + 11, // 16: plrl.PluralServer.GetObservabilityConfig:output_type -> plrl.ObservabilityConfig + 10, // 17: plrl.PluralServer.ProxyAuthentication:output_type -> plrl.ProxyAuthenticationResponse + 14, // [14:18] is the sub-list for method output_type + 10, // [10:14] is the sub-list for method input_type + 10, // [10:10] is the sub-list for extension type_name + 10, // [10:10] is the sub-list for extension extendee + 0, // [0:10] is the sub-list for field type_name } func init() { file_console_proto_init() } @@ -1222,13 +1326,14 @@ func file_console_proto_init() { File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_console_proto_rawDesc), len(file_console_proto_rawDesc)), - NumEnums: 0, + NumEnums: 1, NumMessages: 16, NumExtensions: 0, NumServices: 1, }, GoTypes: file_console_proto_goTypes, DependencyIndexes: file_console_proto_depIdxs, + EnumInfos: file_console_proto_enumTypes, MessageInfos: file_console_proto_msgTypes, }.Build() File_console_proto = out.File diff --git a/lib/console/ai/provider/bedrock.ex b/lib/console/ai/provider/bedrock.ex index d0fd687ad9..389e05d79a 100644 --- a/lib/console/ai/provider/bedrock.ex +++ b/lib/console/ai/provider/bedrock.ex @@ -4,11 +4,12 @@ defmodule Console.AI.Bedrock do """ @behaviour Console.AI.Provider import Console.AI.Provider.Base - alias Console.AI.{Utils, Stream} + alias Console.AI.{Utils, Stream, Provider.TokenExchange} + alias Console.Schema.DeploymentSettings.OauthToken require Logger - defstruct [:access_token, :model_id, :tool_model_id, :region, :embedding_model, :aws_access_key_id, :aws_secret_access_key, :stream] + defstruct [:access_token, :model_id, :tool_model_id, :base_url, :region, :embedding_model, :aws_access_key_id, :aws_secret_access_key, :stream, :token_exchange] @type t :: %__MODULE__{} @@ -23,8 +24,10 @@ defmodule Console.AI.Bedrock do aws_access_key_id: opts.aws_access_key_id, aws_secret_access_key: opts.aws_secret_access_key, access_token: opts.access_token, + base_url: opts.base_url, region: opts.region, - stream: Stream.stream(), + token_exchange: opts.token_exchange, + stream: add_stream(opts), } end @@ -35,10 +38,14 @@ defmodule Console.AI.Bedrock do """ @spec completion(t(), Console.AI.Provider.history, keyword) :: {:ok, binary} | Console.error def completion(%__MODULE__{} = bedrock, messages, opts) do - messages - |> reqllm_messages() - |> generate_text("amazon-bedrock:#{select_model(bedrock, opts[:client])}", bedrock.stream, Keyword.put(provider_options(bedrock), :tools, tools(opts))) - |> reqllm_result() + with {:ok, provider_opts} <- provider_options(bedrock) do + provider_opts = Keyword.put(provider_opts, :tools, tools(opts)) + + messages + |> reqllm_messages() + |> generate_text("amazon-bedrock:#{select_model(bedrock, opts[:client])}", bedrock.stream, provider_opts) + |> reqllm_result() + end end @doc """ @@ -46,24 +53,29 @@ defmodule Console.AI.Bedrock do """ @spec tool_call(t(), Console.AI.Provider.history, [atom], keyword) :: {:ok, binary} | {:ok, [Console.AI.Tool.t]} | Console.error def tool_call(%__MODULE__{} = bedrock, messages, tools, opts) do - provider_opts = Keyword.put(provider_options(bedrock), :tools, reqllm_tools(tools)) + with {:ok, provider_opts} <- provider_options(bedrock) do + provider_opts = Keyword.put(provider_opts, :tools, reqllm_tools(tools)) - messages - |> reqllm_messages() - |> generate_text("amazon-bedrock:#{select_model(bedrock, opts[:client] || :tool)}", bedrock.stream, provider_opts) - |> reqllm_result() - |> tool_calls() + messages + |> reqllm_messages() + |> generate_text("amazon-bedrock:#{select_model(bedrock, opts[:client] || :tool)}", bedrock.stream, provider_opts) + |> reqllm_result() + |> tool_calls() + end end def embeddings(%__MODULE__{} = bedrock, text) do bedrock = choose_region(bedrock) chunked = Utils.chunk(text, 8000) - provider_options(bedrock) - |> Keyword.put(:dimensions, Utils.embedding_dims()) - |> then(&ReqLLM.embed("amazon-bedrock:#{bedrock.embedding_model}", chunked, &1)) - |> case do - {:ok, embeddings} -> {:ok, Enum.zip(chunked, embeddings)} - error -> error + + with {:ok, provider} <- provider_options(bedrock) do + provider + |> Keyword.put(:dimensions, Utils.embedding_dims()) + |> then(&ReqLLM.embed("amazon-bedrock:#{bedrock.embedding_model}", chunked, &1)) + |> case do + {:ok, embeddings} -> {:ok, Enum.zip(chunked, embeddings)} + error -> error + end end end @@ -76,11 +88,24 @@ defmodule Console.AI.Bedrock do def tools?(), do: true - def provider_options(%__MODULE__{region: region, access_token: token} = bedrock) do - [region: region, access_token: token] - |> Enum.concat(if is_nil(token), do: aws_auth(bedrock), else: []) - |> Enum.filter(fn {_, v} -> not is_nil(v) end) + def provider_options(%__MODULE__{region: region, base_url: base_url} = bedrock) do + with {:ok, token} <- api_key(bedrock) do + {:ok, + [region: region, api_key: token, base_url: base_url] + |> Enum.concat(if is_nil(token), do: aws_auth(bedrock), else: []) + |> Enum.filter(fn {_, v} -> not is_nil(v) end)} + end + end + + defp api_key(%__MODULE__{token_exchange: %OauthToken{enabled: true} = token}) do + with {:ok, %OAuth2.AccessToken{access_token: token}} <- TokenExchange.exchange(token.token_url, token.client_id, token.client_secret), + do: {:ok, token} end + defp api_key(%__MODULE__{access_token: token}) when is_binary(token), do: {:ok, token} + defp api_key(_), do: {:ok, nil} + + defp add_stream(%{enable_stream: false}), do: nil + defp add_stream(_), do: Stream.stream() defp choose_region(%__MODULE__{embedding_model: "cohere.embed-english-v3"} = rock), do: %{rock | region: "us-east-1"} diff --git a/lib/console/graphql/deployments/settings.ex b/lib/console/graphql/deployments/settings.ex index 2314c2b5d2..4ddd17eb29 100644 --- a/lib/console/graphql/deployments/settings.ex +++ b/lib/console/graphql/deployments/settings.ex @@ -174,6 +174,7 @@ defmodule Console.GraphQl.Deployments.Settings do input_object :bedrock_ai_attributes do field :model_id, :string, description: @bedrock_model_id_doc field :tool_model_id, :string, description: "Bedrock model or inference profile for tool calls. Same ID formats as modelId." + field :base_url, :string, description: "the base url to use when querying a Bedrock-compatible API" field :access_token, :string, description: "the openai bedrock access token to use" field :region, :string, description: "the aws region the model is hosted in" field :aws_access_key_id, :string, description: "the aws access key id to use (DEPRECATED)" @@ -181,6 +182,8 @@ defmodule Console.GraphQl.Deployments.Settings do field :embedding_model, :string, description: "Bedrock model or inference profile for embeddings. Same ID formats as modelId." field :proxy_models, list_of(:string), description: @bedrock_proxy_models_doc field :deployments, :json, description: @bedrock_deployments_doc + field :enable_stream, :boolean, description: "whether to enable streaming responses" + field :token_exchange, :openai_token_exchange_attributes, description: "OAuth2 client credentials against a token endpoint to obtain access tokens" end input_object :vertex_ai_attributes do @@ -416,12 +419,15 @@ defmodule Console.GraphQl.Deployments.Settings do object :bedrock_ai_settings do field :model_id, :string, description: @bedrock_model_id_doc <> " Omit for Plural defaults." field :tool_model_id, :string, description: "Bedrock model or inference profile for tool calls. Same ID formats as modelId." + field :base_url, :string, description: "the base url to use when querying a Bedrock-compatible API" field :access_key_id, :string, description: "the openai bedrock aws access key id to use (DEPRECATED)", resolve: fn b, _, _ -> {:ok, Map.get(b, :aws_access_key_id)} end field :region, :string, description: "the aws region the model is hosted in" field :embedding_model, :string, description: "Bedrock model or inference profile for embeddings. Same ID formats as modelId." field :proxy_models, list_of(:string), description: @bedrock_proxy_models_doc field :deployments, :map, description: @bedrock_deployments_doc + field :enable_stream, :boolean, description: "whether streaming responses are enabled" + field :token_exchange, :openai_token_exchange, description: "OAuth2 client credentials configured for token endpoint exchange" end @desc "Settings for usage of GCP VertexAI for LLMs" diff --git a/lib/console/grpc/server.ex b/lib/console/grpc/server.ex index 3e6521948c..572e8afa86 100644 --- a/lib/console/grpc/server.ex +++ b/lib/console/grpc/server.ex @@ -126,10 +126,14 @@ defmodule Console.GRPC.Server do toolModelId: bedrock.tool_model_id || defaults[:tool_model], embeddingModelId: bedrock.embedding_model || defaults[:embedding_model], region: bedrock.region, + accessToken: bedrock.access_token, + baseUrl: bedrock.base_url, awsAccessKeyId: bedrock.aws_access_key_id, awsSecretAccessKey: bedrock.aws_secret_access_key, proxyModels: proxy_models(bedrock), - deployments: to_string_map(bedrock.deployments) + deployments: to_string_map(bedrock.deployments), + tokenExchange: to_pb(bedrock.token_exchange), + enableStream: bedrock_enable_stream(bedrock.enable_stream) } end @@ -148,6 +152,9 @@ defmodule Console.GRPC.Server do defp to_pb(_), do: %Plrl.AiConfig{enabled: false} + defp bedrock_enable_stream(false), do: false + defp bedrock_enable_stream(_), do: true + defp openai_method_to_pb(:chat), do: :CHAT defp openai_method_to_pb(:responses), do: :RESPONSES defp openai_method_to_pb(:auto), do: :AUTO diff --git a/lib/console/schema/deployment_settings.ex b/lib/console/schema/deployment_settings.ex index 85199722fc..5be7f43793 100644 --- a/lib/console/schema/deployment_settings.ex +++ b/lib/console/schema/deployment_settings.ex @@ -277,6 +277,7 @@ defmodule Console.Schema.DeploymentSettings do # Foundation model ID (e.g. anthropic.claude-3-5-sonnet-20241022-v2:0) or regional inference profile ID (e.g. us.anthropic.claude-3-5-sonnet-20241022-v2:0). field :model_id, :string field :tool_model_id, :string + field :base_url, :string field :access_token, EncryptedString field :region, :string field :embedding_model, :string @@ -286,6 +287,9 @@ defmodule Console.Schema.DeploymentSettings do field :proxy_models, {:array, :string} # Deprecated for most configs; maps client model ID -> inference profile ID when aliases cannot be inferred (e.g. application profile suffixes). field :deployments, :map + field :enable_stream, :boolean, default: true + + embeds_one :token_exchange, Console.Schema.DeploymentSettings.OauthToken, on_replace: :update end embeds_one :vertex, Vertex, on_replace: :update do @@ -428,7 +432,8 @@ defmodule Console.Schema.DeploymentSettings do defp bedrock_changeset(model, attrs) do model - |> cast(attrs, ~w(model_id tool_model_id access_token region embedding_model aws_access_key_id aws_secret_access_key proxy_models deployments)a) + |> cast(attrs, ~w(model_id tool_model_id base_url access_token region embedding_model aws_access_key_id aws_secret_access_key proxy_models deployments enable_stream)a) + |> cast_embed(:token_exchange) |> validate_required(~w(region)a) end diff --git a/lib/grpc/console.pb.ex b/lib/grpc/console.pb.ex index a586798651..5a21bf6259 100644 --- a/lib/grpc/console.pb.ex +++ b/lib/grpc/console.pb.ex @@ -132,6 +132,9 @@ defmodule Plrl.BedrockConfig do field :awsSecretAccessKey, 7, proto3_optional: true, type: :string field :proxyModels, 8, repeated: true, type: :string field :deployments, 9, repeated: true, type: Plrl.BedrockConfig.DeploymentsEntry, map: true + field :baseUrl, 10, proto3_optional: true, type: :string + field :tokenExchange, 11, proto3_optional: true, type: Plrl.OpenAiTokenExchange + field :enableStream, 12, proto3_optional: true, type: :bool end defmodule Plrl.AzureOpenAiConfig.DeploymentsEntry do diff --git a/proto/console.proto b/proto/console.proto index 2115a9837d..5af68e967f 100644 --- a/proto/console.proto +++ b/proto/console.proto @@ -74,6 +74,9 @@ message BedrockConfig { optional string awsSecretAccessKey = 7; repeated string proxyModels = 8; map deployments = 9; + optional string baseUrl = 10; + optional OpenAiTokenExchange tokenExchange = 11; + optional bool enableStream = 12; } message AzureOpenAiConfig { diff --git a/schema/schema.graphql b/schema/schema.graphql index dff4034ff6..e9bf81a09e 100644 --- a/schema/schema.graphql +++ b/schema/schema.graphql @@ -5463,6 +5463,9 @@ input BedrockAiAttributes { "Bedrock model or inference profile for tool calls. Same ID formats as modelId." toolModelId: String + "the base url to use when querying a Bedrock-compatible API" + baseUrl: String + "the openai bedrock access token to use" accessToken: String @@ -5483,6 +5486,12 @@ input BedrockAiAttributes { "Deprecated for most configurations: prefer regional-prefixed inference profile IDs in modelId or proxyModels (aliases are inferred automatically). Still needed for explicit client model name overrides, application inference profile resource IDs (profile suffix only, not full ARN), or when alias mapping cannot be inferred. Maps client-facing model ID to inference profile ID. Example: {\"anthropic.claude-3-5-sonnet-20241022-v2:0\": \"us.anthropic.claude-3-5-sonnet-20241022-v2:0\"}" deployments: Json + + "whether to enable streaming responses" + enableStream: Boolean + + "OAuth2 client credentials against a token endpoint to obtain access tokens" + tokenExchange: OpenaiTokenExchangeAttributes } input VertexAiAttributes { @@ -5850,6 +5859,9 @@ type BedrockAiSettings { "Bedrock model or inference profile for tool calls. Same ID formats as modelId." toolModelId: String + "the base url to use when querying a Bedrock-compatible API" + baseUrl: String + "the openai bedrock aws access key id to use (DEPRECATED)" accessKeyId: String @@ -5864,6 +5876,12 @@ type BedrockAiSettings { "Deprecated for most configurations: prefer regional-prefixed inference profile IDs in modelId or proxyModels (aliases are inferred automatically). Still needed for explicit client model name overrides, application inference profile resource IDs (profile suffix only, not full ARN), or when alias mapping cannot be inferred. Maps client-facing model ID to inference profile ID. Example: {\"anthropic.claude-3-5-sonnet-20241022-v2:0\": \"us.anthropic.claude-3-5-sonnet-20241022-v2:0\"}" deployments: Map + + "whether streaming responses are enabled" + enableStream: Boolean + + "OAuth2 client credentials configured for token endpoint exchange" + tokenExchange: OpenaiTokenExchange } "Settings for usage of GCP VertexAI for LLMs" diff --git a/test/console/ai/provider/bedrock_test.exs b/test/console/ai/provider/bedrock_test.exs new file mode 100644 index 0000000000..16ae4900bf --- /dev/null +++ b/test/console/ai/provider/bedrock_test.exs @@ -0,0 +1,142 @@ +defmodule Console.AI.Provider.BedrockTest do + @moduledoc false + use Console.DataCase, async: false + use Mimic + + alias Console.AI.Bedrock + alias Console.Schema.DeploymentSettings.OauthToken + + @token_url "https://example.com/oauth2/token" + @client_id "test-client-id" + @client_secret "test-client-secret" + @custom_base_url "https://custom-bedrock.example.com" + @region "us-east-1" + @model "global.anthropic.claude-haiku-4-5-20251001-v1:0" + + setup do + prev = Application.get_env(:oauth2, :adapter) + Application.put_env(:oauth2, :adapter, Tesla.Mock) + + on_exit(fn -> + if prev do + Application.put_env(:oauth2, :adapter, prev) + else + Application.delete_env(:oauth2, :adapter) + end + end) + + :ok + end + + describe "token exchange" do + setup :set_mimic_global + + test "completion exchanges OAuth token, passes base_url in provider opts, and sends Bearer auth via Req" do + oauth_access_token = "oauth-access-token-for-bedrock" + expected_invoke_url = + "https://bedrock-runtime.#{@region}.amazonaws.com/model/#{@model}/invoke" + + stub(Console.Cache, :get, fn _ -> nil end) + + Tesla.Mock.mock(fn env -> + assert env.method == :post + assert env.url == @token_url + assert URI.decode_query(env.body) == %{"grant_type" => "client_credentials"} + + assert header(env.headers, "authorization") == + "Basic " <> Base.encode64(@client_id <> ":" <> @client_secret) + + %Tesla.Env{ + status: 200, + headers: [{"content-type", "application/json"}], + body: + Jason.encode!(%{ + "access_token" => oauth_access_token, + "token_type" => "Bearer", + "expires_in" => 3600 + }) + } + end) + + expect(Console.Cache, :put, fn key, %OAuth2.AccessToken{access_token: ^oauth_access_token}, _opts -> + assert key == {:token_exchange, @token_url, @client_id, @client_secret} + :ok + end) + + bedrock = + Bedrock.new(%{ + access_token: nil, + aws_access_key_id: nil, + aws_secret_access_key: nil, + embedding_model: nil, + region: @region, + base_url: @custom_base_url, + model_id: @model, + tool_model_id: nil, + enable_stream: false, + token_exchange: %OauthToken{ + enabled: true, + token_url: @token_url, + client_id: @client_id, + client_secret: @client_secret + } + }) + + assert {:ok, provider_opts} = Bedrock.provider_options(bedrock) + assert Keyword.get(provider_opts, :api_key) == oauth_access_token + assert Keyword.get(provider_opts, :base_url) == @custom_base_url + + expect(Req, :request, fn %Req.Request{} = request -> + url = request.url |> URI.to_string() + + assert url == expected_invoke_url + + assert request_header(request, "authorization") in [ + "Bearer #{oauth_access_token}", + ["Bearer #{oauth_access_token}"] + ] + + raw_body = %{ + "id" => "msg_bedrock_e2e", + "type" => "message", + "role" => "assistant", + "model" => @model, + "content" => [%{"type" => "text", "text" => "ok"}], + "stop_reason" => "end_turn", + "usage" => %{"input_tokens" => 10, "output_tokens" => 5} + } + + {:ok, decoded} = + ReqLLM.Providers.AmazonBedrock.Anthropic.parse_response(raw_body, %{ + model: @model, + context: request.options[:context] + }) + + {:ok, + %Req.Response{ + status: 200, + headers: %{"content-type" => ["application/json"]}, + body: decoded + }} + end) + + assert {:ok, "ok"} = Bedrock.completion(bedrock, [{:user, "ping"}], []) + end + end + + defp header(headers, wanted) do + wanted = String.downcase(wanted) + + Enum.find_value(headers, fn {k, v} -> + if String.downcase(to_string(k)) == wanted, do: v + end) + end + + defp request_header(%Req.Request{headers: headers}, wanted) do + wanted = String.downcase(wanted) + + Enum.find_value(headers, fn {k, v} -> + if String.downcase(to_string(k)) == wanted, do: v + end) + end +end