fix: ensure ResponseComplete hook always executes

pkg/epp/handlers/response.go

@@ -96,10 +96,6 @@ func (s *StreamingServer) HandleResponseBodyModelStreaming(ctx context.Context,
 			cachedToken = reqCtx.Usage.PromptTokenDetails.CachedTokens
 		}
 		metrics.RecordPromptCachedTokens(reqCtx.IncomingModelName, reqCtx.TargetModelName, cachedToken)
-		_, err := s.director.HandleResponseBodyComplete(ctx, reqCtx)
-		if err != nil {
-			logger.Error(err, "error in HandleResponseBodyComplete")
-		}
 	}
 }
 

pkg/epp/handlers/server.go

@@ -165,6 +165,17 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 		if reqCtx.RequestRunning {
 			metrics.DecRunningRequests(reqCtx.IncomingModelName)
 		}
+
+		// If we scheduled a pod (TargetPod != nil) but never marked the response  as complete (e.g. error, disconnect,
+		// panic), force the completion hooks to run.
+		if reqCtx.TargetPod != nil && !reqCtx.ResponseComplete {
+			// Use a fresh context as the request context might be canceled (Client Disconnect).
+			// We only need logging from the original context.
+			cleanupCtx := log.IntoContext(context.Background(), logger)
+			if _, err := s.director.HandleResponseBodyComplete(cleanupCtx, reqCtx); err != nil {
+				logger.Error(err, "error in HandleResponseBodyComplete")
+			}
+		}
 	}(err, reqCtx)
 
 	for {
@@ -275,6 +286,10 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 				s.HandleResponseBodyModelStreaming(ctx, reqCtx, responseText)
 				if v.ResponseBody.EndOfStream {
 					loggerTrace.Info("stream completed")
+					reqCtx.ResponseComplete = true
+					if _, err := s.director.HandleResponseBodyComplete(ctx, reqCtx); err != nil {
+						logger.Error(err, "error in HandleResponseBodyComplete")
+					}
 
 					reqCtx.ResponseCompleteTimestamp = time.Now()
 					metrics.RecordRequestLatencies(ctx, reqCtx.IncomingModelName, reqCtx.TargetModelName, reqCtx.RequestReceivedTimestamp, reqCtx.ResponseCompleteTimestamp)

pkg/epp/requestcontrol/plugins.go

@@ -52,7 +52,12 @@ type ResponseStreaming interface {
 	ResponseStreaming(ctx context.Context, request *types.LLMRequest, response *Response, targetPod *backend.Pod)
 }
 
-// ResponseComplete is called by the director after the complete response is sent.
+// ResponseComplete is called by the director when the request lifecycle terminates.
+// This occurs after a response is fully sent, OR if the request fails/disconnects after a pod was scheduled.
+//
+// Plugins should assume this is the final cleanup hook for a request.
+//
+// TODO: Consider passing an error or success bool; however, this is a breaking change and is deffered for now.
 type ResponseComplete interface {
 	plugins.Plugin
 	ResponseComplete(ctx context.Context, request *types.LLMRequest, response *Response, targetPod *backend.Pod)