mirrors/pikasTech-HWLAB
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(workbench-runtime): optimize session message summaries (#1958)

Browse Source
This commit is contained in:
Lyon
2026-06-23 08:20:09 +08:00
committed by GitHub
parent 65f75dafec
commit f672255aca
1 changed files with 22 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/workbenchruntime/service.go
+22 -36
View File
@@ -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 {