add tunnel to optimize sse sending (#74)

Author: hyprh

internal/sse/sse.go

@@ -156,3 +156,46 @@ func FormatJSONRPCEvent(w io.Writer, eventType string, id interface{}, data inte
 	}
 	return nil
 }
+
+// EventBatch represents a single event in a batch operation
+type EventBatch struct {
+	EventType string
+	ID        interface{}
+	Data      interface{}
+}
+
+// FormatJSONRPCEventBatch formats multiple JSON-RPC events in a single write operation.
+// This reduces the number of write calls and improves performance for high-frequency event streams.
+// All events are formatted into a single buffer and written once.
+func FormatJSONRPCEventBatch(w io.Writer, events []EventBatch) error {
+	if len(events) == 0 {
+		return nil
+	}
+
+	// Pre-allocate buffer for better performance
+	// Rough estimation: each event ~200-500 bytes
+	var buf bytes.Buffer
+	for _, event := range events {
+		// Create a JSON-RPC response with the data as the result
+		response := jsonrpc.NewNotificationResponse(event.ID, event.Data)
+		// Marshal the entire JSON-RPC envelope
+		jsonData, err := json.Marshal(response)
+		if err != nil {
+			return fmt.Errorf("failed to marshal JSON-RPC SSE event data: %w", err)
+		}
+		// Format according to text/event-stream specification
+		// event: <eventType>
+		// data: <jsonrpc_envelope>
+		// <empty line>
+		if _, err := fmt.Fprintf(&buf, "event: %s\ndata: %s\n\n", event.EventType, string(jsonData)); err != nil {
+			return fmt.Errorf("failed to format JSON-RPC SSE event: %w", err)
+		}
+	}
+
+	// Write the entire batch in one operation
+	if _, err := w.Write(buf.Bytes()); err != nil {
+		return fmt.Errorf("failed to write JSON-RPC SSE event batch: %w", err)
+	}
+
+	return nil
+}

server/server.go

@@ -284,7 +284,7 @@ func (s *A2AServer) parseJSONRPCRequest(w http.ResponseWriter, body io.ReadClose
 
 // routeJSONRPCMethod routes the request to the appropriate handler based on the method.
 func (s *A2AServer) routeJSONRPCMethod(ctx context.Context, w http.ResponseWriter, request jsonrpc.Request) {
-	log.Infof("Received JSON-RPC request (ID: %v, Method: %s)", request.ID, request.Method)
+	log.Debugf("Received JSON-RPC request (ID: %v, Method: %s)", request.ID, request.Method)
 
 	switch request.Method {
 
@@ -660,30 +660,9 @@ func handleSSEStream(
 	// Use request context to detect client disconnection.
 	clientClosed := ctx.Done()
 
-	// --- Event Forwarding Loop ---
-	for {
-		select {
-		case event, ok := <-eventsChan:
-			if !ok {
-				flusher.Flush()
-				return // End the handler.
-			}
-			if err := sendSSEEvent(w, rpcID, &event); err != nil {
-				if err == errUnknownEvent {
-					log.Warnf("Unknown event type received for request ID: %s: %T. Skipping.", rpcID, event)
-					continue
-				}
-				log.Errorf("Error writing SSE JSON-RPC event for request ID: %s (client likely disconnected): %v", rpcID, err)
-				return
-			}
-			// Flush the buffer to ensure the event is sent immediately.
-			flusher.Flush()
-		case <-clientClosed:
-			// Client disconnected (request context canceled).
-			log.Infof("SSE client disconnected for request ID: %s. Closing stream.", rpcID)
-			return // Exit the handler.
-		}
-	}
+	// Use optimized tunnel for batching events
+	tunnel := newSSETunnel(w, flusher, rpcID)
+	tunnel.start(ctx, eventsChan, clientClosed)
 }
 
 func sendSSEEvent(w http.ResponseWriter, rpcID string, event interface{}) error {

server/sse_tunnel.go

@@ -0,0 +1,141 @@
+package server
+
+import (
+	"context"
+	"net/http"
+	"time"
+
+	"trpc.group/trpc-go/trpc-a2a-go/internal/sse"
+	"trpc.group/trpc-go/trpc-a2a-go/log"
+	"trpc.group/trpc-go/trpc-a2a-go/protocol"
+)
+
+const (
+	defaultSSEBatchSize     = 20
+	defaultSSEFlushInterval = 1 * time.Second
+)
+
+// sseTunnel optimizes SSE streaming by batching events before sending
+type sseTunnel struct {
+	w             http.ResponseWriter
+	flusher       http.Flusher
+	rpcID         string
+	batchSize     int
+	flushInterval time.Duration
+	batch         []protocol.StreamingMessageEvent
+}
+
+// newSSETunnel creates a new SSE tunnel with default settings
+func newSSETunnel(w http.ResponseWriter, flusher http.Flusher, rpcID string) *sseTunnel {
+	return &sseTunnel{
+		w:             w,
+		flusher:       flusher,
+		rpcID:         rpcID,
+		batchSize:     defaultSSEBatchSize,
+		flushInterval: defaultSSEFlushInterval,
+		batch:         make([]protocol.StreamingMessageEvent, 0, defaultSSEBatchSize),
+	}
+}
+
+// start runs the SSE tunnel with event batching optimization
+func (t *sseTunnel) start(ctx context.Context, eventsChan <-chan protocol.StreamingMessageEvent, clientClosed <-chan struct{}) {
+	ticker := time.NewTicker(t.flushInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case event, ok := <-eventsChan:
+			if !ok {
+				// Channel closed, flush any remaining events and exit
+				if len(t.batch) > 0 {
+					t.flushBatch()
+				}
+				return
+			}
+
+			// Add event to batch
+			t.batch = append(t.batch, event)
+
+			// Flush if batch is full
+			if len(t.batch) >= t.batchSize {
+				ticker.Reset(t.flushInterval)
+				if !t.flushBatch() {
+					return // Error occurred, exit
+				}
+			}
+
+		case <-ticker.C:
+			// Periodic flush for any accumulated events
+			if len(t.batch) > 0 {
+				if !t.flushBatch() {
+					return // Error occurred, exit
+				}
+			}
+
+		case <-clientClosed:
+			// Client disconnected
+			log.Infof("SSE client disconnected for request ID: %s. Closing stream.", t.rpcID)
+			return
+		}
+	}
+}
+
+// flushBatch sends all events in the current batch as a single write operation
+func (t *sseTunnel) flushBatch() bool {
+	if len(t.batch) == 0 {
+		return true
+	}
+
+	// Convert to batch format for efficient processing
+	events := make([]sse.EventBatch, 0, len(t.batch))
+
+	// Process all events in the batch
+	for _, event := range t.batch {
+		eventType, err := t.getEventType(&event)
+		if err != nil {
+			if err == errUnknownEvent {
+				log.Warnf("Unknown event type received for request ID: %s: %T. Skipping.", t.rpcID, event)
+				continue
+			}
+			log.Errorf("Error determining event type for request ID: %s: %v", t.rpcID, err)
+			return false
+		}
+
+		// Add to batch events
+		events = append(events, sse.EventBatch{
+			EventType: eventType,
+			ID:        t.rpcID,
+			Data:      &event,
+		})
+	}
+
+	// Write the entire batch using optimized batch function
+	if err := sse.FormatJSONRPCEventBatch(t.w, events); err != nil {
+		log.Errorf("Error writing SSE batch for request ID: %s (client likely disconnected): %v", t.rpcID, err)
+		return false
+	}
+
+	// Clear the batch and flush to client
+	t.batch = t.batch[:0]
+	t.flusher.Flush()
+
+	return true
+}
+
+// getEventType determines the SSE event type from a StreamingMessageEvent
+func (t *sseTunnel) getEventType(event *protocol.StreamingMessageEvent) (string, error) {
+	actualEvent := event.Result
+
+	switch actualEvent.(type) {
+	case *protocol.TaskStatusUpdateEvent:
+		return protocol.EventStatusUpdate, nil
+	case *protocol.TaskArtifactUpdateEvent:
+		return protocol.EventArtifactUpdate, nil
+	case *protocol.Message:
+		return protocol.EventMessage, nil
+	case *protocol.Task:
+		return protocol.EventTask, nil
+	default:
+		return "", errUnknownEvent
+	}
+}