From b5376a0d74eeb03908359b4804e8e559182e3620 Mon Sep 17 00:00:00 2001 From: Lyon <88232613+pikasTech@users.noreply.github.com> Date: Mon, 22 Jun 2026 16:43:03 +0800 Subject: [PATCH] fix: cache Workbench session message facts (#1904) --- internal/workbenchruntime/service.go | 75 +++++++++++++++++++++++++++- 1 file changed, 73 insertions(+), 2 deletions(-) diff --git a/internal/workbenchruntime/service.go b/internal/workbenchruntime/service.go index a5f1a6c1..ca1ef8d3 100644 --- a/internal/workbenchruntime/service.go +++ b/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 {