feat(sse): allow specifying the schema for SSE operations

sse/sse.go

@@ -82,6 +82,8 @@ func (s Sender) Comment(comment string) error {
 // the context, input, and a `send` function that can be used to send messages
 // to the client. Flushing is handled automatically as long as the adapter's
 // `BodyWriter` implements `http.Flusher`.
+// The default response schema can be overridden by setting the
+// `op.Responses["200"].Content["text/event-stream"]` field before calling this function.
 func Register[I any](api huma.API, op huma.Operation, eventTypeMap map[string]any, f func(ctx context.Context, input *I, send Sender)) {
 	// Start by defining the SSE schema & operation response.
 	if op.Responses == nil {
@@ -144,8 +146,15 @@ func Register[I any](api huma.API, op huma.Operation, eventTypeMap map[string]an
 			},
 		},
 	}
-	op.Responses["200"].Content["text/event-stream"] = &huma.MediaType{
-		Schema: schema,
+
+	opMediaType := op.Responses["200"].Content["text/event-stream"]
+
+	if opMediaType == nil {
+		op.Responses["200"].Content["text/event-stream"] = &huma.MediaType{
+			Schema: schema,
+		}
+	} else if opMediaType.Schema == nil {
+		opMediaType.Schema = schema
 	}
 
 	// Register the operation with the API, using the built-in streaming

sse/sse_test.go

@@ -93,6 +93,77 @@ type sseTest struct {
 }
 
 var sseTests = []sseTest{
+	{
+		Title: "sse preserves existing text/event-stream content",
+		TestFunc: func(t *testing.T) {
+			_, api := humatest.New(t)
+			existingSchema := &huma.Schema{
+				Title:       "Existing Schema",
+				Description: "This schema should be preserved",
+				Type:        huma.TypeString,
+			}
+
+			sse.Register(api, huma.Operation{
+				OperationID: "sse",
+				Method:      http.MethodGet,
+				Path:        "/sse",
+				Responses: map[string]*huma.Response{
+					"200": {
+						Content: map[string]*huma.MediaType{
+							"text/event-stream": {
+								Schema: existingSchema,
+							},
+						},
+					},
+				},
+			}, map[string]any{
+				"message": &DefaultMessage{},
+			}, func(ctx context.Context, input *struct{}, send sse.Sender) {
+				send.Data(DefaultMessage{Message: "Hello, world!"})
+			})
+
+			o := api.OpenAPI()
+			response := o.Paths["/sse"].Get.Responses["200"]
+			content := response.Content["text/event-stream"]
+
+			assert.Equal(t, "Existing Schema", content.Schema.Title)
+			assert.Equal(t, "This schema should be preserved", content.Schema.Description)
+			assert.Equal(t, huma.TypeString, content.Schema.Type)
+		},
+	},
+	{
+		Title: "sse sets default schema and preserves other fields in text/event-stream content",
+		TestFunc: func(t *testing.T) {
+			_, api := humatest.New(t)
+
+			example := DefaultMessage{Message: "Example message"}
+			sse.Register(api, huma.Operation{
+				OperationID: "sse",
+				Method:      http.MethodGet,
+				Path:        "/sse",
+				Responses: map[string]*huma.Response{
+					"200": {
+						Content: map[string]*huma.MediaType{
+							"text/event-stream": {
+								Example: example,
+							},
+						},
+					},
+				},
+			}, map[string]any{
+				"message": &DefaultMessage{},
+			}, func(ctx context.Context, input *struct{}, send sse.Sender) {
+				send.Data(DefaultMessage{Message: "Hello, world!"})
+			})
+
+			o := api.OpenAPI()
+			response := o.Paths["/sse"].Get.Responses["200"]
+			content := response.Content["text/event-stream"]
+
+			assert.Equal(t, example, content.Example)
+			assert.NotNil(t, content.Schema)
+		},
+	},
 	{
 		Title: "sse",
 		TestFunc: func(t *testing.T) {
@@ -263,6 +334,34 @@ data: {"message":"Hello, world!"}
 				"no body write should occur before the user handler sends an event")
 		},
 	},
+	{
+		Title: "sse adds text/event-stream content when not present",
+		TestFunc: func(t *testing.T) {
+			_, api := humatest.New(t)
+
+			sse.Register(api, huma.Operation{
+				OperationID: "sse",
+				Method:      http.MethodGet,
+				Path:        "/sse",
+			}, map[string]any{
+				"message": &DefaultMessage{},
+			}, func(ctx context.Context, input *struct{}, send sse.Sender) {
+				send.Data(DefaultMessage{Message: "Hello, world!"})
+			})
+
+			o := api.OpenAPI()
+			response := o.Paths["/sse"].Get.Responses["200"]
+			content := response.Content["text/event-stream"]
+
+			// The SSE schema should be added
+			assert.NotNil(t, content, "text/event-stream content should be present")
+			assert.NotNil(t, content.Schema, "schema should be present")
+			assert.Equal(t, "Server Sent Events", content.Schema.Title)
+			assert.Equal(t, huma.TypeArray, content.Schema.Type)
+			assert.NotNil(t, content.Schema.Items, "schema items should be present")
+			assert.NotNil(t, content.Schema.Items.Extensions["oneOf"], "oneOf extension should be present")
+		},
+	},
 	{
 		Title: "sse stable event order in openapi",
 		TestFunc: func(t *testing.T) {