mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(terraform): keep unconfigured allowed_routes plan-known and unsent (#44487)
* feat(terraform): expose key type on virtual keys
* docs(terraform): remove in-tree key type docs
* fix(terraform): preserve server-derived key routes
* fix(terraform): keep unconfigured key routes plan-known and unsent
Two regressions from exposing key_type on litellm_key:
1. Marking allowed_routes Computed makes an omitted attribute unknown at
plan time ("known only after apply"), so any plan that consumes it
before the key exists fails, e.g.
for_each = toset(coalesce(litellm_key.x.allowed_routes, [])).
Computed is dropped again; server-derived routes still land in state
through reads, and a DiffSuppressFunc keyed on the raw config keeps a
config that never declares the attribute from showing a perpetual
removal diff against those routes (a config that shrinks the list or
sets it still diffs).
2. mapResourceDataToKey copies allowed_routes unconditionally and
UpdateKey sends it when non-empty, so once reads materialize the
server's routes into state, every update re-asserts them: an
alias-only rename POSTs allowed_routes (the pre-key_type provider
sent none), and with stale state (-refresh=false) it silently
overwrites routes managed outside Terraform. Updates now omit the
field whenever the raw config does not declare it.
The key_type flow is unchanged: create still sends key_type, the proxy
presets the routes, reads materialize them into state, and plans stay
drift-free.
* fix(terraform): reject allowed_routes alongside a presetting key_type
The proxy derives allowed_routes from the key_type preset and overwrites
whatever the request declared, so a config combining the two could never
match what gets stored: the key came back with the preset routes and
drifted against the declared list on every plan. A CustomizeDiff now
fails the plan with an actionable message when a presetting key_type
(llm_api, management, read_only) is combined with allowed_routes.
key_type "default" presets nothing and keeps declared routes.
* fix(terraform): scope key_type route rejection to create-shaped plans
/key/update stores an explicit allowed_routes verbatim and never reapplies
the key_type preset, so an existing or imported typed key can manage its
routes in place. Only plans that create a key (fresh, or a replacement
that changes key_type) still reject the combination, because there the
preset always overwrites the declared list. A replacement forced by
another ForceNew attribute converges on the next apply, which re-sends
the declared routes.
* fix(terraform): restore declared routes on typed key creation
/key/generate replaces a declared allowed_routes with the key_type
preset while /key/update stores the list verbatim, so any create that
carries both (a fresh key, or a replacement forced by key_type or
another ForceNew attribute) used to leave the key holding the preset
instead of the declared routes until a second apply. When the generate
response does not match the declared list, create now follows up with an
update that re-sends the full create payload against the new key hash,
so the first apply already stores the declared routes. This also
replaces the plan-time rejection of the combination: every config shape
now converges, and existing typed keys keep managing routes in place as
before.
* fix(terraform): delete the key when a route restore fails at create
If /key/generate succeeds but the restore update is rejected, the key
exists server-side while terraform holds no state for it: an active key
with the type preset would be orphaned and a retried apply would mint
another one. The restore failure path now deletes the created key, and a
delete that also fails names the key hash in the error so an operator
can remove it manually.
* fix(terraform): make the route restore surgical and keep supplied keys
Two sharp edges on the create-time route restore:
- Re-sending the full create payload rewrote fields the config never
declared: /key/update is a merge patch, so the empty metadata and
model_rpm_limit/model_tpm_limit maps the restored struct carried would
clear server-applied values such as team-inherited rate limits. The
restore now sends only the routes plus the two fields /key/update
requires non-null (permissions, model_max_budget); every other stored
value is kept.
- /key/generate upserts a config-supplied key value, so a restore
failure on such a key must not delete it: it may be an existing
credential that predates this apply. The compensating delete now runs
only for proxy-minted keys, and the error names the hash either way.
* fix(terraform): echo stored permissions and budgets in route restore
The surgical restore body carried empty permissions and model_max_budget
objects, and /key/update writes fields that are present: a key created
with declared permissions or model budgets next to a presetting key_type
and allowed_routes lost them on the first apply. The restore now echoes
the values /key/generate just stored (falling back to the configured
values when the response omits them), so the only field the restore ever
changes is allowed_routes.
* test(terraform): pin echoed budgets in the route restore
Adds the nonempty model_max_budget case Greptile asked for (the restore
must echo the stored map, never clear it) and drops a comment that
restated its own line.
* test(terraform): assert the declared budget reaches key generation
The budget echo case fed the raw config a malformed JSON string (a
template leftover), so nothing verified the declared budget actually
reached /key/generate. The config now carries the valid JSON and the
generate payload is asserted to match it.
* chore(terraform): trim the restore test preface to the proxy facts
---------
Co-authored-by: Roman Soletskyi <roman@mistral.ai>
This commit is contained in:
parent
d8bc2b78e4
commit
cb76270bd9
5 changed files with 529 additions and 7 deletions
|
|
@ -79,6 +79,7 @@ Here's an example of creating an API key with various options:
|
|||
|
||||
```hcl
|
||||
resource "litellm_key" "example_key" {
|
||||
key_type = "llm_api"
|
||||
models = ["gpt-4", "claude-3.5-sonnet"]
|
||||
max_budget = 100.0
|
||||
user_id = "user123"
|
||||
|
|
@ -123,6 +124,7 @@ resource "litellm_key" "example_key" {
|
|||
|
||||
The <code>litellm_key</code> resource supports the following options:
|
||||
|
||||
- <code>key_type</code>: Choose the key's default route access
|
||||
- <code>models</code>: List of allowed models for this key
|
||||
- <code>max_budget</code>: Maximum budget for the key
|
||||
- <code>user_id</code> and <code>team_id</code>: Associate the key with a user and team
|
||||
|
|
|
|||
|
|
@ -121,6 +121,35 @@ func hoistKeyFieldsStoredInMetadata(info map[string]interface{}) {
|
|||
}
|
||||
}
|
||||
|
||||
// RestoreKeyRoutes re-applies a declared allowed_routes to a freshly
|
||||
// generated key, whose key_type preset replaced the declared list.
|
||||
// /key/update validates permissions and model_max_budget as non-null and
|
||||
// keeps the stored value for every field absent from the body, so this
|
||||
// carries exactly the routes plus those two required objects. Their values
|
||||
// must be the ones just stored for the key: sending empty objects would
|
||||
// clear configured or server-defaulted restrictions, and sending the whole
|
||||
// create payload would rewrite fields the config never declared (such as
|
||||
// team-inherited model rate limits).
|
||||
func (c *Client) RestoreKeyRoutes(keyID string, routes []string, permissions, modelMaxBudget map[string]interface{}) (*Key, error) {
|
||||
if permissions == nil {
|
||||
permissions = map[string]interface{}{}
|
||||
}
|
||||
if modelMaxBudget == nil {
|
||||
modelMaxBudget = map[string]interface{}{}
|
||||
}
|
||||
updateData := map[string]interface{}{
|
||||
"key": keyID,
|
||||
"allowed_routes": routes,
|
||||
"permissions": permissions,
|
||||
"model_max_budget": modelMaxBudget,
|
||||
}
|
||||
resp, err := c.sendRequest("POST", "/key/update", updateData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return c.parseKeyResponse(resp)
|
||||
}
|
||||
|
||||
func (c *Client) UpdateKey(key *Key) (*Key, error) {
|
||||
// Create a new map with only the fields that can be updated
|
||||
updateData := map[string]interface{}{
|
||||
|
|
@ -242,6 +271,10 @@ func (c *Client) parseKeyResponse(resp map[string]interface{}) (*Key, error) {
|
|||
if s, ok := v.(string); ok {
|
||||
createdKey.TokenID = s
|
||||
}
|
||||
case "key_type":
|
||||
if s, ok := v.(string); ok {
|
||||
createdKey.KeyType = s
|
||||
}
|
||||
case "models":
|
||||
if models, ok := v.([]interface{}); ok {
|
||||
createdKey.Models = make([]string, len(models))
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import (
|
|||
"github.com/hashicorp/go-cty/cty"
|
||||
"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
|
||||
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
|
||||
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
|
||||
)
|
||||
|
||||
func resourceKey() *schema.Resource {
|
||||
|
|
@ -34,6 +35,14 @@ func resourceKey() *schema.Resource {
|
|||
Type: schema.TypeString,
|
||||
Computed: true,
|
||||
},
|
||||
"key_type": {
|
||||
Type: schema.TypeString,
|
||||
Optional: true,
|
||||
Computed: true,
|
||||
ForceNew: true,
|
||||
ValidateFunc: validation.StringInSlice([]string{"llm_api", "management", "read_only", "default"}, false),
|
||||
Description: "Type of key that determines its default allowed routes. Changing it creates a new key",
|
||||
},
|
||||
"models": {
|
||||
Type: schema.TypeList,
|
||||
Optional: true,
|
||||
|
|
@ -161,9 +170,10 @@ func resourceKey() *schema.Resource {
|
|||
Elem: &schema.Schema{Type: schema.TypeString},
|
||||
},
|
||||
"allowed_routes": {
|
||||
Type: schema.TypeList,
|
||||
Optional: true,
|
||||
Elem: &schema.Schema{Type: schema.TypeString},
|
||||
Type: schema.TypeList,
|
||||
Optional: true,
|
||||
Elem: &schema.Schema{Type: schema.TypeString},
|
||||
DiffSuppressFunc: suppressUnconfiguredAllowedRoutes,
|
||||
},
|
||||
"allowed_passthrough_routes": {
|
||||
Type: schema.TypeList,
|
||||
|
|
@ -284,12 +294,49 @@ func resourceKeyCreate(ctx context.Context, d *schema.ResourceData, m interface{
|
|||
} else if v := d.Get("key").(string); v != "" {
|
||||
key.Key = v
|
||||
}
|
||||
proxyMintedKey := key.Key == ""
|
||||
|
||||
createdKey, err := c.CreateKey(key)
|
||||
if err != nil {
|
||||
return diag.FromErr(fmt.Errorf("error creating key: %s", err))
|
||||
}
|
||||
|
||||
// /key/generate replaces a declared allowed_routes with the key_type
|
||||
// preset, while /key/update stores the list verbatim. Re-assert the
|
||||
// declared routes right after create so the first apply already leaves
|
||||
// the key with the routes the config asks for (a replacement forced by
|
||||
// any ForceNew attribute would otherwise hand out the preset until a
|
||||
// second apply).
|
||||
if keyTypePresetsRoutes(key.KeyType) && len(key.AllowedRoutes) > 0 &&
|
||||
!slices.Equal(createdKey.AllowedRoutes, key.AllowedRoutes) {
|
||||
// Echo the values the generate just stored for the two fields
|
||||
// /key/update requires non-null; empty objects would clear
|
||||
// configured or server-defaulted restrictions.
|
||||
storedOrConfigured := func(stored, configured map[string]interface{}) map[string]interface{} {
|
||||
if stored != nil {
|
||||
return stored
|
||||
}
|
||||
return configured
|
||||
}
|
||||
if _, err := c.RestoreKeyRoutes(createdKey.TokenID, key.AllowedRoutes,
|
||||
storedOrConfigured(createdKey.Permissions, key.Permissions),
|
||||
storedOrConfigured(createdKey.ModelMaxBudget, key.ModelMaxBudget)); err != nil {
|
||||
// For a proxy-minted key nothing is in state yet, so returning
|
||||
// without cleanup would orphan an active key terraform cannot see
|
||||
// or delete, and a retried apply would mint another one. Delete it
|
||||
// so the retry starts clean; a config-supplied key may predate
|
||||
// this apply, so it is never deleted here. Name the hash either
|
||||
// way so an operator can finish by hand if cleanup fails.
|
||||
if proxyMintedKey {
|
||||
if delErr := c.DeleteKey(createdKey.TokenID); delErr != nil {
|
||||
return diag.FromErr(fmt.Errorf("error restoring allowed_routes over the key_type preset: %s; cleanup failed too, key %s must be deleted manually: %s", err, createdKey.TokenID, delErr))
|
||||
}
|
||||
return diag.FromErr(fmt.Errorf("error restoring allowed_routes over the key_type preset (created key deleted, retry the apply): %s", err))
|
||||
}
|
||||
return diag.FromErr(fmt.Errorf("error restoring allowed_routes over the key_type preset on config-supplied key %s; the key was left in place: %s", createdKey.TokenID, err))
|
||||
}
|
||||
}
|
||||
|
||||
d.SetId(createdKey.TokenID)
|
||||
// Set the write-only key value so it's available during this apply
|
||||
// but will not be persisted to state.
|
||||
|
|
@ -322,6 +369,9 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{
|
|||
|
||||
key := &Key{Key: d.Id()}
|
||||
mapResourceDataToKey(d, key)
|
||||
if allowedRoutesNotConfigured(d) {
|
||||
key.AllowedRoutes = nil
|
||||
}
|
||||
if !d.HasChange("duration") {
|
||||
key.Duration = ""
|
||||
}
|
||||
|
|
@ -373,6 +423,47 @@ func changedMap(d *schema.ResourceData, name string) map[string]interface{} {
|
|||
return d.Get(name).(map[string]interface{})
|
||||
}
|
||||
|
||||
// allowedRoutesNotConfigured reports whether the raw configuration leaves
|
||||
// allowed_routes unset. d.Get cannot answer this: it merges state into
|
||||
// unconfigured attributes, so once a refresh has materialized the server's
|
||||
// routes into state (or left a stale copy behind with -refresh=false), the
|
||||
// configured and unconfigured cases read identically.
|
||||
func allowedRoutesNotConfigured(d *schema.ResourceData) bool {
|
||||
raw, diags := d.GetRawConfigAt(cty.GetAttrPath("allowed_routes"))
|
||||
return !diags.HasError() && raw.IsNull()
|
||||
}
|
||||
|
||||
// keyTypePresetsRoutes reports whether the proxy derives allowed_routes from
|
||||
// this key_type at create time, overwriting whatever the request declared.
|
||||
// "default" (and an unset type) preset nothing.
|
||||
func keyTypePresetsRoutes(keyType string) bool {
|
||||
switch keyType {
|
||||
case "llm_api", "management", "read_only":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Reads copy the server's routes into state so drift on them stays visible,
|
||||
// and /key/update keeps the stored routes whenever allowed_routes is absent
|
||||
// from the payload. Suppressing the diff for a config that never declares the
|
||||
// attribute therefore matches the wire behavior: without suppression the plan
|
||||
// would show a perpetual removal diff against server-derived routes (the
|
||||
// presets a key_type implies, or routes granted directly on the proxy) that
|
||||
// no apply can ever clear. When the raw config is unavailable (helpers that
|
||||
// diff without one), only a whole-list removal shape is suppressed so a
|
||||
// config that shrinks the list still diffs.
|
||||
func suppressUnconfiguredAllowedRoutes(k, old, new string, d *schema.ResourceData) bool {
|
||||
if new != "" && new != "0" {
|
||||
return false
|
||||
}
|
||||
raw, diags := d.GetRawConfigAt(cty.GetAttrPath("allowed_routes"))
|
||||
if !diags.HasError() {
|
||||
return raw.IsNull()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
var errKeyGone = errors.New("no longer exists")
|
||||
|
||||
func plannedKeyMetadata(c *Client, d *schema.ResourceData) (map[string]interface{}, error) {
|
||||
|
|
@ -449,6 +540,7 @@ func resourceKeyDelete(ctx context.Context, d *schema.ResourceData, m interface{
|
|||
}
|
||||
|
||||
func mapResourceDataToKey(d *schema.ResourceData, key *Key) {
|
||||
key.KeyType = d.Get("key_type").(string)
|
||||
key.Models = expandStringList(d.Get("models").([]interface{}))
|
||||
if v, ok := d.GetOk("max_budget"); ok {
|
||||
val := v.(float64)
|
||||
|
|
@ -504,6 +596,9 @@ func mapKeyToResourceData(d *schema.ResourceData, key *Key) {
|
|||
|
||||
// Note: "key" is write-only and must not be set here (Read operations).
|
||||
// It is only set during Create so it is available during apply.
|
||||
if key.KeyType != "" {
|
||||
d.Set("key_type", key.KeyType)
|
||||
}
|
||||
|
||||
if len(key.Models) > 0 {
|
||||
d.Set("models", key.Models)
|
||||
|
|
@ -576,9 +671,7 @@ func mapKeyToResourceData(d *schema.ResourceData, key *Key) {
|
|||
if len(key.EnforcedParams) > 0 {
|
||||
d.Set("enforced_params", key.EnforcedParams)
|
||||
}
|
||||
if len(key.AllowedRoutes) > 0 {
|
||||
d.Set("allowed_routes", key.AllowedRoutes)
|
||||
}
|
||||
d.Set("allowed_routes", append([]string{}, key.AllowedRoutes...))
|
||||
if len(key.AllowedPassthroughRoutes) > 0 {
|
||||
d.Set("allowed_passthrough_routes", key.AllowedPassthroughRoutes)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,13 +3,16 @@ package litellm
|
|||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/hashicorp/go-cty/cty"
|
||||
"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
|
||||
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
|
||||
"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
|
||||
|
|
@ -22,6 +25,7 @@ func newKeyResourceData(t *testing.T, raw map[string]interface{}) *schema.Resour
|
|||
|
||||
func TestMapResourceDataToKeyNewFields(t *testing.T) {
|
||||
d := newKeyResourceData(t, map[string]interface{}{
|
||||
"key_type": "llm_api",
|
||||
"budget_id": "budget-1",
|
||||
"enforced_params": []interface{}{"user"},
|
||||
"allowed_routes": []interface{}{"/chat/completions"},
|
||||
|
|
@ -36,6 +40,9 @@ func TestMapResourceDataToKeyNewFields(t *testing.T) {
|
|||
key := &Key{}
|
||||
mapResourceDataToKey(d, key)
|
||||
|
||||
if key.KeyType != "llm_api" {
|
||||
t.Errorf("KeyType = %q, want llm_api", key.KeyType)
|
||||
}
|
||||
if key.BudgetID != "budget-1" {
|
||||
t.Errorf("BudgetID = %q, want budget-1", key.BudgetID)
|
||||
}
|
||||
|
|
@ -139,6 +146,7 @@ func TestParseKeyResponseNewFields(t *testing.T) {
|
|||
client := NewClient("http://localhost:4000", "test-key", true)
|
||||
resp := map[string]interface{}{
|
||||
"key": "sk-test",
|
||||
"key_type": "llm_api",
|
||||
"budget_id": "budget-1",
|
||||
"enforced_params": []interface{}{"user"},
|
||||
"allowed_routes": []interface{}{"/chat/completions"},
|
||||
|
|
@ -154,6 +162,9 @@ func TestParseKeyResponseNewFields(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("parseKeyResponse returned error: %v", err)
|
||||
}
|
||||
if key.KeyType != "llm_api" {
|
||||
t.Errorf("KeyType = %q, want llm_api", key.KeyType)
|
||||
}
|
||||
if key.BudgetID != "budget-1" || key.OrganizationID != "org-1" || key.ProjectID != "proj-1" {
|
||||
t.Errorf("string fields not parsed: %+v", key)
|
||||
}
|
||||
|
|
@ -183,7 +194,10 @@ func TestCreateKeySendsConfigSuppliedKey(t *testing.T) {
|
|||
defer srv.Close()
|
||||
|
||||
client := NewClient(srv.URL, "test-key", true)
|
||||
d := newKeyResourceData(t, map[string]interface{}{"key": "sk-custom"})
|
||||
d := newKeyResourceData(t, map[string]interface{}{
|
||||
"key": "sk-custom",
|
||||
"key_type": "llm_api",
|
||||
})
|
||||
|
||||
diags := resourceKeyCreate(context.Background(), d, client)
|
||||
if diags.HasError() {
|
||||
|
|
@ -192,11 +206,386 @@ func TestCreateKeySendsConfigSuppliedKey(t *testing.T) {
|
|||
if captured["key"] != "sk-custom" {
|
||||
t.Errorf("create payload key = %v, want sk-custom", captured["key"])
|
||||
}
|
||||
if captured["key_type"] != "llm_api" {
|
||||
t.Errorf("create payload key_type = %v, want llm_api", captured["key_type"])
|
||||
}
|
||||
if d.Id() != "hash-1" {
|
||||
t.Errorf("resource ID = %q, want hash-1", d.Id())
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyTypeRejectsUnknownValue(t *testing.T) {
|
||||
_, errs := resourceKey().Schema["key_type"].ValidateFunc("unrestricted", "key_type")
|
||||
if len(errs) == 0 {
|
||||
t.Fatal("key_type accepted an unknown value")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyTypeChangeForcesReplacement(t *testing.T) {
|
||||
res := resourceKey()
|
||||
priorData := newKeyResourceData(t, map[string]interface{}{"key_type": "default"})
|
||||
priorData.SetId("hash-1")
|
||||
config := terraform.NewResourceConfigRaw(map[string]interface{}{"key_type": "llm_api"})
|
||||
diff, err := res.Diff(context.Background(), priorData.State(), config, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("diff failed: %v", err)
|
||||
}
|
||||
if diff == nil || !diff.RequiresNew() {
|
||||
t.Fatalf("changing key_type must force replacement, diff = %+v", diff)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyTypePresetRoutesDoNotDrift(t *testing.T) {
|
||||
cases := map[string]struct {
|
||||
read *Key
|
||||
config map[string]interface{}
|
||||
}{
|
||||
"llm_api preset": {read: &Key{KeyType: "llm_api", AllowedRoutes: []string{"llm_api_routes"}}, config: map[string]interface{}{"key_type": "llm_api"}},
|
||||
"default no routes": {read: &Key{KeyType: "default"}, config: map[string]interface{}{}},
|
||||
}
|
||||
for name, tc := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
res := resourceKey()
|
||||
priorData := newKeyResourceData(t, map[string]interface{}{})
|
||||
priorData.SetId("hash-1")
|
||||
if err := priorData.Set("server_metadata", serverKeyMetadata(tc.read.Metadata)); err != nil {
|
||||
t.Fatalf("set server_metadata: %v", err)
|
||||
}
|
||||
mapKeyToResourceData(priorData, tc.read)
|
||||
diff, err := res.Diff(context.Background(), priorData.State(), terraform.NewResourceConfigRaw(tc.config), nil)
|
||||
if err != nil {
|
||||
t.Fatalf("diff failed: %v", err)
|
||||
}
|
||||
if diff != nil && !diff.Empty() {
|
||||
t.Fatalf("server-derived allowed_routes must not drift, diff = %+v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A config that omits allowed_routes must stay KNOWN at plan time: marking it
|
||||
// computed makes it "known after apply", which fails any plan that consumes
|
||||
// the attribute (e.g. for_each = toset(coalesce(..., []))) before the key
|
||||
// exists.
|
||||
func TestAllowedRoutesOmittedIsKnownAtPlan(t *testing.T) {
|
||||
res := resourceKey()
|
||||
prior := &terraform.InstanceState{}
|
||||
prior.RawConfig = keyRawConfig(t, nil)
|
||||
config := terraform.NewResourceConfigRaw(map[string]interface{}{"key_alias": "example"})
|
||||
diff, err := res.Diff(context.Background(), prior, config, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("diff failed: %v", err)
|
||||
}
|
||||
for k, attr := range diff.Attributes {
|
||||
if !strings.HasPrefix(k, "allowed_routes") {
|
||||
continue
|
||||
}
|
||||
if attr.NewComputed {
|
||||
t.Fatalf("omitted allowed_routes is computed (unknown) at plan time: %+v", attr)
|
||||
}
|
||||
t.Errorf("omitted allowed_routes produced a plan diff %q = %+v, want none", k, attr)
|
||||
}
|
||||
}
|
||||
|
||||
// Suppressing unconfigured routes must not swallow real config changes: a
|
||||
// config that shrinks the declared list still has to diff.
|
||||
func TestAllowedRoutesConfiguredShrinkStillDiffs(t *testing.T) {
|
||||
res := resourceKey()
|
||||
priorData := newKeyResourceData(t, map[string]interface{}{
|
||||
"allowed_routes": []interface{}{"/a", "/b", "/c"},
|
||||
})
|
||||
priorData.SetId("hash-1")
|
||||
prior := priorData.State()
|
||||
prior.RawConfig = keyRawConfig(t, []string{"/a", "/b"})
|
||||
config := terraform.NewResourceConfigRaw(map[string]interface{}{
|
||||
"allowed_routes": []interface{}{"/a", "/b"},
|
||||
})
|
||||
diff, err := res.Diff(context.Background(), prior, config, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("diff failed: %v", err)
|
||||
}
|
||||
if diff == nil || diff.Attributes["allowed_routes.#"] == nil {
|
||||
t.Fatalf("config shrinking allowed_routes must still diff, diff = %+v", diff)
|
||||
}
|
||||
}
|
||||
|
||||
// The diff for server-derived routes in state is suppressed via the raw
|
||||
// config, the same signal real terraform runs carry.
|
||||
func TestAllowedRoutesUnconfiguredDoesNotDriftWithRawConfig(t *testing.T) {
|
||||
res := resourceKey()
|
||||
priorData := newKeyResourceData(t, map[string]interface{}{
|
||||
"key_alias": "old",
|
||||
"allowed_routes": []interface{}{"/v1/models"},
|
||||
})
|
||||
priorData.SetId("hash-1")
|
||||
prior := priorData.State()
|
||||
prior.RawConfig = keyRawConfig(t, nil)
|
||||
config := terraform.NewResourceConfigRaw(map[string]interface{}{"key_alias": "old"})
|
||||
diff, err := res.Diff(context.Background(), prior, config, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("diff failed: %v", err)
|
||||
}
|
||||
for k := range diff.Attributes {
|
||||
if strings.HasPrefix(k, "allowed_routes") {
|
||||
t.Fatalf("unconfigured allowed_routes drifted at plan: %q = %+v", k, diff.Attributes[k])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func keyRawConfig(t *testing.T, routes []string) cty.Value {
|
||||
t.Helper()
|
||||
return cty.ObjectVal(map[string]cty.Value{
|
||||
"key_alias": cty.StringVal("example"),
|
||||
"allowed_routes": keyRawConfigRoutes(t, routes),
|
||||
})
|
||||
}
|
||||
|
||||
// An update must never POST allowed_routes the config does not declare: the
|
||||
// value d.Get returns is whatever the last refresh stored (stale with
|
||||
// -refresh=false), and /key/update would overwrite externally managed routes
|
||||
// with it. The alias-only rename below must leave the field out.
|
||||
func TestResourceKeyUpdateOmitsUnconfiguredAllowedRoutes(t *testing.T) {
|
||||
var captured map[string]interface{}
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
json.Unmarshal(body, &captured)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if r.URL.Path == "/key/update" {
|
||||
w.Write([]byte(`{"key": "hash-1"}`))
|
||||
return
|
||||
}
|
||||
w.Write([]byte(`{"key":"hash-1","info":{"key_alias":"renamed","allowed_routes":["/v1/models"]}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res := resourceKey()
|
||||
priorData := newKeyResourceData(t, map[string]interface{}{
|
||||
"key_alias": "old",
|
||||
"allowed_routes": []interface{}{"/chat/completions"},
|
||||
})
|
||||
priorData.SetId("hash-1")
|
||||
prior := priorData.State()
|
||||
config := terraform.NewResourceConfigRaw(map[string]interface{}{"key_alias": "renamed"})
|
||||
diff, err := res.Diff(context.Background(), prior, config, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("diff failed: %v", err)
|
||||
}
|
||||
diff.RawConfig = keyRawConfig(t, nil)
|
||||
|
||||
_, diags := res.Apply(context.Background(), prior, diff, NewClient(srv.URL, "test-key", true))
|
||||
if diags.HasError() {
|
||||
t.Fatalf("apply failed: %+v", diags)
|
||||
}
|
||||
if _, present := captured["allowed_routes"]; present {
|
||||
t.Fatalf("update POSTed allowed_routes %v although the config does not declare it", captured["allowed_routes"])
|
||||
}
|
||||
if captured["key_alias"] != "renamed" {
|
||||
t.Errorf("update payload key_alias = %v, want renamed", captured["key_alias"])
|
||||
}
|
||||
}
|
||||
|
||||
// Declaring allowed_routes keeps owning them: an update still re-asserts the
|
||||
// configured routes, matching the pre-key_type behavior.
|
||||
func TestResourceKeyUpdateSendsConfiguredAllowedRoutes(t *testing.T) {
|
||||
var captured map[string]interface{}
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
json.Unmarshal(body, &captured)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if r.URL.Path == "/key/update" {
|
||||
w.Write([]byte(`{"key": "hash-1"}`))
|
||||
return
|
||||
}
|
||||
w.Write([]byte(`{"key":"hash-1","info":{"key_alias":"renamed","allowed_routes":["/v1/models"]}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res := resourceKey()
|
||||
priorData := newKeyResourceData(t, map[string]interface{}{
|
||||
"key_alias": "old",
|
||||
"allowed_routes": []interface{}{"/v1/models"},
|
||||
})
|
||||
priorData.SetId("hash-1")
|
||||
prior := priorData.State()
|
||||
config := terraform.NewResourceConfigRaw(map[string]interface{}{
|
||||
"key_alias": "renamed",
|
||||
"allowed_routes": []interface{}{"/v1/models"},
|
||||
})
|
||||
diff, err := res.Diff(context.Background(), prior, config, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("diff failed: %v", err)
|
||||
}
|
||||
diff.RawConfig = keyRawConfig(t, []string{"/v1/models"})
|
||||
|
||||
_, diags := res.Apply(context.Background(), prior, diff, NewClient(srv.URL, "test-key", true))
|
||||
if diags.HasError() {
|
||||
t.Fatalf("apply failed: %+v", diags)
|
||||
}
|
||||
routes, ok := captured["allowed_routes"].([]interface{})
|
||||
if !ok || len(routes) != 1 || routes[0] != "/v1/models" {
|
||||
t.Fatalf("update payload allowed_routes = %v, want [/v1/models]", captured["allowed_routes"])
|
||||
}
|
||||
}
|
||||
|
||||
// /key/generate replaces declared routes with the key_type preset while
|
||||
// /key/update stores them verbatim, so create must re-assert the declared
|
||||
// list; the restore body may only touch allowed_routes plus the two fields
|
||||
// /key/update requires non-null.
|
||||
func TestCreateKeyRestoresDeclaredRoutesOverPreset(t *testing.T) {
|
||||
cases := map[string]struct {
|
||||
generateRoutes []interface{}
|
||||
suppliedKey string
|
||||
permissions map[string]interface{}
|
||||
generateStores bool
|
||||
generateBudget map[string]interface{}
|
||||
restoreFails bool
|
||||
wantUpdate bool
|
||||
}{
|
||||
"preset overwrote declared": {generateRoutes: []interface{}{"llm_api_routes"}, wantUpdate: true},
|
||||
"declared permissions are echoed": {generateRoutes: []interface{}{"llm_api_routes"}, permissions: map[string]interface{}{"get_server_info": "true"}, generateStores: true, wantUpdate: true},
|
||||
"declared budgets are echoed": {generateRoutes: []interface{}{"llm_api_routes"}, generateBudget: map[string]interface{}{"x": true}, wantUpdate: true},
|
||||
"generate honored declared": {generateRoutes: []interface{}{"/v1/models"}, wantUpdate: false},
|
||||
"restore update rejected": {generateRoutes: []interface{}{"llm_api_routes"}, restoreFails: true, wantUpdate: true},
|
||||
"rejected restore keeps supplied": {generateRoutes: []interface{}{"llm_api_routes"}, suppliedKey: "sk-custom", restoreFails: true, wantUpdate: true},
|
||||
"supplied key still gets restored": {generateRoutes: []interface{}{"llm_api_routes"}, suppliedKey: "sk-custom", wantUpdate: true},
|
||||
}
|
||||
for name, tc := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
var generateBody, updateBody, deleteBody map[string]interface{}
|
||||
updateCalled, deleteCalled := false, false
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
switch r.URL.Path {
|
||||
case "/key/generate":
|
||||
json.Unmarshal(body, &generateBody)
|
||||
extra := ""
|
||||
if tc.generateStores {
|
||||
extra = `, "permissions": {"get_server_info": true}`
|
||||
}
|
||||
if tc.generateBudget != nil {
|
||||
extra += `, "model_max_budget": {"gpt-4o-mini": {"budget_limit": 5, "time_period": "30d"}}`
|
||||
}
|
||||
w.Write([]byte(`{"key": "sk-new", "token_id": "hash-1", "allowed_routes": ["` + tc.generateRoutes[0].(string) + `"]` + extra + `}`))
|
||||
case "/key/update":
|
||||
updateCalled = true
|
||||
json.Unmarshal(body, &updateBody)
|
||||
if tc.restoreFails {
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
w.Write([]byte(`{"error":{"message":"rejected"}}`))
|
||||
return
|
||||
}
|
||||
w.Write([]byte(`{"key": "hash-1"}`))
|
||||
case "/key/delete":
|
||||
deleteCalled = true
|
||||
json.Unmarshal(body, &deleteBody)
|
||||
w.Write([]byte(`{}`))
|
||||
default:
|
||||
w.Write([]byte(`{"key":"hash-1","info":{"key_alias":"typed","key_type":"llm_api","allowed_routes":["/v1/models"]}}`))
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
raw := map[string]interface{}{
|
||||
"key_alias": "typed",
|
||||
"key_type": "llm_api",
|
||||
"allowed_routes": []interface{}{"/v1/models"},
|
||||
}
|
||||
if tc.permissions != nil {
|
||||
raw["permissions"] = tc.permissions
|
||||
}
|
||||
if tc.generateBudget != nil {
|
||||
raw["model_max_budget"] = `{"gpt-4o-mini": {"budget_limit": 5, "time_period": "30d"}}`
|
||||
}
|
||||
if tc.suppliedKey != "" {
|
||||
raw["key"] = tc.suppliedKey
|
||||
}
|
||||
d := newKeyResourceData(t, raw)
|
||||
diags := resourceKeyCreate(context.Background(), d, NewClient(srv.URL, "test-key", true))
|
||||
if tc.restoreFails {
|
||||
if !diags.HasError() {
|
||||
t.Fatal("create succeeded although the restore update was rejected")
|
||||
}
|
||||
if tc.suppliedKey != "" {
|
||||
if deleteCalled {
|
||||
t.Fatal("failed restore deleted a config-supplied key, which /key/generate may have upserted onto an existing credential")
|
||||
}
|
||||
return
|
||||
}
|
||||
if !deleteCalled {
|
||||
t.Fatal("failed restore did not delete the created key; a retried apply would orphan it")
|
||||
}
|
||||
if keys, _ := deleteBody["keys"].([]interface{}); len(keys) != 1 || keys[0] != "hash-1" {
|
||||
t.Errorf("delete payload keys = %v, want [hash-1]", deleteBody["keys"])
|
||||
}
|
||||
if d.Id() != "" {
|
||||
t.Errorf("state recorded id %q for a key the provider deleted", d.Id())
|
||||
}
|
||||
return
|
||||
}
|
||||
if diags.HasError() {
|
||||
t.Fatalf("create failed: %+v", diags)
|
||||
}
|
||||
if generateBody["key_type"] != "llm_api" {
|
||||
t.Errorf("generate payload key_type = %v, want llm_api", generateBody["key_type"])
|
||||
}
|
||||
if tc.generateBudget != nil {
|
||||
want := map[string]interface{}{"gpt-4o-mini": map[string]interface{}{"budget_limit": float64(5), "time_period": "30d"}}
|
||||
if fmt.Sprint(generateBody["model_max_budget"]) != fmt.Sprint(want) {
|
||||
t.Errorf("generate payload model_max_budget = %v, want the declared %v", generateBody["model_max_budget"], want)
|
||||
}
|
||||
}
|
||||
if tc.wantUpdate {
|
||||
if !updateCalled {
|
||||
t.Fatal("create did not re-assert declared allowed_routes after the preset overwrote them")
|
||||
}
|
||||
wantPerms := map[string]interface{}{}
|
||||
if tc.generateStores {
|
||||
wantPerms = map[string]interface{}{"get_server_info": true}
|
||||
}
|
||||
wantBudget := map[string]interface{}{}
|
||||
if tc.generateBudget != nil {
|
||||
wantBudget = map[string]interface{}{"gpt-4o-mini": map[string]interface{}{"budget_limit": float64(5), "time_period": "30d"}}
|
||||
}
|
||||
wantBody := map[string]interface{}{
|
||||
"key": "hash-1",
|
||||
"allowed_routes": []interface{}{"/v1/models"},
|
||||
"permissions": wantPerms,
|
||||
"model_max_budget": wantBudget,
|
||||
}
|
||||
if len(updateBody) != len(wantBody) {
|
||||
t.Fatalf("restore payload = %v, want exactly %v (anything else rewrites fields the config did not declare)", updateBody, wantBody)
|
||||
}
|
||||
for k, v := range wantBody {
|
||||
if fmt.Sprint(updateBody[k]) != fmt.Sprint(v) {
|
||||
t.Errorf("restore payload %s = %v (%T), want %v (%T)", k, updateBody[k], updateBody[k], v, v)
|
||||
}
|
||||
}
|
||||
} else if updateCalled {
|
||||
t.Fatal("create re-asserted routes although the generate already stored the declared list")
|
||||
}
|
||||
if deleteCalled {
|
||||
t.Fatal("create deleted a key although nothing failed")
|
||||
}
|
||||
if got := d.Get("allowed_routes").([]interface{}); len(got) != 1 || got[0] != "/v1/models" {
|
||||
t.Errorf("state allowed_routes = %v, want [/v1/models]", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func keyRawConfigRoutes(t *testing.T, routes []string) cty.Value {
|
||||
t.Helper()
|
||||
if routes == nil {
|
||||
return cty.NullVal(cty.List(cty.String))
|
||||
}
|
||||
vals := make([]cty.Value, len(routes))
|
||||
for i, r := range routes {
|
||||
vals[i] = cty.StringVal(r)
|
||||
}
|
||||
return cty.ListVal(vals)
|
||||
}
|
||||
|
||||
// The proxy validates each model_max_budget entry as a BudgetConfig object and
|
||||
// 500s on a bare number, so the JSON string must reach /key/generate as nested
|
||||
// objects and the proxy's response must map back to equivalent JSON in state.
|
||||
|
|
@ -370,6 +759,7 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) {
|
|||
w.Write([]byte(`{
|
||||
"key": "hash-1",
|
||||
"info": {
|
||||
"key_type": "llm_api",
|
||||
"key_alias": "envelope-alias",
|
||||
"models": ["gpt-4o-mini"],
|
||||
"budget_id": "budget-1",
|
||||
|
|
@ -388,6 +778,9 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) {
|
|||
if key.KeyAlias != "envelope-alias" {
|
||||
t.Errorf("KeyAlias = %q, want envelope-alias (info envelope not unwrapped)", key.KeyAlias)
|
||||
}
|
||||
if key.KeyType != "llm_api" {
|
||||
t.Errorf("KeyType = %q, want llm_api", key.KeyType)
|
||||
}
|
||||
if key.BudgetID != "budget-1" || key.TeamID != "team-1" {
|
||||
t.Errorf("nested fields not parsed: %+v", key)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -131,6 +131,7 @@ type ModelInfo struct {
|
|||
type Key struct {
|
||||
Key string `json:"key,omitempty"`
|
||||
TokenID string `json:"token_id,omitempty"`
|
||||
KeyType string `json:"key_type,omitempty"`
|
||||
Models []string `json:"models"`
|
||||
Spend float64 `json:"spend,omitempty"`
|
||||
MaxBudget *float64 `json:"max_budget,omitempty"`
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue