diff --git a/CHANGELOG.md b/CHANGELOG.md index 53fd9f4..ab2770d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,19 @@ 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 as an `explain` / `ask` provider** (#35, contributed by + @edwardsb). `PGBOT_AI_PROVIDER=bedrock` (alias `mantle`) routes `openai.*` + models through the Responses API and `anthropic.*` models through the + Messages API on `bedrock-mantle..api.aws`; the default is + `openai.gpt-5.6-terra`. Authenticate with `AWS_BEARER_TOKEN_BEDROCK` or with + `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` / `AWS_SESSION_TOKEN`, from + which pgbot mints Bedrock's short-lived bearer token itself — a hundred lines + of SigV4 over the standard library, pinned to AWS's reference generator by a + golden-signature test. No AWS SDK, no config files, no STS or metadata calls: + a profile or SSO login is exported with + `eval "$(aws configure export-credentials --format env)"`. Access keys only + ever go to the Mantle host for the configured region, and Bedrock requests + never follow redirects. ## [0.8.1] - 2026-09-06 diff --git a/README.md b/README.md index 3ed5dbc..c53eebb 100644 --- a/README.md +++ b/README.md @@ -438,7 +438,8 @@ 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). | +| `AWS_BEARER_TOKEN_BEDROCK`, or `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` / `AWS_SESSION_TOKEN` | Enables `ask` / `explain` via AWS Bedrock Mantle with `PGBOT_AI_PROVIDER=bedrock` (never auto-detected). `AWS_REGION` picks the endpoint; `AWS_CREDENTIAL_EXPIRATION` bounds the minted token. | +| `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 +684,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_BEARER_TOKEN_BEDROCK`, or `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` | `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. @@ -698,6 +700,50 @@ 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. +For AWS Bedrock Mantle, select `bedrock` (or its alias `mantle`) and give pgbot +AWS credentials the same way you give it any other key — through the environment. +No AWS SDK, no config files read, no calls to STS or instance metadata: pgbot +mints Bedrock's bearer token itself from the standard three variables. + +```sh +export PGBOT_AI_PROVIDER=bedrock +export AWS_REGION=us-east-1 +# a profile, an SSO login, or an assumed role becomes the three variables: +eval "$(aws configure export-credentials --profile your-profile --format env)" +pgbot ask "What needs attention?" "$DATABASE_URL" + +# Claude uses the Anthropic Messages API automatically: +export PGBOT_AI_MODEL=anthropic.claude-sonnet-5 +pgbot ask "What needs attention?" "$DATABASE_URL" +``` + +Authentication precedence is `PGBOT_AI_API_KEY`, then `AWS_BEARER_TOKEN_BEDROCK` +(a Bedrock API key from the console), then `AWS_ACCESS_KEY_ID` / +`AWS_SECRET_ACCESS_KEY` (plus `AWS_SESSION_TOKEN` for temporary credentials). +An explicit token is sent as-is. From access keys pgbot mints a fresh bearer +token per request, valid for at most 15 minutes and never past +`AWS_CREDENTIAL_EXPIRATION` when that is set (the export command sets it). +`AWS_PROFILE` on its own does not authenticate — export it as above. + +Region precedence is `AWS_REGION`, `AWS_DEFAULT_REGION`, then `us-east-1`. The +default model is `openai.gpt-5.6-terra`. OpenAI GPT models use +`https://bedrock-mantle..api.aws/openai/v1` and the Responses API, 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`. Access keys are only ever sent +to the Mantle HTTPS host for the configured region, and Bedrock requests never +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). Suppressed findings never contribute to the exit code. @@ -1128,8 +1174,11 @@ 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 -say so, naming the provider, host, and model, with a confirmation prompt. A +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. Bedrock's token is minted locally from your +environment credentials; pgbot never reads AWS config files or calls STS or +instance metadata. A local endpoint (Ollama, vLLM, LM Studio on this machine) is identified as local and sends nothing off the box. 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..e70e2e5 --- /dev/null +++ b/internal/ai/bedrock.go @@ -0,0 +1,219 @@ +package ai + +// AWS Bedrock Mantle: OpenAI GPT models through the Responses API and Claude +// through the Messages API, on AWS's endpoints, with AWS credentials. +// +// No AWS SDK. Bedrock authenticates with a bearer token that is nothing more +// than a SigV4-presigned URL of `https://bedrock.amazonaws.com/?Action= +// CallWithBearerToken`, base64-encoded — about a hundred lines of HMAC below. +// Credentials come from the three variables every AWS tool understands +// (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN); a profile, an +// SSO login, or an assumed role becomes those with one command: +// +// eval "$(aws configure export-credentials --profile prod --format env)" +// +// That keeps pgbot's promise: keys only from the environment, no config files +// read, no calls to STS or instance metadata, and one static binary. + +import ( + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "errors" + "fmt" + "net/http" + "net/url" + "sort" + "strconv" + "strings" + "time" +) + +const ( + // A minted token lives this long at most (the service allows up to 12h; + // pgbot makes one call and has no reason to hold a longer-lived secret). + bedrockTokenTTL = 15 * time.Minute + bedrockSTSHost = "bedrock.amazonaws.com" +) + +// awsCredentials are the environment credentials a token is minted from. +type awsCredentials struct { + AccessKeyID string + SecretAccessKey string + SessionToken string // empty for long-lived keys + Expires time.Time // zero when not known +} + +// awsCredentialsFromEnv reads the standard variables. AWS_CREDENTIAL_EXPIRATION +// is what `aws configure export-credentials --format env` emits next to them; +// honoring it keeps a minted token from outliving the credentials behind it. +func awsCredentialsFromEnv() (awsCredentials, error) { + c := awsCredentials{ + AccessKeyID: envOr("AWS_ACCESS_KEY_ID", ""), + SecretAccessKey: envOr("AWS_SECRET_ACCESS_KEY", ""), + SessionToken: envOr("AWS_SESSION_TOKEN", ""), + } + if c.AccessKeyID == "" || c.SecretAccessKey == "" { + return c, errors.New("no AWS credentials for Bedrock — set AWS_BEARER_TOKEN_BEDROCK (a Bedrock API key), " + + "or AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY (and AWS_SESSION_TOKEN) and pgbot mints a short-lived token; " + + "a profile or SSO login exports them with: eval \"$(aws configure export-credentials --format env)\"") + } + if exp := envOr("AWS_CREDENTIAL_EXPIRATION", ""); exp != "" { + t, err := time.Parse(time.RFC3339, exp) + if err != nil { + return c, fmt.Errorf("AWS_CREDENTIAL_EXPIRATION %q is not an RFC 3339 timestamp", exp) + } + c.Expires = t + } + return c, nil +} + +func bedrockModel(model, base, key string, httpc *http.Client) (LanguageModel, error) { + region := firstEnv("AWS_REGION", "AWS_DEFAULT_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 minted bearer token through a redirect. + httpc.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse } + if key == "" { + creds, err := awsCredentialsFromEnv() + if err != nil { + return nil, err + } + 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("AWS credentials are only sent to the Bedrock Mantle HTTPS endpoint for region %s; set AWS_REGION to the endpoint's region", region) + } + httpc.Transport = &bedrockAuth{creds: creds, 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) +} + +// bedrockAuth mints a fresh token per request. Minting is local HMAC work, so +// there is nothing to cache or refresh; the token's TTL is bounded by the +// credentials' own expiry. +type bedrockAuth struct { + creds awsCredentials + region, host string + anthropic bool // Messages API takes x-api-key; the Responses API a Bearer + 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, errors.New("refusing to send AWS credentials outside the configured Mantle endpoint") + } + token, err := bedrockToken(a.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) +} + +// bedrockToken builds the bearer token AWS's own token generators produce: a +// SigV4 query-presigned POST to bedrock.amazonaws.com?Action=CallWithBearerToken +// (empty-payload hash; UNSIGNED-PAYLOAD yields an invalid token), with the +// scheme stripped and "&Version=1" appended, base64-encoded, prefixed. +func bedrockToken(creds awsCredentials, region string, now time.Time) (string, error) { + ttl := bedrockTokenTTL + if !creds.Expires.IsZero() && creds.Expires.Sub(now) < ttl { + ttl = creds.Expires.Sub(now) + } + if ttl < time.Second { + return "", errors.New("AWS credentials have expired; renew your AWS login and export them again") + } + amzDate := now.UTC().Format("20060102T150405Z") + scope := amzDate[:8] + "/" + region + "/bedrock/aws4_request" + params := map[string]string{ + "Action": "CallWithBearerToken", + "X-Amz-Algorithm": "AWS4-HMAC-SHA256", + "X-Amz-Credential": creds.AccessKeyID + "/" + scope, + "X-Amz-Date": amzDate, + "X-Amz-Expires": strconv.FormatInt(int64(ttl/time.Second), 10), + "X-Amz-SignedHeaders": "host", + } + if creds.SessionToken != "" { + params["X-Amz-Security-Token"] = creds.SessionToken + } + query := sigv4Query(params) + emptyPayload := sha256.Sum256(nil) + canonical := strings.Join([]string{ + http.MethodPost, "/", query, + "host:" + bedrockSTSHost, "", // canonical headers, then the blank line + "host", hex.EncodeToString(emptyPayload[:]), + }, "\n") + canonicalHash := sha256.Sum256([]byte(canonical)) + toSign := strings.Join([]string{"AWS4-HMAC-SHA256", amzDate, scope, hex.EncodeToString(canonicalHash[:])}, "\n") + key := []byte("AWS4" + creds.SecretAccessKey) + for _, part := range []string{amzDate[:8], region, "bedrock", "aws4_request"} { + key = hmacSHA256(key, part) + } + signature := hex.EncodeToString(hmacSHA256(key, toSign)) + presigned := bedrockSTSHost + "/?" + query + "&X-Amz-Signature=" + signature + return "bedrock-api-key-" + base64.StdEncoding.EncodeToString([]byte(presigned+"&Version=1")), nil +} + +// sigv4Query renders params as SigV4's canonical query string: keys sorted, +// every key and value RFC 3986-encoded (only unreserved characters bare, hex +// upper-case, space as %20 — not the form encoding net/url produces). +func sigv4Query(params map[string]string) string { + keys := make([]string, 0, len(params)) + for k := range params { + keys = append(keys, k) + } + sort.Strings(keys) + parts := make([]string, 0, len(keys)) + for _, k := range keys { + parts = append(parts, sigv4Escape(k)+"="+sigv4Escape(params[k])) + } + return strings.Join(parts, "&") +} + +func sigv4Escape(s string) string { + const hexDigits = "0123456789ABCDEF" + var b strings.Builder + for i := 0; i < len(s); i++ { + c := s[i] + switch { + case 'A' <= c && c <= 'Z', 'a' <= c && c <= 'z', '0' <= c && c <= '9', c == '-', c == '_', c == '.', c == '~': + b.WriteByte(c) + default: + b.WriteByte('%') + b.WriteByte(hexDigits[c>>4]) + b.WriteByte(hexDigits[c&15]) + } + } + return b.String() +} + +func hmacSHA256(key []byte, data string) []byte { + h := hmac.New(sha256.New, key) + h.Write([]byte(data)) + return h.Sum(nil) +} diff --git a/internal/ai/bedrock_test.go b/internal/ai/bedrock_test.go new file mode 100644 index 0000000..2796527 --- /dev/null +++ b/internal/ai/bedrock_test.go @@ -0,0 +1,246 @@ +package ai + +import ( + "context" + "encoding/base64" + "encoding/json" + "io" + "net/http" + "net/url" + "strings" + "testing" + "time" +) + +func tokenQuery(t *testing.T, token string) url.Values { + t.Helper() + if !strings.HasPrefix(token, "bedrock-api-key-") { + t.Fatalf("missing token prefix in %q", token) + } + 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 := awsCredentials{AccessKeyID: "AKIDEXAMPLE", SecretAccessKey: "dummy-secret", SessionToken: "session/+= token"} + token, err := bedrockToken(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 — the + // proof that the hand-rolled presigner matches the SDK it replaced. + if q.Get("X-Amz-Signature") != "c51d43f3459d73b462fc95dca8da87f70d1a65920d0bfd32f1d6761e47485a2f" { + t.Fatalf("signature differs from AWS reference generator: %s", q.Get("X-Amz-Signature")) + } + if q.Get("Version") != "1" || q.Get("X-Amz-Expires") != "900" || q.Get("X-Amz-Security-Token") != creds.SessionToken { + t.Fatal("incorrect token envelope") + } + if q.Get("X-Amz-Credential") != "AKIDEXAMPLE/20260905/us-east-1/bedrock/aws4_request" || q.Get("X-Amz-SignedHeaders") != "host" { + t.Fatal("incorrect credential scope") + } + + creds.Expires = now.Add(90 * time.Second) + token, err = bedrockToken(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(creds, "us-east-1", now); err == nil { + t.Fatal("expired credentials accepted") + } + creds.Expires = time.Time{} + creds.SessionToken = "" + token, err = bedrockToken(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") + } +} + +// The canonical query must use RFC 3986 escaping, not net/url's form encoding: +// a space is %20 and '+' is %2B, or the signature does not verify. +func TestSigv4Escape(t *testing.T) { + if got := sigv4Escape("session/+= token~-_."); got != "session%2F%2B%3D%20token~-_." { + t.Fatalf("sigv4Escape = %q", got) + } + if got := sigv4Query(map[string]string{"b": "2", "A": "1", "X-Amz-Date": "x"}); got != "A=1&X-Amz-Date=x&b=2" { + t.Fatalf("sigv4Query ordering = %q", got) + } +} + +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_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "AWS_CREDENTIAL_EXPIRATION", + "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION", "AWS_DEFAULT_REGION", "AWS_PROFILE"} { + t.Setenv(k, "") + } +} + +func TestBedrockEnvCredentials(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_REGION", "us-west-2") + t.Setenv("AWS_ACCESS_KEY_ID", "AKIDEXAMPLE") + t.Setenv("AWS_SECRET_ACCESS_KEY", "dummy-secret") + t.Setenv("AWS_SESSION_TOKEN", "session/+= token") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + var client *http.Client + header, other := "Authorization", "x-api-key" + switch m := m.(type) { + case *responsesModel: + client = m.provider.HTTP + case *anthropicModel: + client = m.provider.HTTP + header, other = "x-api-key", "Authorization" + default: + t.Fatalf("unexpected model type %T", m) + } + if m.Provider() != "bedrock" || !strings.Contains(m.Endpoint(), "us-west-2") { + t.Fatal("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) + } + // The header follows the model family, never the URL path. + q := tokenQuery(t, strings.TrimPrefix(r.Header.Get(header), "Bearer ")) + if v := r.Header.Get(other); strings.Contains(v, "bedrock-api-key-") { + t.Errorf("token also sent in %s", other) + } + if !strings.Contains(q.Get("X-Amz-Credential"), "/us-west-2/bedrock/") || q.Get("X-Amz-Security-Token") != "session/+= token" { + t.Error("incorrect signing region or session token") + } + 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_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("an explicit token must be used as-is, with no minting transport") + } + if m.Model() != "openai."+defaultOpenAIModel || m.Endpoint() != "https://bedrock-mantle.us-east-1.api.aws/openai/v1" { + t.Fatalf("defaults: model=%s endpoint=%s", m.Model(), m.Endpoint()) + } + 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") + } + + // No token and no access keys: say exactly what to set, and never touch a + // profile or the AWS config files. + t.Setenv("AWS_BEARER_TOKEN_BEDROCK", "") + t.Setenv("AWS_PROFILE", "some-profile") + t.Setenv("OPENAI_API_KEY", "unrelated-key") + _, err = Resolve() + if err == nil || !strings.Contains(err.Error(), "aws configure export-credentials") { + t.Fatalf("expected a missing-credentials error naming the export command, got %v", err) + } + + t.Setenv("AWS_ACCESS_KEY_ID", "AKIDEXAMPLE") + t.Setenv("AWS_SECRET_ACCESS_KEY", "dummy-secret") + t.Setenv("PGBOT_AI_BASE_URL", "https://example.com/openai/v1") + if _, err := Resolve(); err == nil { + t.Fatal("access keys must not be sent to a non-Mantle endpoint") + } + 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("region mismatch accepted") + } + + // AWS_CREDENTIAL_EXPIRATION (emitted by `aws configure export-credentials`) + // bounds the minted token; an unparseable value is refused rather than ignored. + t.Setenv("PGBOT_AI_BASE_URL", "") + t.Setenv("AWS_CREDENTIAL_EXPIRATION", "not-a-time") + if _, err := Resolve(); err == nil { + t.Fatal("malformed AWS_CREDENTIAL_EXPIRATION accepted") + } + t.Setenv("AWS_CREDENTIAL_EXPIRATION", time.Now().UTC().Add(-time.Minute).Format(time.RFC3339)) + m, err = Resolve() + if err != nil { + t.Fatal(err) + } + if _, err := m.Generate(context.Background(), Call{Prompt: "hello"}); err == nil || !strings.Contains(err.Error(), "expired") { + t.Fatalf("expired credentials should fail at token minting, got %v", err) + } +} 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..ad029fc 100644 --- a/internal/ai/provider.go +++ b/internal/ai/provider.go @@ -3,10 +3,11 @@ // already-computed, PII-free Context and asks a model to explain and prioritize // it in plain language. Everything it emits is labeled as model-generated. // -// 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. +// The model is yours to choose: Gemini, Anthropic, OpenAI, AWS Bedrock Mantle, 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. Even Bedrock's SigV4 token is a hundred lines of HMAC (bedrock.go). package ai import ( @@ -22,9 +23,9 @@ import ( // 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. +// the binary from 23 MB to ~65 MB for one non-streaming POST (aws-sdk-go-v2 alone +// measured +3 MB and 14 modules). The interface is narrowed to the single call +// pgbot makes: one system turn, one user turn, no tools, no streaming. type Provider interface { Name() string LanguageModel(ctx context.Context, modelID string) (LanguageModel, error) @@ -33,7 +34,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..5b88de9 100644 --- a/internal/ai/resolve.go +++ b/internal/ai/resolve.go @@ -10,12 +10,13 @@ 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 — never a flag, never a config file — so they can't leak into shell +// history or the process list. That invariant is enforced here, once, for every +// provider; Bedrock's AWS access keys are environment variables like any other. // // 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 +75,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 +151,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..d8f6e43 100644 --- a/internal/ai/responses.go +++ b/internal/ai/responses.go @@ -17,19 +17,20 @@ const ( // 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. +// xAI, where it is the documented primary interface, and for OpenAI models on +// Bedrock Mantle. // // 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. +// OpenAI, xAI and Mantle 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. 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 @@ -115,6 +116,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) + } + }) + } +}