feat(flags): preload referenced cohorts in flags hypercache

[haacked]

Problem

The Rust feature-flags service queries PostgreSQL via CohortCacheManager for every team's cohorts on each /flags request, even though the cohort definitions rarely change. This adds unnecessary latency and PG load, especially for teams with many cohort-based flags.

Changes

Preloads cohort definitions (including transitive dependencies) into the flags hypercache at cache-write time so the Rust service can skip the separate PG query.

Python (cache write path):

• Extracts cohort IDs from active flag filters (groups only — super_groups and holdout_groups cannot contain cohort properties)
• BFS-loads referenced cohorts with transitive cohort-on-cohort dependency resolution (depth limit 20, cycle-safe)
• Serializes cohorts into the hypercache payload alongside flags and evaluation metadata
• Batch path collects all cohort IDs across teams and loads in a single BFS pass
• Adds post_save/post_delete signal on Cohort to invalidate the flags cache on definition changes (skips recalculation-only saves via _COHORT_RECALCULATION_FIELDS subset check)

Rust (cache read path):

• HypercacheFlagsWrapper gains optional cohorts field with #[serde(default)] for backwards compatibility
• FeatureFlagMatcher consumes preloaded cohorts via Option::take(), falls back to CohortCacheManager when absent
• EvaluationMetadata gains Default and PartialEq derives

Rolling deploy safety:

• Old Rust ignores the unknown cohorts key (no deny_unknown_fields)
• New Rust reading old cache entries defaults cohorts to None and falls back to PG

How did you test this code?

• 20 new Python tests across 4 test classes: cohort extraction, BFS loading (transitive deps, circular deps, depth limit), serialization (all 16 fields), signal handler filtering
• 2 new Rust tests for hypercache cohort deserialization
• Rust test helper updated to propagate cohorts through round-trip

Test plan

The automated unit tests cover extraction, BFS loading, serialization, and signal handler logic. This plan focuses on end-to-end scenarios that exercise the full Python → Redis → Rust pipeline.

Prerequisites

• PostHog dev environment running (hogli start)
• Rust feature-flags service running and connected to the same Redis and Postgres
• Test team created with cohorts (direct, transitive, circular, unreferenced) and flags referencing them

1. Python: cache write path

• _get_feature_flags_for_service returns flags, evaluation_metadata, and cohorts
• Only referenced cohorts included (unreferenced cohort excluded)
• Transitive deps loaded (flag → cohort B → cohort A loads both)
• Circular cohort dependencies (D ↔ E) terminate without infinite loop
• BFS depth limit enforced (chain of 25 cohorts, loader returned only 20)
• Batch path isolates cohorts per team (no cross-team leakage)
• Log messages include flag_count and cohort_count

2. Python: signal handler (cache invalidation)

• Recalculation-only saves (is_calculating) do NOT trigger cache invalidation
• Definition changes (name) DO trigger cache invalidation
• Mixed fields (count + name) DO trigger invalidation
• Cohort deletion triggers invalidation via post_delete
• Signal uses transaction.on_commit, not direct task call
• No-op when FLAGS_REDIS_URL is not configured

3. Rust: flag evaluation with preloaded cohorts

• Preloaded cohort path: /flags with cohorts in cache evaluates correctly (cohort flags: no_condition_match, circular: dependency_cycle_cohort, plain: condition_match)
• Fallback path: cache without cohorts key falls back to CohortCacheManager, identical results

4. Rolling deploy safety

• Old Rust (master) reads new cache (with cohorts): silently ignores unknown key, correct results
• New Rust (branch) reads old cache (without cohorts): falls back to PG, correct results

Publish to changelog?

no

[Copilot]

Pull request overview

This PR reduces /flags request latency and Postgres load in the Rust feature-flags service by preloading only the cohort definitions actually referenced by a team’s active feature flags (including transitive cohort-on-cohort dependencies) into the flags HyperCache payload written by Django.

Changes:

• Python: extracts referenced cohort IDs from active flag filters.groups, BFS-loads cohorts + transitive dependencies (depth-limited), serializes them into the flags hypercache payload, and adds cohort-change signals to invalidate the flags cache.
• Rust: extends the hypercache wrapper to optionally deserialize cohorts (backwards compatible), threads cohorts through fetch/filter flows, and consumes preloaded cohorts in FeatureFlagMatcher (falling back to CohortCacheManager when absent).
• Tests: adds/updates Python unit tests for extraction/loading/serialization + signal filtering, and Rust tests for cohort deserialization and round-trip helpers.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
rust/feature-flags/src/utils/test_graph_utils.rs	Updates test builders to include the new cohorts field in the wrapper/list structures.
rust/feature-flags/src/handler/tests.rs	Updates handler tests to include cohorts: None in constructed wrapper/list values.
rust/feature-flags/src/handler/flags.rs	Plumbs cohorts through fetch_and_filter so downstream evaluation can use it.
rust/feature-flags/src/handler/billing.rs	Updates billing tests for the expanded FeatureFlagList shape.
rust/feature-flags/src/flags/test_helpers.rs	Ensures redis/hypercache test helpers round-trip cohorts alongside flags + metadata.
rust/feature-flags/src/flags/test_flag_matching.rs	Updates flag matching tests to include cohorts: None.
rust/feature-flags/src/flags/flag_service.rs	Updates cache parsing to return (flags, evaluation_metadata, cohorts) and stores cohorts in FlagResult.
rust/feature-flags/src/flags/flag_models.rs	Adds optional cohorts to hypercache wrapper and adds cohorts to runtime-only FeatureFlagList.
rust/feature-flags/src/flags/flag_matching.rs	Consumes preloaded cohorts (via Option::take) in prepare_flag_evaluation_state, with PG fallback.
rust/feature-flags/src/flags/feature_flag_list.rs	Adds cohort deserialization support and tests for presence/absence of cohorts.
posthog/models/feature_flag/test/test_flags_cache.py	Adds tests for cohort extraction, BFS loading (deps/cycles/depth), serialization, and cache invalidation signal behavior.
posthog/models/feature_flag/flags_cache.py	Implements cohort extraction/BFS loading/serialization into the hypercache payload + cohort-change invalidation signals.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[greptile-apps]

Comments Outside Diff (2)

1. posthog/models/feature_flag/test/test_flags_cache.py, line 742-749 (link)

   Standalone test should be folded into the parameterized sibling

   test_skips_cohort_type_only_save tests the same skip-on-recalculation-field logic as the existing test_skips_recalculation_only_save parameterized test, but as a one-off method. This breaks the "prefer parameterised tests" convention used throughout the test suite. The cohort_type case should simply be another entry in the @parameterized.expand list above:

   @parameterized.expand(
       [
           ("is_calculating", "is_calculating", True),
           ("count", "count", 100),
           ("version", "version", 2),
           ("cohort_type", "cohort_type", "behavioral"),
       ]
   )
   @patch("django.db.transaction.on_commit", lambda fn: fn())
   @patch("posthog.tasks.feature_flags.update_team_service_flags_cache")
   def test_skips_recalculation_only_save(self, _name, field, value, mock_task):
       ...

   The standalone test_skips_cohort_type_only_save method can then be removed.

   Context Used: Do not attempt to comment on incorrect alphabetica... (source)

   Prompt To Fix With AI

   This is a comment left during a code review.
   Path: posthog/models/feature_flag/test/test_flags_cache.py
   Line: 742-749
   
   Comment:
   **Standalone test should be folded into the parameterized sibling**
   
   `test_skips_cohort_type_only_save` tests the same skip-on-recalculation-field logic as the existing `test_skips_recalculation_only_save` parameterized test, but as a one-off method. This breaks the "prefer parameterised tests" convention used throughout the test suite. The `cohort_type` case should simply be another entry in the `@parameterized.expand` list above:
   
   ```python
   @parameterized.expand(
       [
           ("is_calculating", "is_calculating", True),
           ("count", "count", 100),
           ("version", "version", 2),
           ("cohort_type", "cohort_type", "behavioral"),
       ]
   )
   @patch("django.db.transaction.on_commit", lambda fn: fn())
   @patch("posthog.tasks.feature_flags.update_team_service_flags_cache")
   def test_skips_recalculation_only_save(self, _name, field, value, mock_task):
       ...
   ```
   
   The standalone `test_skips_cohort_type_only_save` method can then be removed.
   
   **Context Used:** Do not attempt to comment on incorrect alphabetica... ([source](https://app.greptile.com/review/custom-context?memory=instruction-0))
   
   How can I resolve this? If you propose a fix, please make it concise.

   _{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

2. posthog/models/feature_flag/test/test_flags_cache.py, line 665-668 (link)

   Overly loose boundary assertion for depth-limit test

   The BFS algorithm stops at exactly 20 cohorts (one per iteration before hitting _MAX_COHORT_DEPENDENCY_DEPTH). The current assertion <= 20 would still pass if only, say, 5 cohorts were loaded due to an off-by-one regression. A precise equality check pinpoints the expected cut-off behaviour and makes the test a real regression guard:

   Prompt To Fix With AI

   This is a comment left during a code review.
   Path: posthog/models/feature_flag/test/test_flags_cache.py
   Line: 665-668
   
   Comment:
   **Overly loose boundary assertion for depth-limit test**
   
   The BFS algorithm stops at exactly 20 cohorts (one per iteration before hitting `_MAX_COHORT_DEPENDENCY_DEPTH`). The current assertion `<= 20` would still pass if only, say, 5 cohorts were loaded due to an off-by-one regression. A precise equality check pinpoints the expected cut-off behaviour and makes the test a real regression guard:
   
   
   
   How can I resolve this? If you propose a fix, please make it concise.

Prompt To Fix All With AI

This is a comment left during a code review.
Path: posthog/models/feature_flag/flags_cache.py
Line: 99-118

Comment:
**`groups` silently included without explanation**

The `NOTE` comment documents why `cohort_type` is in this set, but `groups` is included without any explanation. `groups` is the *legacy* cohort-condition field (deprecated in favour of `filters`) and it is updated during `calculate_people_ch()` via `update_fields=["...", "groups"]` (see `posthog/models/cohort/cohort.py:347`). Without a note, a future reader may conclude this is a bug (definition field incorrectly suppressing cache invalidation).

```suggestion
_COHORT_RECALCULATION_FIELDS = frozenset(
    [
        "count",
        "version",
        "pending_version",
        "is_calculating",
        "last_calculation",
        "last_calculation_duration_ms",
        "errors_calculating",
        "last_error_at",
        # NOTE: `groups` is the legacy cohort-condition field (deprecated in favour of
        # `filters`).  calculate_people_ch() always saves it in update_fields even when
        # unchanged (see cohort.py:347).  Real definition changes go through a full save
        # (update_fields=None), so they still trigger invalidation.
        "groups",
        "cohort_type",
    ]
)
```

How can I resolve this? If you propose a fix, please make it concise.

---

This is a comment left during a code review.
Path: posthog/models/feature_flag/test/test_flags_cache.py
Line: 742-749

Comment:
**Standalone test should be folded into the parameterized sibling**

`test_skips_cohort_type_only_save` tests the same skip-on-recalculation-field logic as the existing `test_skips_recalculation_only_save` parameterized test, but as a one-off method. This breaks the "prefer parameterised tests" convention used throughout the test suite. The `cohort_type` case should simply be another entry in the `@parameterized.expand` list above:

```python
@parameterized.expand(
    [
        ("is_calculating", "is_calculating", True),
        ("count", "count", 100),
        ("version", "version", 2),
        ("cohort_type", "cohort_type", "behavioral"),
    ]
)
@patch("django.db.transaction.on_commit", lambda fn: fn())
@patch("posthog.tasks.feature_flags.update_team_service_flags_cache")
def test_skips_recalculation_only_save(self, _name, field, value, mock_task):
    ...
```

The standalone `test_skips_cohort_type_only_save` method can then be removed.

**Context Used:** Do not attempt to comment on incorrect alphabetica... ([source](https://app.greptile.com/review/custom-context?memory=instruction-0))

How can I resolve this? If you propose a fix, please make it concise.

---

This is a comment left during a code review.
Path: posthog/models/feature_flag/test/test_flags_cache.py
Line: 665-668

Comment:
**Overly loose boundary assertion for depth-limit test**

The BFS algorithm stops at exactly 20 cohorts (one per iteration before hitting `_MAX_COHORT_DEPENDENCY_DEPTH`). The current assertion `<= 20` would still pass if only, say, 5 cohorts were loaded due to an off-by-one regression. A precise equality check pinpoints the expected cut-off behaviour and makes the test a real regression guard:

```suggestion
        assert len(result) == 20  # BFS loads exactly _MAX_COHORT_DEPENDENCY_DEPTH cohorts
```

How can I resolve this? If you propose a fix, please make it concise.

_{Reviews (1): Last reviewed commit: "fix(flags): improve cohort preloading ob..." | Re-trigger Greptile}

[matheus-vb]

🚀