diff --git a/CHANGELOG.md b/CHANGELOG.md index 53fd9f4..63acf96 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,14 @@ separately by `model.SchemaVersion` (currently 1.2.0). alone is enough to pick a database. pgx's `ParseConfig` already reads `PGSERVICEFILE` (or the libpq default path); this just stops pgbot from erroring out before pgx gets a chance to. +- **AWS Bedrock Mantle for `ask` and `explain`** (#35). Select + `PGBOT_AI_PROVIDER=bedrock` (alias `mantle`) to route OpenAI models through + Responses and Anthropic models through Messages. Authenticate with + `PGBOT_AI_API_KEY` / `AWS_BEARER_TOKEN_BEDROCK`, or use the AWS SDK credential + chain for profiles, SSO, and role credentials. IAM tokens are signed locally, + restricted to the configured regional Mantle HTTPS host, and never forwarded + through redirects. Adds the AWS SDK for credential resolution and signing; + inference continues to use the existing HTTP clients. ## [0.8.1] - 2026-09-06 diff --git a/README.md b/README.md index 3ed5dbc..ae3d322 100644 --- a/README.md +++ b/README.md @@ -438,7 +438,7 @@ one SSH connection serves the whole run. Raise `--timeout` if the link is slow. | `GEMINI_API_KEY` / `GOOGLE_API_KEY` | Enables `ask` / `explain` via Google Gemini. | | `ANTHROPIC_API_KEY` | Enables `ask` / `explain` via Anthropic. | | `XAI_API_KEY` / `GROK_API_KEY` | Enables `ask` / `explain` via xAI. | -| `PGBOT_AI_PROVIDER` | `gemini`, `anthropic`, `openai`, or `xai` — picks one when several keys are set (auto-detection tries OpenAI first). | +| `PGBOT_AI_PROVIDER` | `gemini`, `anthropic`, `openai`, `xai`, or `bedrock` (alias `mantle`) — picks one when several keys are set (auto-detection tries OpenAI first). | | `PGBOT_AI_MODEL` / `PGBOT_AI_BASE_URL` / `PGBOT_AI_API_KEY` | Model, endpoint, and key override for whichever provider is selected; the way to reach an OpenAI-compatible service (OpenRouter, Groq, Ollama, vLLM, …). | | `PGBOT_AI_REASONING_EFFORT` | `none`, `low`, `medium`, `high`, `xhigh`, or `max` for reasoning models (OpenAI's default here is `xhigh`). | | `PGBOT_OPENAI_MODEL` / `PGBOT_OPENAI_URL` | Still honored: OpenAI-scoped model/endpoint override. | @@ -683,6 +683,7 @@ not require confirmation. | Anthropic | `ANTHROPIC_API_KEY` | `claude-opus-5` | `/v1/messages` | | OpenAI | `OPENAI_API_KEY` | `gpt-5.6-terra` | `/chat/completions` | | xAI | `XAI_API_KEY` / `GROK_API_KEY` | `grok-4.6` | `/responses` | +| Bedrock Mantle | AWS credential chain or `AWS_BEARER_TOKEN_BEDROCK` | `openai.gpt-5.6-terra` | Responses (GPT) / Messages (Claude) | The OpenAI provider also supports compatible services such as OpenRouter, Groq, Together, DeepSeek, Mistral, Ollama, vLLM, and LM Studio. @@ -696,7 +697,51 @@ Use `PGBOT_AI_PROVIDER` to select a provider explicitly. `PGBOT_AI_MODEL`, `PGBOT_AI_BASE_URL`, `PGBOT_AI_API_KEY`, and `PGBOT_AI_REASONING_EFFORT` override its defaults. Existing `PGBOT_GEMINI_MODEL` and `PGBOT_GEMINI_URL` and `PGBOT_OPENAI_MODEL` and `PGBOT_OPENAI_URL` settings remain supported. Keys -are read only from environment variables. +are read only from environment variables, except for Bedrock IAM credentials +resolved through the AWS SDK. + +For AWS Bedrock Mantle, select `bedrock` (or its alias `mantle`) and an AWS +profile. pgbot resolves IAM credentials itself through the AWS SDK's default +credential chain, including environment credentials, shared profiles, SSO, +`aws login` sessions, and workload roles. No AWS CLI subprocess is used. + +```sh +export PGBOT_AI_PROVIDER=bedrock +export AWS_PROFILE=your-profile +export AWS_REGION=us-east-1 +export PGBOT_AI_MODEL=openai.gpt-5.6-terra +pgbot ask "What needs attention?" --url "$DATABASE_URL" + +# Claude uses the Anthropic Messages API automatically: +export PGBOT_AI_MODEL=anthropic.claude-sonnet-5 +pgbot ask "What needs attention?" --url "$DATABASE_URL" +``` + +Authentication precedence is `PGBOT_AI_API_KEY`, then +`AWS_BEARER_TOKEN_BEDROCK`, then IAM credential resolution. Explicit tokens +bypass IAM and are not refreshed. Unset both token variables to use +`AWS_PROFILE`. For IAM authentication, the SDK caches and refreshes credentials; +pgbot mints a fresh bearer token per request, valid for at most 15 minutes and +never longer than the credentials. Renew an expired interactive login with +`aws login` or `aws sso login` as appropriate for your profile. + +Region precedence is `AWS_REGION`, `AWS_DEFAULT_REGION`, the AWS profile's +region (when using IAM), then `us-east-1`. OpenAI GPT models use the base URL +`https://bedrock-mantle..api.aws/openai/v1`, as documented by +[AWS for OpenAI GPT models](https://aws.amazon.com/blogs/machine-learning/get-started-with-openai-gpt-5-6-sol-terra-and-luna-on-amazon-bedrock/). +Models beginning with `anthropic.` use +`https://bedrock-mantle..api.aws/anthropic` and the Messages API. +Set `PGBOT_AI_MODEL` to the exact Bedrock model ID available to your account +and region, including for GPT-6 models. `PGBOT_AI_BASE_URL` overrides the base +URL, without the final `/responses` or `/v1/messages`. IAM authentication +requires a Mantle HTTPS hostname matching the configured region. Bedrock +requests do not follow redirects. + +Responses requests set `store=false`, omit sampling temperature for GPT-5/6 +reasoning models, and allow at least 32,000 output tokens (including hidden +reasoning). `PGBOT_AI_REASONING_EFFORT` is optional for Responses; when unset, +the service chooses its default. Claude keeps the existing Messages request +shape. Neither protocol retries inference automatically. **Exit codes** (a stable contract for CI): `0` clean · `1` warnings · `2` critical findings · `3` connection/execution failure · `64` usage error (bad flags/args). @@ -1128,11 +1173,18 @@ package is scoped. Use `npx @pgbot/cli`. Nothing leaves the machine unless you ask for it: every command except the AI layer is entirely local. The only commands that make an outbound call are `pgbot explain` and `pgbot ask`, which send the same PII-free Context to your configured -model — Gemini, Anthropic, OpenAI, xAI, or an OpenAI-compatible endpoint — and +model — Gemini, Anthropic, OpenAI, xAI, AWS Bedrock Mantle, or an +OpenAI-compatible endpoint — and say so, naming the provider, host, and model, with a confirmation prompt. A local endpoint (Ollama, vLLM, LM Studio on this machine) is identified as local and sends nothing off the box. +With Bedrock IAM authentication, the AWS SDK credential chain may also contact +AWS STS, SSO/OIDC, or container/instance metadata services to obtain or refresh +credentials. These authentication requests do not include the findings Context; +the Mantle bearer token is signed locally. Supplying `PGBOT_AI_API_KEY` or +`AWS_BEARER_TOKEN_BEDROCK` bypasses the AWS credential chain. + That Context is PII-free by construction: `pg_stat_statements` text is normalized (`$1` placeholders), and the one raw-SQL source (`pg_stat_activity` for blocking chains) is scrubbed of string/numeric literals, emails, and UUIDs before it can diff --git a/go.mod b/go.mod index 3a3d902..24d9bbf 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,8 @@ toolchain go1.27.1 require ( github.com/BurntSushi/toml v1.6.0 + github.com/aws/aws-sdk-go-v2 v1.46.0 + github.com/aws/aws-sdk-go-v2/config v1.33.3 github.com/charmbracelet/lipgloss v1.1.0 github.com/invopop/jsonschema v0.14.0 github.com/jackc/pgx/v5 v5.10.0 @@ -19,6 +21,18 @@ require ( ) require ( + github.com/aws/aws-sdk-go-v2/credentials v1.20.3 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.19.2 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.5.2 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.2 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.5.2 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.19 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.14.2 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.9.0 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.37.0 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.42.0 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.49.0 // indirect + github.com/aws/smithy-go v1.28.1 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/buger/jsonparser v1.1.2 // indirect diff --git a/go.sum b/go.sum index 065afa4..a61d90d 100644 --- a/go.sum +++ b/go.sum @@ -1,6 +1,34 @@ github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk= github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= github.com/apparentlymart/go-textseg/v13 v13.0.0/go.mod h1:ZK2fH7c4NqDTLtiYLvIkEghdlcqw7yxLeM89kiTRPUo= +github.com/aws/aws-sdk-go-v2 v1.46.0 h1:1kt7m/EKcEHt5mlyyxx9cSlMddRPIKbjb6DIQsu4HPk= +github.com/aws/aws-sdk-go-v2 v1.46.0/go.mod h1:bttEH6JqnUL8LepvDVfdrds/fZ5bCIxzpe3abyUrhDU= +github.com/aws/aws-sdk-go-v2/config v1.33.3 h1:h090b3O5S17bF87/0ysHZuIT/7DCb4EBRFQX2PMVPCw= +github.com/aws/aws-sdk-go-v2/config v1.33.3/go.mod h1:YYDB1kTejxbfAbEVUqgCtkVp26xvNCHev9cLKABMGAk= +github.com/aws/aws-sdk-go-v2/credentials v1.20.3 h1:tToOYM/LXev4NpfWlIYGDvBvjHmJ3HXpRU9ppl+pM6k= +github.com/aws/aws-sdk-go-v2/credentials v1.20.3/go.mod h1:wfGneWyncO7p67wqXV2IQhPk14JqIc25woKlaArT3WI= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.19.2 h1:Ldv7RPHs7qwwTscRjAl3YBud32f3BvdAGRmSvAx5L38= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.19.2/go.mod h1:XyK6UV8xbo66ysVqLd2783C09pBYHOm8aKTRV5DVJ30= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.5.2 h1:q/PSLGuRWCChWg+dLnb9dWOnrCxJtnboXbBtFoqqRrI= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.5.2/go.mod h1:TD1jvU2LvXkJexct5vBqcd8QlNXh5EmRUeL/Z32p0n4= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.2 h1:6fl86IPqKEXoySqiOWdfgbEp9OVbn44zTfEICNEBDhY= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.8.2/go.mod h1:63HDfhFkdzBpI8WGXTSKUHPKS6mqldj4u3LJW7RZtSU= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.5.2 h1:XMgIRS+uW9F3yFKnXGRrI9pkHi99CXTmoz2kz2/TGBA= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.5.2/go.mod h1:vorxDzK+n3jiv9a5ST/LG0Eu9cSv1CRdKTpG6pDMs+M= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.19 h1:bAdDl/HkGCcGPoe25ToSHEw23VIxt6CT5fLcg111BKg= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.19/go.mod h1:KaUzbLxv4CeSxh6ZCl9B4m7CuFenS8kUEaDs+f/DQr4= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.14.2 h1:ZtHYnumr6QyxhzEzNZwzQTFJEXOswrZqTTkRxthwvr4= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.14.2/go.mod h1:a1NXrYpBd311gBzn1UI5UzJyyvXktM4xNh/ydPiPpqY= +github.com/aws/aws-sdk-go-v2/service/signin v1.9.0 h1:c3k+k/CS4L+sAIH6fxikL+g5g2LpeNczaoyjjw1iMKI= +github.com/aws/aws-sdk-go-v2/service/signin v1.9.0/go.mod h1:AGIoQg99fBrOIQnF78TLx4lj18mc4gZ0hJx1UaLIFM4= +github.com/aws/aws-sdk-go-v2/service/sso v1.37.0 h1:+rqBaOq7jzInjY8M12hr+zEe85JpRll9BjMx38r33Ok= +github.com/aws/aws-sdk-go-v2/service/sso v1.37.0/go.mod h1:XFlVwUsw3sYh8Hw37umYVJnrcWrwXWRxbNw/JY0bblw= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.42.0 h1:hzM3GslEAOBcLn3DHH6ENToFi+vXP+n02W+x6zejAIM= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.42.0/go.mod h1:588e7skMkYIYkSUseT8E3WKFfbAfrA1bj3Zf+qxJtJY= +github.com/aws/aws-sdk-go-v2/service/sts v1.49.0 h1:N7Ey8obY3uSui+cxl0OUzFlFmkxSucoJnrniFhw+cLc= +github.com/aws/aws-sdk-go-v2/service/sts v1.49.0/go.mod h1:zMBwjSf4Pt8a1OHYiZ5rPD0PJRK1kQrUaxhS/Dbld8E= +github.com/aws/smithy-go v1.28.1 h1:R/nXH00c8qcfCzQVELtRw+eLQWtzv+VAIEFJ1/xxXlQ= +github.com/aws/smithy-go v1.28.1/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= diff --git a/internal/ai/anthropic.go b/internal/ai/anthropic.go index b5d9153..e7534bc 100644 --- a/internal/ai/anthropic.go +++ b/internal/ai/anthropic.go @@ -28,12 +28,18 @@ const ( // AnthropicProvider talks to the Messages API. type AnthropicProvider struct { + Label string APIKey string BaseURL string HTTP *http.Client } -func (p *AnthropicProvider) Name() string { return "anthropic" } +func (p *AnthropicProvider) Name() string { + if p.Label != "" { + return p.Label + } + return "anthropic" +} func (p *AnthropicProvider) LanguageModel(_ context.Context, modelID string) (LanguageModel, error) { if modelID == "" { @@ -47,7 +53,7 @@ type anthropicModel struct { model string } -func (m *anthropicModel) Provider() string { return "anthropic" } +func (m *anthropicModel) Provider() string { return m.provider.Name() } func (m *anthropicModel) Model() string { return m.model } func (m *anthropicModel) Endpoint() string { return m.provider.BaseURL } diff --git a/internal/ai/bedrock.go b/internal/ai/bedrock.go new file mode 100644 index 0000000..5f070d1 --- /dev/null +++ b/internal/ai/bedrock.go @@ -0,0 +1,117 @@ +package ai + +import ( + "context" + "crypto/sha256" + "encoding/base64" + "fmt" + "net/http" + "net/url" + "strconv" + "strings" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/aws/signer/v4" + "github.com/aws/aws-sdk-go-v2/config" +) + +func bedrockModel(model, base, key string, httpc *http.Client) (LanguageModel, error) { + region := firstEnv("AWS_REGION", "AWS_DEFAULT_REGION") + var cfg aws.Config + if key == "" { + var err error + cfg, err = config.LoadDefaultConfig(context.Background()) + if err != nil { + return nil, fmt.Errorf("load AWS configuration: %w", err) + } + if region == "" { + region = cfg.Region + } + } + if region == "" { + region = "us-east-1" + } + if model == "" { + model = "openai." + defaultOpenAIModel + } + anthropic := strings.HasPrefix(model, "anthropic.") + if base == "" { + base = "https://bedrock-mantle." + region + ".api.aws" + if anthropic { + base += "/anthropic" + } else { + base += "/openai/v1" + } + } + base = trimURL(base) + // Never forward a supplied or IAM-derived bearer token through a redirect. + httpc.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse } + if key == "" { + u, err := url.Parse(base) + if err != nil || u.Scheme != "https" || u.User != nil || u.RawQuery != "" || u.Fragment != "" || u.Host != "bedrock-mantle."+region+".api.aws" { + return nil, fmt.Errorf("IAM authentication requires a Bedrock Mantle HTTPS endpoint matching AWS region %s; set AWS_REGION to the endpoint region", region) + } + httpc.Transport = &bedrockAuth{credentials: cfg.Credentials, region: region, host: u.Host, anthropic: anthropic, next: http.DefaultTransport} + } + if anthropic { + p := &AnthropicProvider{APIKey: key, BaseURL: base, HTTP: httpc, Label: "bedrock"} + return p.LanguageModel(context.Background(), model) + } + p := &ResponsesProvider{APIKey: key, BaseURL: base, HTTP: httpc, Label: "bedrock", ReasoningEffort: envOr("PGBOT_AI_REASONING_EFFORT", "")} + return p.LanguageModel(context.Background(), model) +} + +// The SDK caches and refreshes credentials. Minting per request is local signing, +// so no separate token cache or refresh goroutine is needed. +type bedrockAuth struct { + credentials aws.CredentialsProvider + region, host string + anthropic bool + next http.RoundTripper +} + +func (a *bedrockAuth) RoundTrip(req *http.Request) (*http.Response, error) { + if req.URL.Scheme != "https" || req.URL.Host != a.host { + return nil, fmt.Errorf("refusing to send AWS credentials outside the configured Mantle endpoint") + } + creds, err := a.credentials.Retrieve(req.Context()) + if err != nil { + return nil, fmt.Errorf("resolve AWS credentials (AWS_PROFILE or the default credential chain): %w", err) + } + token, err := bedrockToken(req.Context(), creds, a.region, time.Now().UTC()) + if err != nil { + return nil, err + } + clone := req.Clone(req.Context()) + if a.anthropic { + clone.Header.Set("x-api-key", token) + } else { + clone.Header.Set("Authorization", "Bearer "+token) + } + return a.next.RoundTrip(clone) +} + +func bedrockToken(ctx context.Context, creds aws.Credentials, region string, now time.Time) (string, error) { + ttl := 15 * time.Minute + if creds.CanExpire && creds.Expires.Sub(now) < ttl { + ttl = creds.Expires.Sub(now) + } + if ttl < time.Second { + return "", fmt.Errorf("AWS credentials have expired; renew your AWS login") + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://bedrock.amazonaws.com/?Action=CallWithBearerToken", nil) + if err != nil { + return "", err + } + query := req.URL.Query() + query.Set("X-Amz-Expires", strconv.FormatInt(int64(ttl/time.Second), 10)) + req.URL.RawQuery = query.Encode() + // This request is only presigned, never sent. The empty payload hash matters: + // UNSIGNED-PAYLOAD yields a different signature and an invalid bearer token. + signed, _, err := v4.NewSigner().PresignHTTP(ctx, creds, req, fmt.Sprintf("%x", sha256.Sum256(nil)), "bedrock", region, now) + if err != nil { + return "", fmt.Errorf("sign Bedrock token: %w", err) + } + return "bedrock-api-key-" + base64.StdEncoding.EncodeToString([]byte(strings.TrimPrefix(signed, "https://")+"&Version=1")), nil +} diff --git a/internal/ai/bedrock_test.go b/internal/ai/bedrock_test.go new file mode 100644 index 0000000..b5867aa --- /dev/null +++ b/internal/ai/bedrock_test.go @@ -0,0 +1,212 @@ +package ai + +import ( + "context" + "encoding/base64" + "encoding/json" + "io" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" +) + +func tokenQuery(t *testing.T, token string) url.Values { + t.Helper() + if !strings.HasPrefix(token, "bedrock-api-key-") { + t.Fatal("missing token prefix") + } + data, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(token, "bedrock-api-key-")) + if err != nil { + t.Fatal(err) + } + u, err := url.Parse("https://" + string(data)) + if err != nil { + t.Fatal(err) + } + if u.Host != "bedrock.amazonaws.com" || u.Path != "/" { + t.Fatal("incorrect signing target") + } + return u.Query() +} + +func TestBedrockToken(t *testing.T) { + now := time.Date(2026, 9, 5, 12, 0, 0, 0, time.UTC) + creds := aws.Credentials{AccessKeyID: "AKIDEXAMPLE", SecretAccessKey: "dummy-secret", SessionToken: "session/+= token"} + token, err := bedrockToken(context.Background(), creds, "us-east-1", now) + if err != nil { + t.Fatal(err) + } + q := tokenQuery(t, token) + // Golden signature from AWS's Python aws-bedrock-token-generator with these + // dummy credentials, frozen timestamp, region, and 900-second expiry. + if q.Get("X-Amz-Signature") != "c51d43f3459d73b462fc95dca8da87f70d1a65920d0bfd32f1d6761e47485a2f" { + t.Fatal("signature differs from AWS reference generator") + } + if q.Get("Version") != "1" || q.Get("X-Amz-Expires") != "900" || q.Get("X-Amz-Security-Token") != creds.SessionToken { + t.Fatal("incorrect token envelope") + } + creds.CanExpire = true + creds.Expires = now.Add(90 * time.Second) + token, err = bedrockToken(context.Background(), creds, "us-east-1", now) + if err != nil { + t.Fatal(err) + } + if tokenQuery(t, token).Get("X-Amz-Expires") != "90" { + t.Fatal("token must not outlive credentials") + } + creds.Expires = now + if _, err := bedrockToken(context.Background(), creds, "us-east-1", now); err == nil { + t.Fatal("expired credentials accepted") + } + creds.CanExpire = false + creds.SessionToken = "" + token, err = bedrockToken(context.Background(), creds, "us-east-1", now) + if err != nil { + t.Fatal(err) + } + if _, exists := tokenQuery(t, token)["X-Amz-Security-Token"]; exists { + t.Fatal("static credentials must omit session token") + } +} + +type bedrockTestTransport func(*http.Request) (*http.Response, error) + +func (f bedrockTestTransport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +func isolateAWS(t *testing.T) { + t.Helper() + clearEnv(t) + for _, k := range []string{"AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "AWS_WEB_IDENTITY_TOKEN_FILE", "AWS_ROLE_ARN", "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", "AWS_CONTAINER_CREDENTIALS_FULL_URI"} { + t.Setenv(k, "") + } + t.Setenv("AWS_EC2_METADATA_DISABLED", "true") + t.Setenv("AWS_CONFIG_FILE", filepath.Join(t.TempDir(), "config")) + t.Setenv("AWS_SHARED_CREDENTIALS_FILE", filepath.Join(t.TempDir(), "credentials")) +} + +func TestBedrockIAMProfile(t *testing.T) { + for _, tc := range []struct { + model, base, path string + }{ + {"openai.gpt-5.6-terra", "", "/openai/v1/responses"}, + {"anthropic.claude-sonnet-5", "", "/anthropic/v1/messages"}, + {"anthropic.claude-sonnet-5", "https://bedrock-mantle.us-west-2.api.aws", "/v1/messages"}, + {"openai.gpt-5.6-terra", "https://bedrock-mantle.us-west-2.api.aws/anthropic", "/anthropic/responses"}, + } { + t.Run(tc.model+tc.path, func(t *testing.T) { + model := tc.model + isolateAWS(t) + t.Setenv("PGBOT_AI_PROVIDER", "bedrock") + t.Setenv("PGBOT_AI_MODEL", model) + t.Setenv("PGBOT_AI_BASE_URL", tc.base) + t.Setenv("AWS_PROFILE", "test-profile") + if err := os.WriteFile(os.Getenv("AWS_CONFIG_FILE"), []byte("[profile test-profile]\nregion = us-west-2\naws_access_key_id = AKIDEXAMPLE\naws_secret_access_key = dummy-secret\naws_session_token = session/+= token\n"), 0600); err != nil { + t.Fatal(err) + } + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + var client *http.Client + header := "Authorization" + switch m := m.(type) { + case *responsesModel: + client = m.provider.HTTP + case *anthropicModel: + client = m.provider.HTTP + header = "x-api-key" + default: + t.Fatalf("unexpected model type %T", m) + } + if m.Provider() != "bedrock" || !strings.Contains(m.Endpoint(), "us-west-2") { + t.Fatal("profile region or provider label lost") + } + auth := client.Transport.(*bedrockAuth) + calls := 0 + auth.next = bedrockTestTransport(func(r *http.Request) (*http.Response, error) { + calls++ + if r.URL.Path != tc.path { + t.Errorf("incorrect API path: %s", r.URL.Path) + } + q := tokenQuery(t, strings.TrimPrefix(r.Header.Get(header), "Bearer ")) + if !strings.Contains(q.Get("X-Amz-Credential"), "/us-west-2/bedrock/") { + t.Error("incorrect signing region") + } + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Error(err) + } + if body["model"] != model { + t.Error("model override lost") + } + if header == "x-api-key" && r.Header.Get("anthropic-version") != anthropicVersion { + t.Error("missing Anthropic version") + } + return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{"status":"completed","stop_reason":"end_turn","content":[{"type":"text","text":"OK"}],"output":[{"type":"message","content":[{"type":"output_text","text":"OK"}]}]}`))}, nil + }) + out, err := m.Generate(context.Background(), Call{Prompt: "hello"}) + if err != nil { + t.Fatal(err) + } + if out.Text != "OK" || calls != 1 { + t.Fatal("generation failed") + } + if client.CheckRedirect == nil || client.CheckRedirect(nil, nil) != http.ErrUseLastResponse { + t.Fatal("credentialed redirects must be disabled") + } + r, _ := http.NewRequest("POST", "https://example.com/responses", nil) + if _, err := auth.RoundTrip(r); err == nil || calls != 1 { + t.Fatal("credentials sent outside Mantle") + } + }) + } +} + +func TestBedrockAuthConfiguration(t *testing.T) { + isolateAWS(t) + t.Setenv("PGBOT_AI_PROVIDER", "bedrock") + t.Setenv("AWS_PROFILE", "nonexistent") + t.Setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-override") + t.Setenv("PGBOT_AI_API_KEY", "explicit-override") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + p := m.(*responsesModel).provider + if p.APIKey != "explicit-override" || p.HTTP.Transport != nil { + t.Fatal("explicit token must bypass IAM") + } + t.Setenv("PGBOT_AI_API_KEY", "") + m, err = Resolve() + if err != nil { + t.Fatal(err) + } + if m.(*responsesModel).provider.APIKey != "bedrock-override" { + t.Fatal("Bedrock token override lost") + } + t.Setenv("AWS_BEARER_TOKEN_BEDROCK", "") + t.Setenv("AWS_PROFILE", "") + t.Setenv("OPENAI_API_KEY", "unrelated-key") + m, err = Resolve() + if err != nil { + t.Fatal(err) + } + if _, err = m.Generate(context.Background(), Call{Prompt: "hello"}); err == nil || !strings.Contains(err.Error(), "resolve AWS credentials") { + t.Fatalf("expected missing IAM credentials error, got %v", err) + } + t.Setenv("PGBOT_AI_BASE_URL", "https://example.com/openai/v1") + if _, err := Resolve(); err == nil { + t.Fatal("IAM authentication must reject non-Mantle endpoints") + } + t.Setenv("AWS_REGION", "us-east-1") + t.Setenv("PGBOT_AI_BASE_URL", "https://bedrock-mantle.us-west-2.api.aws/openai/v1") + if _, err := Resolve(); err == nil { + t.Fatal("IAM region mismatch accepted") + } +} diff --git a/internal/ai/openai.go b/internal/ai/openai.go index 5f54592..09d3bd7 100644 --- a/internal/ai/openai.go +++ b/internal/ai/openai.go @@ -85,7 +85,8 @@ func reasoningModel(id string) bool { if i := strings.LastIndex(id, "/"); i >= 0 { // strip an "openai/" vendor prefix id = id[i+1:] } - for _, p := range []string{"gpt-5", "o1", "o3", "o4"} { + id = strings.TrimPrefix(id, "openai.") // Bedrock model IDs + for _, p := range []string{"gpt-5", "gpt-6", "o1", "o3", "o4"} { if strings.HasPrefix(id, p) { return true } diff --git a/internal/ai/provider.go b/internal/ai/provider.go index 28a681c..a22ab06 100644 --- a/internal/ai/provider.go +++ b/internal/ai/provider.go @@ -5,8 +5,8 @@ // // The model is yours to choose: Gemini, Anthropic, OpenAI, or any OpenAI-compatible // endpoint (OpenRouter, Groq, Together, DeepSeek, xAI, Mistral, Ollama, vLLM, -// LM Studio). Each provider is a few hundred lines of net/http so pgbot keeps its -// single-static-binary, minimal-dependency promise — no vendor SDKs. +// LM Studio), plus AWS Bedrock Mantle. Inference uses net/http; AWS credential +// resolution and signing use the AWS SDK. package ai import ( @@ -18,13 +18,8 @@ import ( // Provider is a named source of language models. // -// This mirrors charmbracelet/fantasy's Provider/LanguageModel pair on purpose, so -// a fantasy-backed implementation could drop in later — but we implement it over -// net/http instead of depending on fantasy, which pulls the real vendor SDKs -// (anthropic-sdk-go, openai-go, google.golang.org/genai, aws-sdk-go-v2) and takes -// the binary from 23 MB to ~65 MB for one non-streaming POST. The interface is -// narrowed to the single call pgbot makes: one system turn, one user turn, no -// tools, no streaming. +// Providers use net/http for inference. Bedrock additionally uses the AWS SDK +// for its credential chain and SigV4 signing; it reuses the same wire clients. type Provider interface { Name() string LanguageModel(ctx context.Context, modelID string) (LanguageModel, error) @@ -33,7 +28,7 @@ type Provider interface { // LanguageModel is one model at one endpoint, ready to answer a single turn. type LanguageModel interface { Generate(ctx context.Context, c Call) (*Response, error) - Provider() string // "gemini" | "openai" | "anthropic" | "xai" + Provider() string // "gemini" | "openai" | "anthropic" | "xai" | "bedrock" Model() string // resolved model id — shown in the AI banner Endpoint() string // base URL we POST to — powers the consent prompt } diff --git a/internal/ai/resolve.go b/internal/ai/resolve.go index d059683..c62680b 100644 --- a/internal/ai/resolve.go +++ b/internal/ai/resolve.go @@ -10,12 +10,12 @@ import ( ) // Resolve builds the model to use from the environment. Keys come ONLY from the -// environment — never a flag — so they can't leak into shell history or the -// process list. That invariant is enforced here, once, for every provider. +// environment (or the AWS credential chain for Bedrock), never a flag, so they +// cannot leak into shell history or the process list. // // Precedence: // -// PGBOT_AI_PROVIDER explicit: gemini | openai | anthropic | xai +// PGBOT_AI_PROVIDER explicit: gemini | openai | anthropic | xai | bedrock // otherwise auto-detected from whichever key is set, // OpenAI first to preserve existing behavior // PGBOT_AI_MODEL model id (else the provider's default) @@ -74,6 +74,12 @@ func Resolve() (LanguageModel, error) { } p = &AnthropicProvider{APIKey: key, BaseURL: trimURL(base), HTTP: httpc} + case "bedrock", "mantle": + if key == "" { + key = firstEnv("AWS_BEARER_TOKEN_BEDROCK") + } + return bedrockModel(model, base, key, httpc) + case "xai", "grok", "responses": // The Responses API, which xAI documents as its primary interface. Same // endpoint shape at OpenAI, so PGBOT_AI_PROVIDER=responses + an OpenAI key @@ -144,7 +150,7 @@ func Resolve() (LanguageModel, error) { } default: - return nil, fmt.Errorf("unknown PGBOT_AI_PROVIDER %q (want gemini, openai, anthropic, or xai)", name) + return nil, fmt.Errorf("unknown PGBOT_AI_PROVIDER %q (want gemini, openai, anthropic, xai, responses, bedrock, or mantle)", name) } // A local endpoint (Ollama, vLLM, LM Studio) usually has no key at all, and diff --git a/internal/ai/resolve_test.go b/internal/ai/resolve_test.go index 0e0763f..dfaface 100644 --- a/internal/ai/resolve_test.go +++ b/internal/ai/resolve_test.go @@ -13,6 +13,7 @@ func clearEnv(t *testing.T) { "PGBOT_AI_PROVIDER", "PGBOT_AI_MODEL", "PGBOT_AI_BASE_URL", "PGBOT_AI_API_KEY", "PGBOT_AI_REASONING_EFFORT", "GEMINI_API_KEY", "GOOGLE_API_KEY", "ANTHROPIC_API_KEY", "OPENAI_API_KEY", "OPENROUTER_API_KEY", "XAI_API_KEY", "GROK_API_KEY", + "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION", "AWS_DEFAULT_REGION", "PGBOT_GEMINI_MODEL", "PGBOT_GEMINI_URL", "PGBOT_OPENAI_MODEL", "PGBOT_OPENAI_URL", } { t.Setenv(k, "") @@ -180,7 +181,7 @@ func TestResolve_remoteEndpointRequiresKey(t *testing.T) { func TestResolve_unknownProvider(t *testing.T) { clearEnv(t) - t.Setenv("PGBOT_AI_PROVIDER", "bedrock") + t.Setenv("PGBOT_AI_PROVIDER", "unknown-provider") t.Setenv("PGBOT_AI_API_KEY", "k") if _, err := Resolve(); err == nil || !strings.Contains(err.Error(), "unknown") { t.Errorf("unknown provider should be rejected clearly, got %v", err) diff --git a/internal/ai/responses.go b/internal/ai/responses.go index 4e53f4c..28a20d2 100644 --- a/internal/ai/responses.go +++ b/internal/ai/responses.go @@ -15,21 +15,15 @@ const ( defaultXAIURL = "https://api.x.ai/v1" ) -// ResponsesProvider speaks the Responses API (POST /responses) — the newer -// surface both xAI and OpenAI prefer over /chat/completions. pgbot uses it for -// xAI, where it is the documented primary interface. -// -// It is deliberately NOT the default for the OpenAI-compatible world: only -// OpenAI and xAI implement /responses, while Ollama, vLLM, LM Studio, Groq, -// Together, DeepSeek and Mistral implement only /chat/completions. This provider -// is additive — OpenAIProvider stays the compatibility path. +// ResponsesProvider speaks POST /responses for xAI, OpenAI, and Bedrock Mantle. +// OpenAIProvider remains the /chat/completions compatibility path. type ResponsesProvider struct { APIKey string BaseURL string HTTP *http.Client // Label is the provider name shown in the consent prompt and AI banner - // ("xai", "openai") — the endpoint is shared, the vendor is not. + // ("xai", "openai", "bedrock") — the endpoint is shared, the vendor is not. Label string // ReasoningEffort is sent as reasoning.effort when set. Left empty by default @@ -70,7 +64,7 @@ type responsesRequest struct { // Store is explicitly false. The Responses API defaults it to TRUE, which // retains the request server-side for later retrieval — a quiet downgrade of // the disclosure `pgbot explain` asks the user to consent to. We send the - // findings once and keep nothing on the vendor's side. + // findings without enabling stored conversation state. Store bool `json:"store"` MaxOutputTokens *int64 `json:"max_output_tokens,omitempty"` Temperature *float64 `json:"temperature,omitempty"` @@ -115,6 +109,9 @@ func (m *responsesModel) Generate(ctx context.Context, c Call) (*Response, error MaxOutputTokens: &limit, Temperature: c.Temperature, } + if reasoningModel(m.model) { + reqBody.Temperature = nil + } if e := m.provider.ReasoningEffort; e != "" { reqBody.Reasoning = &reasoningCfg{Effort: e} } diff --git a/internal/ai/responses_test.go b/internal/ai/responses_test.go index ec542a2..125d9c5 100644 --- a/internal/ai/responses_test.go +++ b/internal/ai/responses_test.go @@ -243,3 +243,55 @@ func TestResponses_maxOutputTokensFloor(t *testing.T) { t.Errorf("max_output_tokens = %v with a 50000 hint; want the hint to win above the floor", sent) } } + +func TestBedrockResponses(t *testing.T) { + for _, model := range []string{"openai.gpt-5.6-terra", "openai.gpt-6-astra"} { + t.Run(model, func(t *testing.T) { + clearEnv(t) + t.Setenv("PGBOT_AI_PROVIDER", "bedrock") + t.Setenv("AWS_BEARER_TOKEN_BEDROCK", "dummy-token") + t.Setenv("AWS_REGION", "us-west-2") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + if m.Provider() != "bedrock" || m.Model() != "openai."+defaultOpenAIModel || m.Endpoint() != "https://bedrock-mantle.us-west-2.api.aws/openai/v1" { + t.Fatal("incorrect Bedrock defaults") + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" || r.URL.Path != "/openai/v1/responses" || r.Header.Get("Authorization") != "Bearer dummy-token" { + t.Error("incorrect request route or authentication") + } + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Error(err) + } + if body["model"] != model || body["store"] != false || body["max_output_tokens"] != float64(reasoningTokenFloor) || body["instructions"] != "system" || body["input"] != "report" { + t.Errorf("incorrect request: %v", body) + } + if _, ok := body["temperature"]; ok { + t.Error("reasoning model must omit temperature") + } + if body["reasoning"].(map[string]any)["effort"] != "low" { + t.Error("missing effort") + } + io.WriteString(w, `{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"Healthy."}]}]}`) + })) + defer srv.Close() + t.Setenv("PGBOT_AI_BASE_URL", srv.URL+"/openai/v1/") + t.Setenv("PGBOT_AI_MODEL", model) + t.Setenv("PGBOT_AI_REASONING_EFFORT", "low") + m, err = Resolve() + if err != nil { + t.Fatal(err) + } + out, err := m.Generate(context.Background(), Call{System: "system", Prompt: "report", Temperature: f64(0.2), MaxOutputTokens: i64(8192)}) + if err != nil { + t.Fatal(err) + } + if out.Text != "Healthy." { + t.Errorf("unexpected answer: %q", out.Text) + } + }) + } +}