fix: cache Workbench session message facts (#1904)

internal/workbenchruntime/service.go

@@ -25,6 +25,7 @@ const (
 	sessionSummaryQueryAttemptTimeout = 2 * time.Second
 	sessionSummaryQueryRetryBackoff   = 150 * time.Millisecond
 	sessionsSummaryCacheClass         = "sessions.summary"
+	sessionMessagePageCacheClass      = "session.message_page"
 	terminalTurnCacheClass            = "turn.terminal.snapshot"
 	terminalTracePageCacheClass       = "trace.terminal.events"
 )
@@ -53,6 +54,7 @@ type factsCachePlan struct {
 	Class            string
 	Key              string
 	Meta             map[string]any
+	SessionID        string
 	TraceID          string
 	ProjectionSeq    int64
 	ProjectionStatus string
@@ -90,6 +92,7 @@ var workbenchRuntimeReadIndexes = []string{
 	"CREATE INDEX IF NOT EXISTS idx_workbench_sessions_status_updated ON workbench_sessions(status, updated_at DESC, session_id)",
 	"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_trace_updated ON workbench_turns(trace_id, updated_at DESC)",
 	"CREATE INDEX IF NOT EXISTS idx_workbench_projection_checkpoints_trace_updated ON workbench_projection_checkpoints(trace_id, updated_at DESC)",
 }
@@ -439,11 +442,43 @@ func factsResponsePayload(facts map[string][]any, persistence map[string]any) ma
 func (s *Server) factsCachePlan(ctx context.Context, query factQuery) factsCachePlan {
 	class := factsCacheClass(query)
 	meta := factsCacheMetadata(class, "disabled")
-	plan := factsCachePlan{Class: class, Meta: meta, TraceID: strings.TrimSpace(query.TraceID)}
+	plan := factsCachePlan{Class: class, Meta: meta, SessionID: strings.TrimSpace(query.SessionID), TraceID: strings.TrimSpace(query.TraceID)}
 	if s == nil || s.cache == nil || !s.config.Cache.Enabled || class == "" {
 		return plan
 	}
 	meta["enabled"] = true
+	if class == sessionMessagePageCacheClass {
+		if plan.SessionID == "" {
+			setFactsCacheStatus(meta, "skipped")
+			meta["skipReason"] = "missing_session"
+			return plan
+		}
+		revision, err := s.sessionMessagePageRevision(ctx, plan.SessionID)
+		if err != nil {
+			setFactsCacheStatus(meta, "skipped")
+			meta["skipReason"] = "session_revision_unavailable"
+			return plan
+		}
+		if revision == "" {
+			setFactsCacheStatus(meta, "skipped")
+			meta["skipReason"] = "missing_session_revision"
+			return plan
+		}
+		cursor := fmt.Sprintf("limit:%d|families:%s", boundedLimit(query.Limit, 0, 1000), strings.Join(requestedFactFamilies(query), ","))
+		authority := fmt.Sprintf("contract=workbench-runtime-facts-v1|class=%s", class)
+		key, err := s.cache.BuildKey(cacheKeyParts{Class: class, ActorID: "workbench-runtime", SessionID: plan.SessionID, Cursor: cursor, ProjectionRevision: revision, Authority: authority})
+		if err != nil {
+			setFactsCacheStatus(meta, "skipped")
+			meta["skipReason"] = classifyCacheError(err)
+			return plan
+		}
+		meta["projectionRevision"] = revision
+		meta["revisionSource"] = "session_message_page_updated_at"
+		meta["invalidationMode"] = "session_revision_key"
+		plan.Key = key
+		setFactsCacheStatus(meta, "miss")
+		return plan
+	}
 	if plan.TraceID == "" {
 		setFactsCacheStatus(meta, "skipped")
 		meta["skipReason"] = "missing_trace"
@@ -484,6 +519,33 @@ func (s *Server) factsCachePlan(ctx context.Context, query factQuery) factsCache
 	return plan
 }
 
+func (s *Server) sessionMessagePageRevision(ctx context.Context, sessionID string) (string, error) {
+	sessionID = strings.TrimSpace(sessionID)
+	if sessionID == "" {
+		return "", nil
+	}
+	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
+	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)
+	})
+	if err != nil {
+		return "", wrapQueryStage("workbench_runtime.db.session_message_revision", err)
+	}
+	raw := strings.Join([]string{sessionID, nullableSQLString(sessionUpdatedAt), nullableSQLString(messagesUpdatedAt), nullableSQLString(partsUpdatedAt)}, "|")
+	return "session-message-page:" + cacheHash(raw), nil
+}
+
+func nullableSQLString(value sql.NullString) string {
+	if !value.Valid {
+		return ""
+	}
+	return strings.TrimSpace(value.String)
+}
+
 type traceProjectionState struct {
 	ProjectedSeq     int64
 	ProjectionStatus string
@@ -514,10 +576,13 @@ 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.TraceID) == "" {
 		return ""
 	}
-	families := requestedFactFamilies(query)
 	if stringSetEqual(families, []string{"traceEvents"}) {
 		return terminalTracePageCacheClass
 	}
@@ -561,6 +626,12 @@ func stringSetEqual(values []string, expected []string) bool {
 }
 
 func factsCacheTTL(plan factsCachePlan, config cacheConfig) (time.Duration, bool, string) {
+	if plan.Class == sessionMessagePageCacheClass {
+		if config.RunningPageTTL <= 0 {
+			return 0, true, "session_message_page_ttl_disabled"
+		}
+		return config.RunningPageTTL, true, "session_message_page"
+	}
 	stableProjection := plan.ProjectionStatus == "caught-up" || plan.ProjectionStatus == "blocked" || plan.ProjectionStatus == "stalled"
 	terminal := isTerminal(plan.Status) || stableProjection
 	if !terminal {