From ac59d284b184b26a32925db1f1437f8be7c54aeb Mon Sep 17 00:00:00 2001
From: Virgil <virgil@lethean.io>
Date: Wed, 1 Apr 2026 16:18:10 +0000
Subject: [PATCH] feat(api): document rate limit headers on all responses

Co-Authored-By: Virgil <virgil@lethean.io>
---
 openapi.go      | 10 +++++-----
 openapi_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/openapi.go b/openapi.go
index 257a392..c455d8a 100644
--- a/openapi.go
+++ b/openapi.go
@@ -193,7 +193,7 @@ func operationResponses(dataSchema map[string]any) map[string]any {
 					"schema": envelopeSchema(nil),
 				},
 			},
-			"headers": standardResponseHeaders(),
+			"headers": mergeHeaders(standardResponseHeaders(), rateLimitSuccessHeaders()),
 		},
 		"401": map[string]any{
 			"description": "Unauthorised",
@@ -202,7 +202,7 @@ func operationResponses(dataSchema map[string]any) map[string]any {
 					"schema": envelopeSchema(nil),
 				},
 			},
-			"headers": standardResponseHeaders(),
+			"headers": mergeHeaders(standardResponseHeaders(), rateLimitSuccessHeaders()),
 		},
 		"403": map[string]any{
 			"description": "Forbidden",
@@ -211,7 +211,7 @@ func operationResponses(dataSchema map[string]any) map[string]any {
 					"schema": envelopeSchema(nil),
 				},
 			},
-			"headers": standardResponseHeaders(),
+			"headers": mergeHeaders(standardResponseHeaders(), rateLimitSuccessHeaders()),
 		},
 		"429": map[string]any{
 			"description": "Too many requests",
@@ -229,7 +229,7 @@ func operationResponses(dataSchema map[string]any) map[string]any {
 					"schema": envelopeSchema(nil),
 				},
 			},
-			"headers": standardResponseHeaders(),
+			"headers": mergeHeaders(standardResponseHeaders(), rateLimitSuccessHeaders()),
 		},
 	}
 }
@@ -263,7 +263,7 @@ func healthResponses() map[string]any {
 					"schema": envelopeSchema(nil),
 				},
 			},
-			"headers": standardResponseHeaders(),
+			"headers": mergeHeaders(standardResponseHeaders(), rateLimitSuccessHeaders()),
 		},
 	}
 }
diff --git a/openapi_test.go b/openapi_test.go
index 433ef19..f1e823e 100644
--- a/openapi_test.go
+++ b/openapi_test.go
@@ -85,6 +85,20 @@ func TestSpecBuilder_Good_EmptyGroups(t *testing.T) {
 	if _, ok := headers["X-RateLimit-Reset"]; !ok {
 		t.Fatal("expected X-RateLimit-Reset header on /health 429 response")
 	}
+	health504 := healthResponses["504"].(map[string]any)
+	health504Headers := health504["headers"].(map[string]any)
+	if _, ok := health504Headers["X-Request-ID"]; !ok {
+		t.Fatal("expected X-Request-ID header on /health 504 response")
+	}
+	if _, ok := health504Headers["X-RateLimit-Limit"]; !ok {
+		t.Fatal("expected X-RateLimit-Limit header on /health 504 response")
+	}
+	if _, ok := health504Headers["X-RateLimit-Remaining"]; !ok {
+		t.Fatal("expected X-RateLimit-Remaining header on /health 504 response")
+	}
+	if _, ok := health504Headers["X-RateLimit-Reset"]; !ok {
+		t.Fatal("expected X-RateLimit-Reset header on /health 504 response")
+	}
 
 	// Verify system tag exists.
 	tags := spec["tags"].([]any)
@@ -278,6 +292,22 @@ func TestSpecBuilder_Good_SecuredResponses(t *testing.T) {
 	if _, ok := headers["X-RateLimit-Reset"]; !ok {
 		t.Fatal("expected X-RateLimit-Reset header in secured operation 429 response")
 	}
+	for _, code := range []string{"400", "401", "403", "504"} {
+		resp := responses[code].(map[string]any)
+		respHeaders := resp["headers"].(map[string]any)
+		if _, ok := respHeaders["X-Request-ID"]; !ok {
+			t.Fatalf("expected X-Request-ID header in secured operation %s response", code)
+		}
+		if _, ok := respHeaders["X-RateLimit-Limit"]; !ok {
+			t.Fatalf("expected X-RateLimit-Limit header in secured operation %s response", code)
+		}
+		if _, ok := respHeaders["X-RateLimit-Remaining"]; !ok {
+			t.Fatalf("expected X-RateLimit-Remaining header in secured operation %s response", code)
+		}
+		if _, ok := respHeaders["X-RateLimit-Reset"]; !ok {
+			t.Fatalf("expected X-RateLimit-Reset header in secured operation %s response", code)
+		}
+	}
 }
 
 func TestSpecBuilder_Good_EnvelopeWrapping(t *testing.T) {