From fc09614f243da59233a1ab901d432726fb160b21 Mon Sep 17 00:00:00 2001
From: Lyon <88232613+pikasTech@users.noreply.github.com>
Date: Fri, 26 Jun 2026 02:46:37 +0800
Subject: [PATCH] fix: cache workbench session read pages (#2172)

---
 internal/workbenchruntime/cache_test.go | 10 +++++++++-
 internal/workbenchruntime/service.go    | 18 ++++++++++++------
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/internal/workbenchruntime/cache_test.go b/internal/workbenchruntime/cache_test.go
index 170f53f7..a16fdb89 100644
--- a/internal/workbenchruntime/cache_test.go
+++ b/internal/workbenchruntime/cache_test.go
@@ -331,13 +331,21 @@ func TestFactsCacheClassAndTTL(t *testing.T) {
 	if got := factsCacheClass(turnSnapshot); got != terminalTurnCacheClass {
 		t.Fatalf("turn snapshot class=%q", got)
 	}
+	sessionMessagePage := factQuery{SessionID: "ses_1", Families: []string{"sessions", "messages", "parts", "turns", "checkpoints"}, Limit: 100}
+	if got := factsCacheClass(sessionMessagePage); got != sessionMessagePageCacheClass {
+		t.Fatalf("session message page class=%q", got)
+	}
 	uncacheable := factQuery{TraceID: "trc_1", Families: []string{"sessions", "messages"}}
 	if got := factsCacheClass(uncacheable); got != "" {
 		t.Fatalf("unexpected cache class=%q", got)
 	}
 
 	config := cacheConfig{TerminalTurnTTL: 5 * time.Minute, TerminalTracePageTTL: 10 * time.Minute, RunningPageTTL: 5 * time.Second}
-	ttl, unstable, reason := factsCacheTTL(factsCachePlan{Class: terminalTracePageCacheClass, ProjectionStatus: "caught-up", Status: "completed"}, config)
+	ttl, unstable, reason := factsCacheTTL(factsCachePlan{Class: sessionMessagePageCacheClass}, config)
+	if ttl != 5*time.Second || !unstable || reason != "session_message_page" {
+		t.Fatalf("session message page ttl=%s unstable=%v reason=%s", ttl, unstable, reason)
+	}
+	ttl, unstable, reason = factsCacheTTL(factsCachePlan{Class: terminalTracePageCacheClass, ProjectionStatus: "caught-up", Status: "completed"}, config)
 	if ttl != 10*time.Minute || unstable || reason != "terminal_trace_page" {
 		t.Fatalf("terminal trace ttl=%s unstable=%v reason=%s", ttl, unstable, reason)
 	}
diff --git a/internal/workbenchruntime/service.go b/internal/workbenchruntime/service.go
index c3de380b..0e919bc5 100644
--- a/internal/workbenchruntime/service.go
+++ b/internal/workbenchruntime/service.go
@@ -105,7 +105,9 @@ var workbenchRuntimeReadIndexes = []string{
 	"CREATE INDEX IF NOT EXISTS idx_workbench_messages_session_updated ON workbench_messages(session_id, updated_at ASC, message_id)",
 	"CREATE INDEX IF NOT EXISTS idx_workbench_messages_session_role_updated ON workbench_messages(session_id, role, updated_at ASC, message_id)",
 	"CREATE INDEX IF NOT EXISTS idx_workbench_parts_session_updated ON workbench_parts(session_id, updated_at ASC, part_id)",
+	"CREATE INDEX IF NOT EXISTS idx_workbench_turns_session_updated ON workbench_turns(session_id, updated_at DESC, turn_id)",
 	"CREATE INDEX IF NOT EXISTS idx_workbench_turns_trace_updated ON workbench_turns(trace_id, updated_at DESC)",
+	"CREATE INDEX IF NOT EXISTS idx_workbench_projection_checkpoints_session_updated ON workbench_projection_checkpoints(session_id, updated_at DESC, trace_id)",
 	"CREATE INDEX IF NOT EXISTS idx_workbench_projection_checkpoints_trace_updated ON workbench_projection_checkpoints(trace_id, updated_at DESC)",
 }
 
@@ -548,15 +550,17 @@ func (s *Server) sessionMessagePageRevision(ctx context.Context, sessionID strin
 	const sqlText = `SELECT
 		COALESCE((SELECT MAX(updated_at)::text FROM workbench_sessions WHERE session_id = $1), ''),
 		COALESCE((SELECT MAX(updated_at)::text FROM workbench_messages WHERE session_id = $1), ''),
-		COALESCE((SELECT MAX(updated_at)::text FROM workbench_parts WHERE session_id = $1), '')`
-	var sessionUpdatedAt, messagesUpdatedAt, partsUpdatedAt sql.NullString
+		COALESCE((SELECT MAX(updated_at)::text FROM workbench_parts WHERE session_id = $1), ''),
+		COALESCE((SELECT MAX(updated_at)::text FROM workbench_turns WHERE session_id = $1), ''),
+		COALESCE((SELECT MAX(updated_at)::text FROM workbench_projection_checkpoints WHERE session_id = $1), '')`
+	var sessionUpdatedAt, messagesUpdatedAt, partsUpdatedAt, turnsUpdatedAt, checkpointsUpdatedAt sql.NullString
 	err := s.withOtelInternalSpan(ctx, "workbench_runtime.db.session_message_revision", s.dbQueryAttrs("workbench_session_message_revision", 1, nil), func(spanCtx context.Context) error {
-		return s.db.QueryRowContext(spanCtx, sqlText, sessionID).Scan(&sessionUpdatedAt, &messagesUpdatedAt, &partsUpdatedAt)
+		return s.db.QueryRowContext(spanCtx, sqlText, sessionID).Scan(&sessionUpdatedAt, &messagesUpdatedAt, &partsUpdatedAt, &turnsUpdatedAt, &checkpointsUpdatedAt)
 	})
 	if err != nil {
 		return "", wrapQueryStage("workbench_runtime.db.session_message_revision", err)
 	}
-	raw := strings.Join([]string{sessionID, nullableSQLString(sessionUpdatedAt), nullableSQLString(messagesUpdatedAt), nullableSQLString(partsUpdatedAt)}, "|")
+	raw := strings.Join([]string{sessionID, nullableSQLString(sessionUpdatedAt), nullableSQLString(messagesUpdatedAt), nullableSQLString(partsUpdatedAt), nullableSQLString(turnsUpdatedAt), nullableSQLString(checkpointsUpdatedAt)}, "|")
 	return "session-message-page:" + cacheHash(raw), nil
 }
 
@@ -598,8 +602,10 @@ func (s *Server) traceProjectionState(ctx context.Context, traceID string) (trac
 
 func factsCacheClass(query factQuery) string {
 	families := requestedFactFamilies(query)
-	if strings.TrimSpace(query.SessionID) != "" && query.AfterProjectedSeq == 0 && stringSetEqual(families, []string{"sessions", "messages", "parts"}) {
-		return sessionMessagePageCacheClass
+	if strings.TrimSpace(query.SessionID) != "" && query.AfterProjectedSeq == 0 {
+		if stringSetEqual(families, []string{"sessions", "messages", "parts"}) || stringSetEqual(families, []string{"sessions", "messages", "parts", "turns", "checkpoints"}) {
+			return sessionMessagePageCacheClass
+		}
 	}
 	if strings.TrimSpace(query.TraceID) == "" {
 		return ""