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

Author: michaeljguarino

assets/src/generated/graphql.ts

@@ -1504,16 +1504,22 @@ export type BedrockAiAttributes = {
   awsAccessKeyId?: InputMaybe<Scalars['String']['input']>;
   /** the aws secret access key to use (DEPRECATED) */
   awsSecretAccessKey?: InputMaybe<Scalars['String']['input']>;
+  /** the base url to use when querying a Bedrock-compatible API */
+  baseUrl?: InputMaybe<Scalars['String']['input']>;
   /** 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<Scalars['Json']['input']>;
   /** Bedrock model or inference profile for embeddings. Same ID formats as modelId. */
   embeddingModel?: InputMaybe<Scalars['String']['input']>;
+  /** whether to enable streaming responses */
+  enableStream?: InputMaybe<Scalars['Boolean']['input']>;
   /** 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<Scalars['String']['input']>;
   /** 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<Array<InputMaybe<Scalars['String']['input']>>>;
   /** the aws region the model is hosted in */
   region?: InputMaybe<Scalars['String']['input']>;
+  /** OAuth2 client credentials against a token endpoint to obtain access tokens */
+  tokenExchange?: InputMaybe<OpenaiTokenExchangeAttributes>;
   /** Bedrock model or inference profile for tool calls. Same ID formats as modelId. */
   toolModelId?: InputMaybe<Scalars['String']['input']>;
 };
@@ -1523,16 +1529,22 @@ export type BedrockAiSettings = {
   __typename?: 'BedrockAiSettings';
   /** the openai bedrock aws access key id to use (DEPRECATED) */
   accessKeyId?: Maybe<Scalars['String']['output']>;
+  /** the base url to use when querying a Bedrock-compatible API */
+  baseUrl?: Maybe<Scalars['String']['output']>;
   /** 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<Scalars['Map']['output']>;
   /** Bedrock model or inference profile for embeddings. Same ID formats as modelId. */
   embeddingModel?: Maybe<Scalars['String']['output']>;
+  /** whether streaming responses are enabled */
+  enableStream?: Maybe<Scalars['Boolean']['output']>;
   /** 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<Scalars['String']['output']>;
   /** 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<Array<Maybe<Scalars['String']['output']>>>;
   /** the aws region the model is hosted in */
   region?: Maybe<Scalars['String']['output']>;
+  /** OAuth2 client credentials configured for token endpoint exchange */
+  tokenExchange?: Maybe<OpenaiTokenExchange>;
   /** Bedrock model or inference profile for tool calls. Same ID formats as modelId. */
   toolModelId?: Maybe<Scalars['String']['output']>;
 };

go/client/models_gen.go

@@ -749,13 +749,22 @@ 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,
 			ProxyModels:        lo.ToSlicePtr(in.Bedrock.ProxyModels),
 			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

go/controller/api/v1alpha1/deploymentsettings_types.go

@@ -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

go/controller/api/v1alpha1/zz_generated.deepcopy.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) {

go/controller/config/crd/bases/deployments.plural.sh_deploymentsettings.yaml

@@ -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