Cache the rebuilt Duo chat history per workflow

Eduardo Bonetrequested to merge
631643-cache-the-rebuilt-chat-history into master

What does this MR do and why?

With the feature flag on, a duoMessages read of a workflow's latest checkpoint reuses a cached chat history and folds only the blobs written since the last read. Today every request folds every ui_chat_log blob on the chain, even when nothing changed, and a cache hit cuts that CPU time by about 70% in a local benchmark. An entry stays valid while each checkpoint it covers keeps its blob count, because a retry can add a blob to an old checkpoint.

Design notes

Today every request inflates and parses every ui_chat_log blob on the ancestor chain, including one full chat log snapshot per thread, even when nothing changed since the last request. In a 24-hour Kibana sample, 512 requests took over 5 seconds, and 70% of their time was Ruby CPU. The slowest workflow (7993706) has 723 ui_chat_log blobs, with 75.7 MB of zlib snapshots.

New class Gitlab::DuoWorkflow::ChatHistoryCache in ee/lib/gitlab/duo_workflow/chat_history_cache.rb. One entry per workflow. Key duo_workflows/chat_history/1/<workflow_id> (the 1 is a VERSION constant, bumped when ChannelValuesReconstructor#channel_history changes what it returns). Value: { blob_counts:, max_blob_id:, history: }, where blob_counts maps each thread_ts on the chain to its blob count. TTL 1 hour.

Read path per page: Workflow.chat_log_blob_counts_for counts the ui_chat_log blobs per thread_ts for the page's workflows. It replaces master's chat_log_thread_ts_for, which read the same index entries with DISTINCT, so the cache adds no query. The chain filter runs in Ruby. One Rails.cache.read_multi per page, and one write_multi per page, only when new blobs were folded and every one of them was loaded. A blob the chain fetch missed (a lagging replica, or a delete) would otherwise stay out of every later read.

Resume rule: every thread_ts in the entry must still be on the chain, with the same blob count. A changed count means a retry or a force rewrite added a blob to a cached checkpoint, and ChannelValuesReconstructor#ordered_blobs_for can then replace a blob in the middle of the fold. A fork drops a cached thread_ts from the chain. Both cases rebuild in full. The newer checkpoints load through the chain fetch. If one of their blobs has a lower id than max_blob_id (it took its id before a cached blob and committed after it), it sorts into the cached part of the fold. The history then rebuilds in full, with one extra chain fetch for that workflow.

Alternative counts: ChannelValuesReconstructor#channel_history takes a new from: argument (the cached history) and re-applies alternative_count to it, through a new private method count_alternatives. A retry on another branch changes the counts without changing this chain's blob ids.

Only the latest top-level checkpoint (thread_ts equal to latest_checkpoint_header.thread_ts) uses the cache. Older checkpoints on the same page (firstCheckpoint, duoWorkflowEvents with first: N) fold their whole chain as before, so they never replace the single entry. The cached checkpoint folds only the blobs of its newer checkpoints, because the page may already have loaded the cached blobs for an older checkpoint.

Queries: the cache adds no query. A page where every row hits skips the chain fetch, which loads the TOASTed data (about 76 MB for workflow 7993706), so the blob table sees one index-only query. A miss runs the same two blob queries as master. The request spec checks both counts.

Feature flag duo_workflow_cache_chat_history: type gitlab_com_derisk, default off, actor is the root namespace of the workflow's project or namespace. With the flag off, the code path is the same as master: no extra query and no Redis call.

A miss costs the same as today: first read, eviction, and parallel requests for the same workflow (up to 5 parallel requests for one workflow appeared in the Kibana sample). Parallel writes: the last write wins, and each entry is consistent with its own blob ids.

Security: the cache key is the workflow id. The cache is read only after GraphQL has authorized the workflow. The cached history is the same for every reader with access to the workflow. The additional_context filter (build_duo_message) runs after the cache read. Cached values are plain Hash, Array, and String objects from JSON.

Scenario CPU time Allocated objects
Flag off (same as master) 0.74-0.86 s 0.68M
Cache miss 0.77 s 0.68M
Cache hit 0.22 s 0.61M
One new checkpoint since the cached read 0.26-0.27 s 0.62M

Local benchmark on a production-shaped copy of a long chat (801 messages), on top of #631639 (closed), which made the uncached fold about 3x cheaper. Local CPU runs about 6x faster than production, so compare the ratios, not the seconds. In every scenario the messages matched the uncached read (same SHA-256 digest of the serialized messages). The cache entry for this workflow is 464 KB in Redis (Rails.cache marshals and zlib-compresses it).

This is one of five tasks under the parent issue. The others are #631639 (closed) (parse only the new part of each snapshot; it changes channel_history internals, so bump VERSION if its output changes), #631640 (closed) (skip the JSON stream validator), #631641 (gateway writes a delta next to each snapshot), and #631642 (frontend checks the last message first). This cache helps on hits; those tasks make misses cheaper.

Overlap: !258751 (merged) edits the comment above Workflow.history_blobs_for_workflows, next to the new method added here. Whichever MR merges second gets a small rebase conflict.

The MR also changes table_size from small to medium in db/docs/p_duo_workflows_checkpoint_blobs.yml. The database dictionary sizes a partitioned table by its largest partition, which is 36 GB on GitLab.com. The database review found the stale value.

References

Screenshots or screen recordings

No UI change.

How to set up and validate locally

This needs Duo Agent Platform set up on GDK with incremental checkpoints enabled.

  1. In a rails console, enable the flag: Feature.enable(:duo_workflow_cache_chat_history). Expect: no error.
  2. Open a project on GDK, for example http://gdk.test:3000/gitlab-org/gitlab-test, open the Duo Agentic Chat panel, and send three or four prompts. Expect: the prompts and replies show in the panel.
  3. Reload the page with the panel open. Expect: the panel loads the chat history.
  4. In the rails console, run Rails.cache.read(['duo_workflows', 'chat_history', 1, Ai::DuoWorkflows::Workflow.last.id])[:blob_counts].values.sum. Expect: a number equal to the count of ui_chat_log blobs on the chat's current branch.
  5. Send one more prompt, then reload the page. Expect: the new messages show, and the sum from step 4 grows.
  6. Disable the flag with Feature.disable(:duo_workflow_cache_chat_history) and reload. Expect: the same messages as in step 5.
  7. Without a Duo Agent Platform setup, run bin/rspec ee/spec/lib/gitlab/duo_workflow/chat_history_cache_spec.rb. Expect: all examples pass.
Edge cases covered by specs
  • A cache hit loads no blob and writes nothing.
  • A new checkpoint folds only its own blobs onto the cached history.
  • Blobs the page loaded for an older checkpoint aren't folded into the cached history again (unit spec, and a request spec that selects latestCheckpoint and firstCheckpoint together).
  • An alternative count that changed after the history was cached is re-applied.
  • A retry that adds a blob to a cached checkpoint rebuilds the history in full.
  • A latest checkpoint that forks off the cached chain rebuilds the history in full.
  • With the flag off, the page folds the chain on every read, even with a cached entry.
  • A pending blob the chain fetch didn't return keeps the entry unwritten, so the next read folds it.
  • Two copies of one workflow that see different latest checkpoints each get their own history; only the newest is cached.
  • A newer checkpoint's blob with a lower id than the cached ones rebuilds the history in full.
  • Each safety check (re-applied counts, the blob-count rule, the lower-id rule, the missing-blob rule, the pending filter, newest-only caching) was removed in turn, and in each case the spec written for it failed.
Query plans (Database Lab)

Source: Database Lab gitlab-production-main, 2026-10-01 session. joe reset ran before each cold plan. It clears the shared buffers but not the OS file cache, so a second cold run on the same pages shows a low read time. Compare buffers.

This MR adds no query. It changes master's Workflow.chat_log_thread_ts_for (renamed chat_log_blob_counts_for) from DISTINCT to GROUP BY with COUNT(*). The non-batched Workflow#chat_log_blob_counts (formerly #chat_log_thread_ts) runs the same query for one workflow. Both forms use the same index-only scan on idx_duo_wf_checkpoint_blobs_dedup and read the same buffers.

Query Densest workflow (8009873, 1,112 blobs) Page of 12 workflows (94 blobs)
Master, DISTINCT 298 buffers; 246.6 ms cold, 3.3 ms warm 108 buffers; 75.3 ms cold, 2.8 ms warm
This MR, GROUP BY 298 buffers; 162.1 ms cold (earlier run with the OS cache cold), 3.2 ms warm 108 buffers; 97.7 ms cold, 2.7 ms warm

The id pluck from the first version of this MR read 1,319 buffers on the densest workflow (791 ms cold), because no index holds both id and thread_ts. This version removes it.

Densest workflow, before (master):

SELECT DISTINCT "p_duo_workflows_checkpoint_blobs"."workflow_id", "p_duo_workflows_checkpoint_blobs"."thread_ts" FROM "p_duo_workflows_checkpoint_blobs" WHERE "p_duo_workflows_checkpoint_blobs"."project_id" = 69979615 AND "p_duo_workflows_checkpoint_blobs"."workflow_id" = 8009873 AND "p_duo_workflows_checkpoint_blobs"."workflow_created_at" = '2026-09-16 12:29:53.788194+00' AND "p_duo_workflows_checkpoint_blobs"."channel" = 'ui_chat_log'
 Unique  (cost=0.56..3.77 rows=1 width=45) (actual time=7.040..244.591 rows=1112 loops=1)
   Buffers: shared hit=36 read=262 dirtied=4
   I/O Timings: read=240.760 write=0.000
   ->  Index Only Scan using p_duo_workflows_checkpoint_b_project_id_workflow_id_threa_idx22 on gitlab_partitions_dynamic.p_duo_workflows_checkpoint_blobs_20260916 p_duo_workflows_checkpoint_blobs  (cost=0.56..3.76 rows=1 width=45) (actual time=7.038..242.900 rows=1112 loops=1)
         Index Cond: ((p_duo_workflows_checkpoint_blobs.project_id = 69979615) AND (p_duo_workflows_checkpoint_blobs.workflow_id = 8009873) AND (p_duo_workflows_checkpoint_blobs.channel = 'ui_chat_log'::text) AND (p_duo_workflows_checkpoint_blobs.workflow_created_at = '2026-09-16 12:29:53.788194+00'::timestamp with time zone))
         Heap Fetches: 54
         Buffers: shared hit=36 read=262 dirtied=4
         I/O Timings: read=240.760 write=0.000

Densest workflow, after:

SELECT COUNT(*) AS "count_all", "p_duo_workflows_checkpoint_blobs"."workflow_id" AS "p_duo_workflows_checkpoint_blobs_workflow_id", "p_duo_workflows_checkpoint_blobs"."thread_ts" AS "p_duo_workflows_checkpoint_blobs_thread_ts" FROM "p_duo_workflows_checkpoint_blobs" WHERE "p_duo_workflows_checkpoint_blobs"."project_id" = 69979615 AND "p_duo_workflows_checkpoint_blobs"."workflow_id" = 8009873 AND "p_duo_workflows_checkpoint_blobs"."workflow_created_at" = '2026-09-16 12:29:53.788194+00' AND "p_duo_workflows_checkpoint_blobs"."channel" = 'ui_chat_log' GROUP BY "p_duo_workflows_checkpoint_blobs"."workflow_id", "p_duo_workflows_checkpoint_blobs"."thread_ts"
 GroupAggregate  (cost=0.56..3.78 rows=1 width=53) (actual time=0.392..7.061 rows=1112 loops=1)
   Group Key: p_duo_workflows_checkpoint_blobs.thread_ts
   Buffers: shared hit=36 read=262 dirtied=4
   I/O Timings: read=5.610 write=0.000
   ->  Index Only Scan using p_duo_workflows_checkpoint_b_project_id_workflow_id_threa_idx22 on gitlab_partitions_dynamic.p_duo_workflows_checkpoint_blobs_20260916 p_duo_workflows_checkpoint_blobs  (cost=0.56..3.76 rows=1 width=45) (actual time=0.386..6.726 rows=1112 loops=1)
         Index Cond: ((p_duo_workflows_checkpoint_blobs.project_id = 69979615) AND (p_duo_workflows_checkpoint_blobs.workflow_id = 8009873) AND (p_duo_workflows_checkpoint_blobs.channel = 'ui_chat_log'::text) AND (p_duo_workflows_checkpoint_blobs.workflow_created_at = '2026-09-16 12:29:53.788194+00'::timestamp with time zone))
         Heap Fetches: 54
         Buffers: shared hit=36 read=262 dirtied=4
         I/O Timings: read=5.610 write=0.000

Page of 12 workflows from one project (gitlab-org/gitlab), after:

SELECT COUNT(*) AS "count_all", "p_duo_workflows_checkpoint_blobs"."workflow_id" AS "p_duo_workflows_checkpoint_blobs_workflow_id", "p_duo_workflows_checkpoint_blobs"."thread_ts" AS "p_duo_workflows_checkpoint_blobs_thread_ts" FROM "p_duo_workflows_checkpoint_blobs" WHERE "p_duo_workflows_checkpoint_blobs"."project_id" = 278964 AND "p_duo_workflows_checkpoint_blobs"."workflow_id" IN (8813740,8813738,8813737,8813726,8813720,8813701,8813682,8813676,8813669,8813666,8813661,8813660) AND "p_duo_workflows_checkpoint_blobs"."workflow_created_at" IN ('2026-09-30 04:24:44.628669+00','2026-09-30 04:24:41.529101+00','2026-09-30 04:24:37.558732+00','2026-09-30 04:24:20.603075+00','2026-09-30 04:23:30.25961+00','2026-09-30 04:22:45.211646+00','2026-09-30 04:21:41.528188+00','2026-09-30 04:21:23.411918+00','2026-09-30 04:20:38.187516+00','2026-09-30 04:20:16.032999+00','2026-09-30 04:20:00.495868+00','2026-09-30 04:19:59.794348+00') AND "p_duo_workflows_checkpoint_blobs"."channel" = 'ui_chat_log' GROUP BY "p_duo_workflows_checkpoint_blobs"."workflow_id", "p_duo_workflows_checkpoint_blobs"."thread_ts"
 GroupAggregate  (cost=0.42..25.86 rows=1 width=53) (actual time=0.359..2.256 rows=94 loops=1)
   Group Key: p_duo_workflows_checkpoint_blobs.workflow_id, p_duo_workflows_checkpoint_blobs.thread_ts
   Buffers: shared hit=14 read=94 dirtied=76
   I/O Timings: read=1.336 write=0.000
   ->  Index Only Scan using p_duo_workflows_checkpoint_b_project_id_workflow_id_threa_idx36 on gitlab_partitions_dynamic.p_duo_workflows_checkpoint_blobs_20260930 p_duo_workflows_checkpoint_blobs  (cost=0.42..25.84 rows=1 width=45) (actual time=0.334..2.214 rows=94 loops=1)
         Index Cond: ((p_duo_workflows_checkpoint_blobs.project_id = 278964) AND (p_duo_workflows_checkpoint_blobs.workflow_id = ANY ('{8813740,8813738,8813737,8813726,8813720,8813701,8813682,8813676,8813669,8813666,8813661,8813660}'::bigint[])) AND (p_duo_workflows_checkpoint_blobs.channel = 'ui_chat_log'::text) AND (p_duo_workflows_checkpoint_blobs.workflow_created_at = ANY ('{"2026-09-30 04:24:44.628669+00","2026-09-30 04:24:41.529101+00","2026-09-30 04:24:37.558732+00","2026-09-30 04:24:20.603075+00","2026-09-30 04:23:30.25961+00","2026-09-30 04:22:45.211646+00","2026-09-30 04:21:41.528188+00","2026-09-30 04:21:23.411918+00","2026-09-30 04:20:38.187516+00","2026-09-30 04:20:16.032999+00","2026-09-30 04:20:00.495868+00","2026-09-30 04:19:59.794348+00"}'::timestamp with time zone[])))
         Heap Fetches: 94
         Buffers: shared hit=14 read=94 dirtied=76
         I/O Timings: read=1.336 write=0.000

Chain fetch after a partial hit, on the densest workflow with its three newest checkpoints (1 blob each). This is the usual read once the cache is warm. Cold (after joe reset): 4.8 ms, 12 buffers (6 read). Warm: 4.0 ms, same plan.

SELECT "p_duo_workflows_checkpoint_blobs".* FROM "p_duo_workflows_checkpoint_blobs" WHERE "p_duo_workflows_checkpoint_blobs"."project_id" = 69979615 AND "p_duo_workflows_checkpoint_blobs"."workflow_id" = 8009873 AND "p_duo_workflows_checkpoint_blobs"."workflow_created_at" = '2026-09-16 12:29:53.788194+00' AND "p_duo_workflows_checkpoint_blobs"."thread_ts" IN ('1f1b691e-47f2-6f02-8480-10b6431696eb', '1f1b691d-b19b-6717-847f-856ab423fd78', '1f1b691b-1413-6efd-847e-44f0eee6fa68') AND "p_duo_workflows_checkpoint_blobs"."channel" = 'ui_chat_log' ORDER BY "p_duo_workflows_checkpoint_blobs"."id" ASC
Cold (after joe reset): 4.812 ms (planning 3.717 ms, execution 1.095 ms), 12 buffers (6 read)
 Sort  (cost=7.75..7.75 rows=1 width=348) (actual time=1.054..1.056 rows=3 loops=1)
   Sort Key: p_duo_workflows_checkpoint_blobs.id
   Sort Method: quicksort  Memory: 26kB
   Buffers: shared hit=6 read=6
   I/O Timings: read=0.980 write=0.000
   ->  Index Scan using p_duo_workflows_checkpoint_b_project_id_workflow_id_threa_idx22 on gitlab_partitions_dynamic.p_duo_workflows_checkpoint_blobs_20260916 p_duo_workflows_checkpoint_blobs  (cost=0.56..7.74 rows=1 width=348) (actual time=1.017..1.037 rows=3 loops=1)
         Index Cond: ((p_duo_workflows_checkpoint_blobs.project_id = 69979615) AND (p_duo_workflows_checkpoint_blobs.workflow_id = 8009873) AND (p_duo_workflows_checkpoint_blobs.thread_ts = ANY ('{1f1b691e-47f2-6f02-8480-10b6431696eb,1f1b691d-b19b-6717-847f-856ab423fd78,1f1b691b-1413-6efd-847e-44f0eee6fa68}'::text[])) AND (p_duo_workflows_checkpoint_blobs.channel = 'ui_chat_log'::text) AND (p_duo_workflows_checkpoint_blobs.workflow_created_at = '2026-09-16 12:29:53.788194+00'::timestamp with time zone))
         Buffers: shared hit=3 read=6
         I/O Timings: read=0.980 write=0.000

Warm: 4.005 ms (planning 3.894 ms, execution 0.111 ms), 12 buffers, same plan

The chain fetch (Workflow.history_blobs_for_workflows) keeps its master shape. On a page where every row hits, it doesn't run. When new blobs exist, its thread_ts list holds only the checkpoints the cache doesn't cover. The rare rebuild after a lower-id blob runs Workflow#history_blobs_for for that workflow, the existing single-workflow chain fetch.

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Edited by Eduardo Bonet

Merge request reports