fix(terraform): preserve undeclared server-side key metadata on update (#40514)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

terraform/provider/litellm/client.go

@@ -80,16 +80,24 @@ func (c *Client) UpdateKey(key *Key) (*Key, error) {
 	updateData := map[string]interface{}{
 		"key":              key.Key,
 		"team_id":          key.TeamID,
-		"metadata":         key.Metadata,
 		"key_alias":        key.KeyAlias,
 		"aliases":          key.Aliases,
 		"permissions":      key.Permissions,
 		"model_max_budget": key.ModelMaxBudget,
-		"model_rpm_limit":  key.ModelRPMLimit,
-		"model_tpm_limit":  key.ModelTPMLimit,
 		"blocked":          key.Blocked,
 	}
 
+	// The proxy keeps the stored metadata only when the field is absent, so nil means omit.
+	if key.Metadata != nil {
+		updateData["metadata"] = key.Metadata
+	}
+	if key.ModelRPMLimit != nil {
+		updateData["model_rpm_limit"] = key.ModelRPMLimit
+	}
+	if key.ModelTPMLimit != nil {
+		updateData["model_tpm_limit"] = key.ModelTPMLimit
+	}
+
 	// The proxy rejects an empty-string budget_duration with a 400, so only
 	// send it when set.
 	if key.BudgetDuration != "" {

terraform/provider/litellm/resource_key.go

@@ -223,6 +223,7 @@ func resourceKeyRead(ctx context.Context, d *schema.ResourceData, m interface{})
 		return nil
 	}
 
+	key.Metadata = declaredKeyMetadata(key.Metadata, d.Get("metadata").(map[string]interface{}))
 	mapKeyToResourceData(d, key)
 	return nil
 }
@@ -232,15 +233,71 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{
 
 	key := &Key{Key: d.Id()}
 	mapResourceDataToKey(d, key)
+	key.ModelRPMLimit = changedMap(d, "model_rpm_limit")
+	key.ModelTPMLimit = changedMap(d, "model_tpm_limit")
 
-	_, err := c.UpdateKey(key)
+	metadata, err := plannedKeyMetadata(c, d)
 	if err != nil {
 		return diag.FromErr(fmt.Errorf("error updating key: %s", err))
 	}
+	key.Metadata = metadata
+
+	if _, err := c.UpdateKey(key); err != nil {
+		return diag.FromErr(fmt.Errorf("error updating key: %s", err))
+	}
 
 	return resourceKeyRead(ctx, d, m)
 }
 
+func changedMap(d *schema.ResourceData, name string) map[string]interface{} {
+	if !d.HasChange(name) {
+		return nil
+	}
+	return d.Get(name).(map[string]interface{})
+}
+
+func plannedKeyMetadata(c *Client, d *schema.ResourceData) (map[string]interface{}, error) {
+	if !d.HasChange("metadata") {
+		return nil, nil
+	}
+	current, err := c.GetKey(d.Id())
+	if err != nil {
+		return nil, err
+	}
+	if current == nil {
+		return nil, fmt.Errorf("key %s no longer exists", d.Id())
+	}
+	oldDeclared, newDeclared := d.GetChange("metadata")
+	return mergeKeyMetadata(current.Metadata, oldDeclared.(map[string]interface{}), newDeclared.(map[string]interface{})), nil
+}
+
+func declaredKeyMetadata(server, declared map[string]interface{}) map[string]interface{} {
+	if server == nil {
+		return nil
+	}
+	result := make(map[string]interface{}, len(declared))
+	for k := range declared {
+		if v, ok := server[k]; ok {
+			result[k] = v
+		}
+	}
+	return result
+}
+
+func mergeKeyMetadata(server, oldDeclared, newDeclared map[string]interface{}) map[string]interface{} {
+	result := make(map[string]interface{}, len(server)+len(newDeclared))
+	for k, v := range server {
+		result[k] = v
+	}
+	for k := range oldDeclared {
+		delete(result, k)
+	}
+	for k, v := range newDeclared {
+		result[k] = v
+	}
+	return result
+}
+
 func resourceKeyDelete(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics {
 	c := m.(*Client)
 

terraform/provider/litellm/resource_key_test.go

@@ -6,9 +6,11 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
 	"testing"
 
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
 )
 
 func newKeyResourceData(t *testing.T, raw map[string]interface{}) *schema.ResourceData {
@@ -254,3 +256,134 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) {
 		t.Errorf("RPMLimit not parsed: %+v", key.RPMLimit)
 	}
 }
+
+// fakeKeyProxy serves /key/info from stored metadata and applies /key/update
+// the way the proxy does: an absent "metadata" keeps the stored map, a
+// present one replaces it wholesale.
+type fakeKeyProxy struct {
+	metadata map[string]interface{}
+	updates  []map[string]interface{}
+}
+
+func (p *fakeKeyProxy) handler() http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		switch r.URL.Path {
+		case "/key/info":
+			json.NewEncoder(w).Encode(map[string]interface{}{
+				"key":  "hash-1",
+				"info": map[string]interface{}{"key_alias": "alias-1", "models": []string{"gpt-4o-mini"}, "metadata": p.metadata},
+			})
+		case "/key/update":
+			var body map[string]interface{}
+			json.NewDecoder(r.Body).Decode(&body)
+			p.updates = append(p.updates, body)
+			if m, ok := body["metadata"].(map[string]interface{}); ok {
+				p.metadata = m
+			}
+			json.NewEncoder(w).Encode(map[string]interface{}{"key": "hash-1", "metadata": p.metadata})
+		default:
+			http.NotFound(w, r)
+		}
+	}
+}
+
+func applyKeyUpdate(t *testing.T, client *Client, stateAttrs map[string]string, config map[string]interface{}) *terraform.InstanceState {
+	t.Helper()
+	r := resourceKey()
+	state := &terraform.InstanceState{ID: "hash-1", Attributes: stateAttrs}
+	diff, err := r.Diff(context.Background(), state, terraform.NewResourceConfigRaw(config), client)
+	if err != nil {
+		t.Fatalf("Diff returned error: %v", err)
+	}
+	if diff == nil {
+		t.Fatalf("expected a non-empty diff between %v and %v", stateAttrs, config)
+	}
+	newState, diags := r.Apply(context.Background(), state, diff, client)
+	if diags.HasError() {
+		t.Fatalf("Apply returned error: %v", diags)
+	}
+	return newState
+}
+
+func TestKeyUpdateWithoutMetadataChangePreservesServerMetadata(t *testing.T) {
+	proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "server_side": "x", "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": float64(5)}}}
+	srv := httptest.NewServer(proxy.handler())
+	defer srv.Close()
+	client := NewClient(srv.URL, "test-key", true)
+
+	newState := applyKeyUpdate(t, client,
+		map[string]string{"key_alias": "alias-1", "max_budget": "10", "metadata.%": "1", "metadata.a": "1"},
+		map[string]interface{}{"key_alias": "alias-1", "max_budget": 20, "metadata": map[string]interface{}{"a": "1"}},
+	)
+
+	if len(proxy.updates) != 1 {
+		t.Fatalf("expected one /key/update call, got %d", len(proxy.updates))
+	}
+	for _, field := range []string{"metadata", "model_rpm_limit", "model_tpm_limit"} {
+		if _, present := proxy.updates[0][field]; present {
+			t.Errorf("unchanged %q was sent on /key/update: %v", field, proxy.updates[0][field])
+		}
+	}
+	if proxy.metadata["server_side"] != "x" {
+		t.Errorf("server-side metadata lost: %v", proxy.metadata)
+	}
+	if got := newState.Attributes["metadata.%"]; got != "1" {
+		t.Errorf("state metadata should hold only the declared entry, got %v", newState.Attributes)
+	}
+	if got := newState.Attributes["metadata.a"]; got != "1" {
+		t.Errorf("metadata.a = %q, want 1", got)
+	}
+}
+
+func TestKeyUpdateWithMetadataChangeMergesOverServerMetadata(t *testing.T) {
+	proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "b": "2", "server_side": "x"}}
+	srv := httptest.NewServer(proxy.handler())
+	defer srv.Close()
+	client := NewClient(srv.URL, "test-key", true)
+
+	applyKeyUpdate(t, client,
+		map[string]string{"key_alias": "alias-1", "metadata.%": "2", "metadata.a": "1", "metadata.b": "2"},
+		map[string]interface{}{"key_alias": "alias-1", "metadata": map[string]interface{}{"a": "2", "c": "3"}},
+	)
+
+	want := map[string]interface{}{"a": "2", "c": "3", "server_side": "x"}
+	if !reflect.DeepEqual(proxy.metadata, want) {
+		t.Errorf("metadata after update = %v, want %v", proxy.metadata, want)
+	}
+}
+
+func TestKeyUpdateSendsChangedModelLimits(t *testing.T) {
+	proxy := &fakeKeyProxy{metadata: map[string]interface{}{}}
+	srv := httptest.NewServer(proxy.handler())
+	defer srv.Close()
+	client := NewClient(srv.URL, "test-key", true)
+
+	applyKeyUpdate(t, client,
+		map[string]string{"key_alias": "alias-1", "model_rpm_limit.%": "1", "model_rpm_limit.gpt-4o-mini": "5"},
+		map[string]interface{}{"key_alias": "alias-1", "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 7}},
+	)
+
+	got, ok := proxy.updates[0]["model_rpm_limit"].(map[string]interface{})
+	if !ok || got["gpt-4o-mini"] != float64(7) {
+		t.Errorf("changed model_rpm_limit not sent: %v", proxy.updates[0])
+	}
+}
+
+func TestKeyReadKeepsOnlyDeclaredMetadata(t *testing.T) {
+	proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "server_side": "x"}}
+	srv := httptest.NewServer(proxy.handler())
+	defer srv.Close()
+	client := NewClient(srv.URL, "test-key", true)
+
+	d := newKeyResourceData(t, map[string]interface{}{"metadata": map[string]interface{}{"a": "1"}})
+	d.SetId("hash-1")
+	if diags := resourceKeyRead(context.Background(), d, client); diags.HasError() {
+		t.Fatalf("Read returned error: %v", diags)
+	}
+
+	want := map[string]interface{}{"a": "1"}
+	if got := d.Get("metadata"); !reflect.DeepEqual(got, want) {
+		t.Errorf("metadata in state = %v, want %v", got, want)
+	}
+}