From f672255aca87d4bcfbe3939fa247cff4685cc28e Mon Sep 17 00:00:00 2001
From: Lyon <88232613+pikasTech@users.noreply.github.com>
Date: Tue, 23 Jun 2026 08:20:09 +0800
Subject: [PATCH] fix(workbench-runtime): optimize session message summaries
 (#1958)

---
 internal/workbenchruntime/service.go | 58 +++++++++++-----------------
 1 file changed, 22 insertions(+), 36 deletions(-)

diff --git a/internal/workbenchruntime/service.go b/internal/workbenchruntime/service.go
index ef7857e0..b92f689d 100644
--- a/internal/workbenchruntime/service.go
+++ b/internal/workbenchruntime/service.go
@@ -1562,24 +1562,37 @@ func (s *Server) queryMessageSummaries(ctx context.Context, sessionIDs []string)
 		return []any{}, nil
 	}
 	args := []any{}
-	placeholders := []string{}
+	valueRows := []string{}
 	for _, sid := range cleaned {
 		args = append(args, sid)
-		placeholders = append(placeholders, fmt.Sprintf("$%d", len(args)))
+		valueRows = append(valueRows, fmt.Sprintf("($%d::text)", len(args)))
 	}
 	bySession := map[string]map[string]any{}
-	countSQL := fmt.Sprintf("SELECT session_id, COUNT(*) AS message_count FROM workbench_messages WHERE session_id IN (%s) GROUP BY session_id", strings.Join(placeholders, ", "))
-	countStage := "workbench_runtime.db.message_counts"
-	err := s.withOtelInternalSpan(ctx, countStage, s.dbQueryAttrs("workbench_messages", len(args), map[string]any{"db.index.expected": "idx_workbench_messages_session_updated"}), func(spanCtx context.Context) error {
-		rows, err := s.db.QueryContext(spanCtx, countSQL, args...)
+	summarySQL := fmt.Sprintf(`SELECT requested.session_id, COALESCE(message_counts.message_count, 0) AS message_count, first_user.message_json
+FROM (VALUES %s) AS requested(session_id)
+LEFT JOIN LATERAL (
+	SELECT COUNT(*) AS message_count
+	FROM workbench_messages
+	WHERE session_id = requested.session_id
+) AS message_counts ON TRUE
+LEFT JOIN LATERAL (
+	SELECT message_json
+	FROM workbench_messages
+	WHERE session_id = requested.session_id AND role = 'user'
+	ORDER BY updated_at ASC, message_id ASC
+	LIMIT 1
+) AS first_user ON TRUE`, strings.Join(valueRows, ", "))
+	summaryStage := "workbench_runtime.db.message_summaries"
+	err := s.withOtelInternalSpan(ctx, summaryStage, s.dbQueryAttrs("workbench_messages", len(args), map[string]any{"db.index.expected": "idx_workbench_messages_session_updated|idx_workbench_messages_session_role_updated", "db.query.shape": "values_lateral"}), func(spanCtx context.Context) error {
+		rows, err := s.db.QueryContext(spanCtx, summarySQL, args...)
 		if err != nil {
 			return err
 		}
 		defer rows.Close()
 		for rows.Next() {
-			var sessionID sql.NullString
+			var sessionID, firstUserMessageJSON sql.NullString
 			var messageCount sql.NullInt64
-			if err := rows.Scan(&sessionID, &messageCount); err != nil {
+			if err := rows.Scan(&sessionID, &messageCount, &firstUserMessageJSON); err != nil {
 				return err
 			}
 			if !sessionID.Valid || strings.TrimSpace(sessionID.String) == "" {
@@ -1587,33 +1600,6 @@ func (s *Server) queryMessageSummaries(ctx context.Context, sessionIDs []string)
 			}
 			item := map[string]any{"sessionId": sessionID.String, "messageCount": nullableInt(messageCount), "valuesRedacted": true}
 			bySession[sessionID.String] = item
-		}
-		return rows.Err()
-	})
-	if err != nil {
-		return nil, wrapQueryStage(countStage, err)
-	}
-	firstUserSQL := fmt.Sprintf("SELECT DISTINCT ON (session_id) session_id, message_json FROM workbench_messages WHERE session_id IN (%s) AND role = 'user' ORDER BY session_id ASC, updated_at ASC, message_id ASC", strings.Join(placeholders, ", "))
-	firstUserStage := "workbench_runtime.db.first_user_messages"
-	err = s.withOtelInternalSpan(ctx, firstUserStage, s.dbQueryAttrs("workbench_messages", len(args), map[string]any{"db.index.expected": "idx_workbench_messages_session_role_updated"}), func(spanCtx context.Context) error {
-		rows, err := s.db.QueryContext(spanCtx, firstUserSQL, args...)
-		if err != nil {
-			return err
-		}
-		defer rows.Close()
-		for rows.Next() {
-			var sessionID, firstUserMessageJSON sql.NullString
-			if err := rows.Scan(&sessionID, &firstUserMessageJSON); err != nil {
-				return err
-			}
-			if !sessionID.Valid || strings.TrimSpace(sessionID.String) == "" {
-				continue
-			}
-			item := bySession[sessionID.String]
-			if item == nil {
-				item = map[string]any{"sessionId": sessionID.String, "messageCount": 0, "valuesRedacted": true}
-				bySession[sessionID.String] = item
-			}
 			if firstUserMessageJSON.Valid && strings.TrimSpace(firstUserMessageJSON.String) != "" {
 				var message map[string]any
 				if json.Unmarshal([]byte(firstUserMessageJSON.String), &message) == nil && len(message) > 0 {
@@ -1624,7 +1610,7 @@ func (s *Server) queryMessageSummaries(ctx context.Context, sessionIDs []string)
 		return rows.Err()
 	})
 	if err != nil {
-		return nil, wrapQueryStage(firstUserStage, err)
+		return nil, wrapQueryStage(summaryStage, err)
 	}
 	result := []any{}
 	for _, sid := range cleaned {