From 44903b3bdefa78ae57c461aab343d9d1573d9b18 Mon Sep 17 00:00:00 2001
From: Erik Bogado <erikbogado@gmail.com>
Date: Mon, 7 Sep 2026 03:36:23 -0300
Subject: [PATCH] feat(terraform): expose generated key output

---
 terraform/provider/docs/resources/key.md      |  4 +++
 terraform/provider/litellm/resource_key.go    | 13 ++++---
 .../provider/litellm/resource_key_test.go     | 36 +++++++++++++++++++
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/terraform/provider/docs/resources/key.md b/terraform/provider/docs/resources/key.md
index 5094b77cbec..dfbb551e91f 100644
--- a/terraform/provider/docs/resources/key.md
+++ b/terraform/provider/docs/resources/key.md
@@ -132,3 +132,7 @@ $ terraform import litellm_key.example 12345
 ```
 
 This allows you to import existing keys into your Terraform state, enabling management of keys that were created outside of Terraform.
+
+### Generated key output
+
+`generated_key` is a sensitive computed string containing the key returned at creation. It remains in Terraform state across refreshes so dependent resources can use it. Protect state as secret material. Importing an existing token hash cannot recover the original key; `generated_key` is null after import. The optional write-only `key` argument still supplies a custom key at creation
diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go
index 0d8674f2d4c..fbc06eab511 100644
--- a/terraform/provider/litellm/resource_key.go
+++ b/terraform/provider/litellm/resource_key.go
@@ -25,6 +25,11 @@ func resourceKey() *schema.Resource {
 				WriteOnly: true,
 				Sensitive: true,
 			},
+			"generated_key": {
+				Type:      schema.TypeString,
+				Computed:  true,
+				Sensitive: true,
+			},
 			"token_id": {
 				Type:     schema.TypeString,
 				Computed: true,
@@ -192,7 +197,7 @@ func resourceKeyCreate(ctx context.Context, d *schema.ResourceData, m interface{
 	// A config-supplied key value becomes the key itself; when absent the
 	// proxy generates one. Write-only attributes are invisible to d.Get in
 	// real Terraform runs, so read the raw config first.
-	if raw, err := d.GetRawConfigAt(cty.GetAttrPath("key")); err == nil && !raw.IsNull() && raw.Type() == cty.String && raw.AsString() != "" {
+	if raw, err := d.GetRawConfigAt(cty.GetAttrPath("key")); err == nil && raw.IsKnown() && !raw.IsNull() && raw.Type() == cty.String && raw.AsString() != "" {
 		key.Key = raw.AsString()
 	} else if v := d.Get("key").(string); v != "" {
 		key.Key = v
@@ -204,9 +209,9 @@ func resourceKeyCreate(ctx context.Context, d *schema.ResourceData, m interface{
 	}
 
 	d.SetId(createdKey.TokenID)
-	// Set the write-only key value so it's available during this apply
-	// but will not be persisted to state.
-	d.Set("key", createdKey.Key)
+	if err := d.Set("generated_key", createdKey.Key); err != nil {
+		return diag.FromErr(fmt.Errorf("error storing generated key: %w", err))
+	}
 	return resourceKeyRead(ctx, d, m)
 }
 
diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go
index 91f0061a9ef..8ce11e63054 100644
--- a/terraform/provider/litellm/resource_key_test.go
+++ b/terraform/provider/litellm/resource_key_test.go
@@ -3,6 +3,8 @@ package litellm
 import (
 	"context"
 	"encoding/json"
+	"github.com/hashicorp/go-cty/cty"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -254,3 +256,37 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) {
 		t.Errorf("RPMLimit not parsed: %+v", key.RPMLimit)
 	}
 }
+
+func TestKeyCreateUsesWriteOnlyRawConfig(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		if r.URL.Path == "/key/generate" {
+			var payload map[string]interface{}
+			if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
+				t.Error(err)
+			}
+			if payload["key"] != "sk-custom-test" {
+				t.Errorf("custom key dropped: %v", payload["key"])
+			}
+			w.Write([]byte(`{"key":"sk-custom-test","token_id":"hash-1"}`))
+			return
+		}
+		w.Write([]byte(`{"key":"hash-1","info":{"key_alias":"remote"}}`))
+	}))
+	defer srv.Close()
+	d, err := schema.InternalMap(resourceKey().Schema).Data(nil, &terraform.InstanceDiff{
+		RawConfig: cty.ObjectVal(map[string]cty.Value{"key": cty.StringVal("sk-custom-test")}),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if d.Get("key") != "" {
+		t.Fatal("test must exercise a write-only key absent from d.Get")
+	}
+	if diags := resourceKeyCreate(context.Background(), d, NewClient(srv.URL, "test-key", false)); diags.HasError() {
+		t.Fatal(diags)
+	}
+	if d.Id() != "hash-1" || d.Get("generated_key") != "sk-custom-test" || d.Get("key_alias") != "remote" {
+		t.Fatal("create did not preserve generated key and refresh remote state")
+	}
+}