chore(middleware): use huma middleware

DivyMohan14 · DivyMohan14 · commit d90bbcb97e8e · 2025-09-12T00:16:57.000+05:30
diff --git a/lib/httpapi/server.go b/lib/httpapi/server.go
@@ -189,23 +189,6 @@ func NewServer(ctx context.Context, config ServerConfig) (*Server, error) {
 	})
 	router.Use(corsMiddleware.Handler)
 
-	// Add SSE middleware to prevent proxy buffering
-	sseMiddleware := func(next http.Handler) http.Handler {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			if strings.HasSuffix(r.URL.Path, "/events") || strings.HasSuffix(r.URL.Path, "/screen") {
-				// Disable proxy buffering for SSE endpoints
-				w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
-				w.Header().Set("Pragma", "no-cache")
-				w.Header().Set("Expires", "0")
-				w.Header().Set("X-Accel-Buffering", "no") // nginx
-				w.Header().Set("X-Proxy-Buffering", "no") // generic proxy
-				w.Header().Set("Connection", "keep-alive")
-			}
-			next.ServeHTTP(w, r)
-		})
-	}
-	router.Use(sseMiddleware)
-
 	humaConfig := huma.DefaultConfig("AgentAPI", version.Version)
 	humaConfig.Info.Description = "HTTP API for Claude Code, Goose, and Aider.\n\nhttps://github.com/coder/agentapi"
 	api := humachi.New(router, humaConfig)
@@ -280,6 +263,19 @@ func hostAuthorizationMiddleware(allowedHosts []string, badHostHandler http.Hand
 	}
 }
 
+// sseMiddleware creates middleware that prevents proxy buffering for SSE endpoints
+func sseMiddleware(ctx huma.Context, next func(huma.Context)) {
+	// Disable proxy buffering for SSE endpoints
+	ctx.SetHeader("Cache-Control", "no-cache, no-store, must-revalidate")
+	ctx.SetHeader("Pragma", "no-cache")
+	ctx.SetHeader("Expires", "0")
+	ctx.SetHeader("X-Accel-Buffering", "no") // nginx
+	ctx.SetHeader("X-Proxy-Buffering", "no") // generic proxy
+	ctx.SetHeader("Connection", "keep-alive")
+
+	next(ctx)
+}
+
 func (s *Server) StartSnapshotLoop(ctx context.Context) {
 	s.conversation.StartSnapshotLoop(ctx)
 	go func() {
@@ -316,6 +312,7 @@ func (s *Server) registerRoutes() {
 		Path:        "/events",
 		Summary:     "Subscribe to events",
 		Description: "The events are sent as Server-Sent Events (SSE). Initially, the endpoint returns a list of events needed to reconstruct the current state of the conversation and the agent's status. After that, it only returns events that have occurred since the last event was sent.\n\nNote: When an agent is running, the last message in the conversation history is updated frequently, and the endpoint sends a new message update event each time.",
+		Middlewares: []func(huma.Context, func(huma.Context)){sseMiddleware},
 	}, map[string]any{
 		// Mapping of event type name to Go struct for that event.
 		"message_update": MessageUpdateBody{},
@@ -328,6 +325,7 @@ func (s *Server) registerRoutes() {
 		Path:        "/internal/screen",
 		Summary:     "Subscribe to screen",
 		Hidden:      true,
+		Middlewares: []func(huma.Context, func(huma.Context)){sseMiddleware},
 	}, map[string]any{
 		"screen": ScreenUpdateBody{},
 	}, s.subscribeScreen)
