From 4b7a06e299badd7b8357946314ef47f13ecb1c29 Mon Sep 17 00:00:00 2001
From: Lyon <88232613+pikasTech@users.noreply.github.com>
Date: Sun, 21 Jun 2026 22:08:41 +0800
Subject: [PATCH] fix: harden workbench runtime session summary query (#1843)

---
 internal/workbenchruntime/otel.go    | 36 +++++++++++++++++++++++++++-
 internal/workbenchruntime/service.go | 29 +++++++++++++++++-----
 2 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/internal/workbenchruntime/otel.go b/internal/workbenchruntime/otel.go
index d056a966..f53bbe37 100644
--- a/internal/workbenchruntime/otel.go
+++ b/internal/workbenchruntime/otel.go
@@ -174,7 +174,7 @@ func (s *Server) withOtelInternalSpan(ctx context.Context, name string, attrs ma
 	err := fn(ctx)
 	errorCode := ""
 	if err != nil {
-		errorCode = otelFirst(queryErrorStage(err), sqlErrorCode(err), "query_failed")
+		errorCode = otelFirst(sqlErrorCode(err), sqlCauseCode(err), queryErrorStage(err), "query_failed")
 	}
 	attributes := map[string]any{
 		"hwlab.runtime.stage": name,
@@ -185,10 +185,44 @@ func (s *Server) withOtelInternalSpan(ctx context.Context, name string, attrs ma
 			attributes[key] = value
 		}
 	}
+	addOtelErrorCauseAttributes(attributes, err)
 	s.emitOtelSpanAsync(name, trace, spanID, trace.ServerSpanID, otelSpanKindInternal, startedAt, attributes, 0, errorCode)
 	return err
 }
 
+func addOtelErrorCauseAttributes(attributes map[string]any, err error) {
+	if err == nil {
+		return
+	}
+	if stage := queryErrorStage(err); stage != "" {
+		attributes["causeStage"] = stage
+	}
+	if code := sqlErrorCode(err); code != "" {
+		attributes["causeCode"] = code
+		return
+	}
+	if code := sqlCauseCode(err); code != "" {
+		attributes["causeCode"] = code
+	}
+}
+
+func sqlCauseCode(err error) string {
+	if err == nil {
+		return ""
+	}
+	message := strings.ToLower(err.Error())
+	switch {
+	case strings.Contains(message, "invalid input syntax for type json"):
+		return "invalid_json_text"
+	case strings.Contains(message, "cannot cast") && strings.Contains(message, "json"):
+		return "invalid_json_cast"
+	case strings.Contains(message, "context deadline exceeded"):
+		return "query_timeout"
+	default:
+		return ""
+	}
+}
+
 func (s *Server) emitOtelSpanAsync(name string, trace otelTraceContext, spanID string, parentSpanID string, kind int, startedAt time.Time, attrs map[string]any, httpStatus int, errorCode string) {
 	go s.emitOtelSpan(context.Background(), name, trace, spanID, parentSpanID, kind, startedAt, attrs, httpStatus, errorCode)
 }
diff --git a/internal/workbenchruntime/service.go b/internal/workbenchruntime/service.go
index 9a17e5f7..191bfae8 100644
--- a/internal/workbenchruntime/service.go
+++ b/internal/workbenchruntime/service.go
@@ -520,11 +520,11 @@ func (s *Server) querySessionSummaries(ctx context.Context, stage string, sqlTex
 		}
 		defer rows.Close()
 		for rows.Next() {
-			var sessionID, ownerUserID, projectID, conversationID, threadID, status, lastTraceID, sourceEventID, providerProfile sql.NullString
+			var sessionID, ownerUserID, projectID, conversationID, threadID, status, lastTraceID, sourceEventID, sessionJSON sql.NullString
 			var projectedSeq, sourceSeq sql.NullInt64
 			var terminal, sealed sql.NullBool
 			var createdAt, updatedAt sql.NullTime
-			if err := rows.Scan(&sessionID, &ownerUserID, &projectID, &conversationID, &threadID, &status, &lastTraceID, &projectedSeq, &sourceSeq, &sourceEventID, &terminal, &sealed, &createdAt, &updatedAt, &providerProfile); err != nil {
+			if err := rows.Scan(&sessionID, &ownerUserID, &projectID, &conversationID, &threadID, &status, &lastTraceID, &projectedSeq, &sourceSeq, &sourceEventID, &terminal, &sealed, &createdAt, &updatedAt, &sessionJSON); err != nil {
 				return err
 			}
 			if !sessionID.Valid || strings.TrimSpace(sessionID.String) == "" {
@@ -549,9 +549,9 @@ func (s *Server) querySessionSummaries(ctx context.Context, stage string, sqlTex
 				"valuesPrinted":  false,
 				"valuesRedacted": true,
 			}
-			if providerProfile.Valid && strings.TrimSpace(providerProfile.String) != "" {
-				item["providerProfile"] = providerProfile.String
-				item["sessionJson"] = map[string]any{"providerProfile": providerProfile.String, "valuesRedacted": true, "secretMaterialStored": false}
+			if providerProfile := providerProfileFromSessionJSON(sessionJSON); providerProfile != "" {
+				item["providerProfile"] = providerProfile
+				item["sessionJson"] = map[string]any{"providerProfile": providerProfile, "valuesRedacted": true, "secretMaterialStored": false}
 			}
 			result = append(result, item)
 		}
@@ -820,10 +820,27 @@ func sessionSummarySelectClause() string {
 		"sealed",
 		"created_at",
 		"updated_at",
-		"COALESCE(NULLIF(session_json::jsonb ->> 'providerProfile', ''), NULLIF(session_json::jsonb #>> '{sessionJson,providerProfile}', '')) AS provider_profile",
+		"session_json",
 	}, ", ")
 }
 
+func providerProfileFromSessionJSON(raw sql.NullString) string {
+	if !raw.Valid || strings.TrimSpace(raw.String) == "" {
+		return ""
+	}
+	var payload map[string]any
+	if err := json.Unmarshal([]byte(raw.String), &payload); err != nil {
+		return ""
+	}
+	if value := strings.TrimSpace(text(payload["providerProfile"])); value != "" {
+		return value
+	}
+	if value := strings.TrimSpace(text(objectMap(payload["sessionJson"])["providerProfile"])); value != "" {
+		return value
+	}
+	return ""
+}
+
 func boundedLimit(value int, fallback int, max int) int {
 	if value <= 0 {
 		return fallback