fix(cohorts): guard cohort recalculation Redis calls

[posthog]

Problem

• A transient Redis error during a cohort recalculation could leave the cohort stuck with is_calculating=True, which blocks recalculation until the flag is cleared and feeds stale membership to flag targeting and insights.
• _on_cohort_changed does a bare cache.get. Under autocommit its on_commit callback runs synchronously inside Cohort.save(), so a Redis error escapes the save.
• The recovery helper save_recovery_bookkeeping caught only InterfaceError and OperationalError, so that Redis error passed through and masked the real calculation error. That is the exact outcome the helper exists to prevent.
• A second Redis call, in run_cohort_query, could discard an already-computed recalculation. The result returns before best-effort stats bookkeeping revokes and reschedules a Celery task, and an unguarded Redis error there threw the result away.

Changes

• A cohort no longer gets stuck calculating after a Redis blip on the recalculation path.
• _on_cohort_changed now wraps its cache work in try/except, logs failed_to_refresh_cohort_dependency_cache, and continues. The dependency caches rebuild on the next change or warm pass.
• run_cohort_query now wraps the delayed-task revoke and reschedule. A Redis error logs cohort_stats_collection_scheduling_failed and keeps the computed result; the delayed task still runs after its original countdown.
• save_recovery_bookkeeping now catches any exception, not only Postgres connection errors, so a recovery write never masks the real failure. This is the defense-in-depth layer behind the _on_cohort_changed fix.

How did you test this code?

• test_cohort_save_survives_redis_error_in_dependency_refresh: a non-recalculation save fires _on_cohort_changed synchronously with cache.get raising; guards that the Redis error no longer escapes Cohort.save().
• Extended test_original_error_surfaces_when_reset_calculating_save_fails with a redis_error case; guards that the widened except still lets the real error propagate for a non-connection failure.
• Ran the two dependency and recovery test classes locally against Postgres, ClickHouse, and Redis: 4 passed.
• Not run: the full backend suite, left to CI.

Automatic notifications

• Publish to changelog?

Docs update

None. Internal recovery behavior, no user-facing docs.

🤖 Agent context

Autonomy: Fully autonomous

• Authored by PostHog Desktop (Claude Opus 4.8) from an inbox report.
• Skills invoked: /writing-tests, /writing-pr-descriptions, /writing-simplified-technical-english.
• The root-cause fix is the guard in _on_cohort_changed. Widening save_recovery_bookkeeping is a second layer that catches any future unguarded Redis call on the same save path.
• No customer material informed this change. The report identified the two code paths and their error signatures; the fix and tests were written against the code alone.

---

Created with PostHog Desktop from this inbox report.

[posthog]

🦔 PostHog Review reviewed this pull request

Found 0 must fix, 2 should fix, 1 consider.

Published 3 findings (view the review).

Resolved comments: 1 fixed, 1 declined, 1 already settled

[trunk-io]

Merging to master in this repository is managed by Trunk.

• To merge this pull request, check the box to the left or comment /trunk merge below.

After your PR is submitted to the merge queue, this comment will be automatically updated with its status. If the PR fails, failure details will also be posted here

[github-actions]

🤖 CI report

⚠️ Trunk lane — backend Python lane

This PR is assigned to the backend Python lane. It runs backend Python tests and may merge in parallel with PRs in other lanes.

⚠️ Backend coverage — 88.0% of changed backend lines covered — 6 uncovered

🧪 Backend test coverage

Patch coverage — changed backend lines (products + core): ██████████████████░░ 88.0% (47 / 53)

File	Patch	Uncovered changed lines
products/cohorts/backend/models/util.py	50.0%	237–239, 242, 246–247

🤖 Agents: add a test covering the lines above, or note why under "How did you test this code?". Machine-readable gap list: the patch-coverage artifact on this run (gh run download 33301196531 -n patch-coverage), or the coverage-data block at the end of this comment.

Per-product line coverage (touched products)

Product	Coverage	Lines
platform_features	██░░░░░░░░░░░░░░░░░░ 12.1%	7 / 58
warehouse_sources_queue	██████░░░░░░░░░░░░░░ 29.4%	91 / 309
demo	███████████░░░░░░░░░ 56.9%	1,518 / 2,666
data_tools	█████████████░░░░░░░ 66.2%	90 / 136
ai_gateway	███████████████░░░░░ 75.0%	9 / 12
batch_exports	████████████████░░░░ 79.8%	19,636 / 24,593
apm	█████████████████░░░ 83.3%	1,234 / 1,481
cdp	█████████████████░░░ 85.5%	4,300 / 5,028
actions	█████████████████░░░ 86.6%	716 / 827
wizard	█████████████████░░░ 86.7%	1,090 / 1,257
data_modeling	█████████████████░░░ 86.7%	9,654 / 11,133
notebooks	██████████████████░░ 87.8%	9,368 / 10,671
product_tours	██████████████████░░ 87.9%	1,303 / 1,482
cohorts	██████████████████░░ 88.5%	7,410 / 8,377
data_warehouse	██████████████████░░ 88.6%	12,768 / 14,411
business_knowledge	██████████████████░░ 89.1%	4,440 / 4,984
dashboards	██████████████████░░ 89.2%	6,387 / 7,157
engineering_analytics	██████████████████░░ 89.4%	7,614 / 8,516
tasks	██████████████████░░ 89.9%	61,753 / 68,718
streamlit_apps	██████████████████░░ 90.2%	2,500 / 2,773
exports	██████████████████░░ 90.2%	8,509 / 9,434
links	██████████████████░░ 90.6%	183 / 202
managed_warehouse	██████████████████░░ 90.7%	8,563 / 9,436
visual_review	██████████████████░░ 90.9%	6,880 / 7,567
conversations	██████████████████░░ 90.9%	20,694 / 22,760
mcp_analytics	██████████████████░░ 91.0%	4,231 / 4,652
alerts	██████████████████░░ 91.2%	5,216 / 5,719
error_tracking	██████████████████░░ 91.6%	12,601 / 13,758
notifications	██████████████████░░ 92.0%	1,075 / 1,168
canvas	██████████████████░░ 92.1%	4,693 / 5,096
early_access_features	███████████████████░ 92.5%	1,323 / 1,430
managed_migrations	███████████████████░ 92.6%	1,556 / 1,681
stamphog	███████████████████░ 92.6%	5,510 / 5,951
slack_app	███████████████████░ 92.6%	12,202 / 13,173
web_analytics	███████████████████░ 93.0%	18,002 / 19,359
context_layer	███████████████████░ 93.1%	2,313 / 2,485
surveys	███████████████████░ 93.3%	6,037 / 6,468
mcp_store	███████████████████░ 93.4%	7,549 / 8,084
reminders	███████████████████░ 93.4%	468 / 501
posthog_ai	███████████████████░ 93.5%	1,372 / 1,468
marketing_analytics	███████████████████░ 93.5%	17,322 / 18,526
ai_observability	███████████████████░ 93.7%	17,981 / 19,187
approvals	███████████████████░ 93.7%	3,492 / 3,725
product_analytics	███████████████████░ 93.8%	7,654 / 8,164
billing_alerts	███████████████████░ 93.9%	2,078 / 2,212
legal_documents	███████████████████░ 94.1%	1,970 / 2,093
workflows	███████████████████░ 94.3%	9,759 / 10,354
review_hog	███████████████████░ 94.3%	9,973 / 10,574
tracing	███████████████████░ 94.4%	2,678 / 2,837
endpoints	███████████████████░ 94.4%	8,920 / 9,448
customer_analytics	███████████████████░ 95.0%	17,542 / 18,464
annotations	███████████████████░ 95.2%	795 / 835
experiments	███████████████████░ 95.5%	29,265 / 30,654
replay_vision	███████████████████░ 95.6%	24,892 / 26,047
skills	███████████████████░ 95.6%	5,330 / 5,577
messaging	███████████████████░ 95.8%	3,591 / 3,750
logs	███████████████████░ 95.9%	13,457 / 14,039
growth	███████████████████░ 95.9%	7,531 / 7,855
feature_flags	███████████████████░ 96.1%	19,441 / 20,220
access_control	███████████████████░ 96.3%	5,458 / 5,669
revenue_analytics	███████████████████░ 96.3%	1,841 / 1,911
user_interviews	███████████████████░ 96.5%	2,639 / 2,735
data_quality	███████████████████░ 97.0%	3,685 / 3,799
warehouse_sources	███████████████████░ 97.2%	403,992 / 415,698
analytics_platform	████████████████████ 98.1%	2,286 / 2,330
data_catalog	████████████████████ 98.3%	3,234 / 3,289
metrics	████████████████████ 98.3%	3,453 / 3,511
pulse	████████████████████ 98.4%	2,017 / 2,049
field_notes	████████████████████ 99.4%	158 / 159

Report-only. Patch coverage = changed backend lines covered vs origin/master. Sorted lowest first.
Known gaps: lines covered only by Temporal tests show as uncovered; core line numbers may drift if master changed the same file.

[stamphog]

Contained, well-tested defensive error handling around Redis/Celery calls in the cohort recalculation recovery path — no data model, API, billing, auth, or dependency changes, and the diff matches the description with no unresolved concerns.

Gate mechanics and policy version

Gate		Result
prerequisites	✓	all clear
deny-list	✓	no deny categories matched
size	✓	70L, 2F substantive, 103L/4F incl. docs/generated/snapshots — within ceiling
tier	✓	T1-agent / T1c-medium (103L, 4F, single-area, fix)
stamphog 2.0.0b4		.stamphog/policy.yml @ a5826f0 · reviewed head a5826f0

[posthog]

PostHog Review alpha 🦔 If you find any issues helpful - please reply "valid", "invalid", etc., for evaluation purposes 🙏

[posthog]

PostHog Review

Found 2 should fix, 1 consider.

[trunk-io]

   

Failed Test	Failure Summary	Logs
test_csv_import_records_unmatched_distinct_ids	The database connection was closed, causing failures when attempting to access the database during tests.	Logs ↗︎
test_csv_import_records_unmatched_emails	The database connection was closed, causing failures when attempting to access the database during tests.	Logs ↗︎
test_calculating_cohort_clickhouse	The database connection was closed, causing failures when attempting to access the database during tests.	Logs ↗︎
test_properties_flat_includes_realtime_style_behavioral_leaf_0_multiple_times_suffixed_value	The database connection was closed, causing failures when attempting to access the database during tests.	Logs ↗︎

... and 23 more

View Full Report ↗︎ ⋅ Docs

New commits were pushed — dismissing the stamphog approval from an earlier head. This PR no longer qualifies for automatic review.

[haacked]

Nice fix. One blocking issue inline, plus a few suggestions.

[haacked]

blocking: If delayed_task.revoke() succeeds and collect_cohort_query_stats.apply_async then raises, the countdown task is already cancelled and nothing replaces it. That CohortCalculationHistory row never gets its queries payload, and collect_cohort_query_stats is the only thing that writes that field. The comment at util.py:236 says the delayed task still runs after its original countdown, which holds when the state read or the revoke fails, but not in this case.

No test covers this block. test_survives_broker_error_scheduling_initial_stats_task (test_util.py:1405) looks like it should, but its side_effect raises on every apply_async call. The first one at util.py:220 fails, so delayed_task stays None, and the if delayed_task and ... gate at util.py:233 skips this entire block. The test's own mock_apply.assert_called_once() is what proves the block never ran. Every other test has settings.TEST true, which skips the scheduling altogether.

Publishing before revoking makes the comment true in every ordering. A publish failure leaves the countdown task alone, and a revoke failure just means both tasks run, which collect_cohort_query_stats already handles by returning early when history.queries is populated.

             try:
-                if delayed_task.state in ["PENDING", "RECEIVED"]:
-                    delayed_task.revoke()  # Cancel the delayed task
-
-                # Run immediately since the query already completed
                 collect_cohort_query_stats.apply_async(
                     args=[cohort_tag, cohort_id, start_time.isoformat(), history.id, query],
                     countdown=COHORT_STATS_COLLECTION_DELAY_SECONDS,
                 )
+
+                # Publish before revoking so a failure here leaves the countdown task in place.
+                if delayed_task.state in ["PENDING", "RECEIVED"]:
+                    delayed_task.revoke()
             except Exception as error:

Add a test case where the first apply_async returns a task and the second raises. Set state to "PENDING" on that returned mock. A bare MagicMock fails the in ["PENDING", "RECEIVED"] check and leaves the revoke line uncovered.

[haacked]

suggestion: This case is named as a Redis error escaping the save's post_save signals, but the call it patches cannot produce one. _safe_reset_calculating_state (cohort.py:654) writes through a queryset .update(), and Django's QuerySet.update() fires no model signals.

The signal-firing call site is different: the bookkeeping self.save(...) at cohort.py:736 does fire signals and has no test.

Suggested change

	("redis_error", Exception("redis unavailable")),
	("non_connection_error", Exception("bookkeeping write failed")),

Fix the comment at test_util.py:1367 the same way.

[haacked]

suggestion: Both new guards log cohort_stats_collection_scheduling_failed with the same two fields, so a line in production does not say which one fired, and the two mean different things. This one fires before the query runs and means no stats task was ever scheduled, so that recalculation gets no stats at all. The one at util.py:247 fires after the result already exists and means the scheduled task was not pulled forward, so the stats still land, just after the full countdown.

Neither passes exc_info, so whoever reads the log gets a stringified error and no traceback. save_recovery_bookkeeping logs its own swallow with exc_info=True at util.py:176.

Suggested change

	logger.warning("cohort_stats_collection_scheduling_failed", cohort_id=cohort_id, error=str(error))
	logger.warning(
	"cohort_stats_initial_scheduling_failed", cohort_id=cohort_id, error=str(error), exc_info=True
	)

For the one at util.py:247, rename the event to cohort_stats_reschedule_failed. Alternatively, keep the shared event name and add a phase="initial" / phase="reschedule" field if you would rather have one name to query on.

[haacked]

suggestion: The name reads as an event handler that runs after something else invalidated the behavioral cache, but this function is the thing doing the invalidating. The _on_ prefix indicates a reaction to an event, like _on_cohort_changed which runs when a cohort changes. Someone tracing where the behavioral cache gets cleared will scroll past this one.

_safe_save_cohort_state (cohort.py:1255) is the nearby name for this exact shape, a wrapper that calls the real thing, logs, and swallows. Nothing outside this file references the new name, so the rename touches the definition and the two transaction.on_commit registrations at dependencies.py:473 and dependencies.py:591.

Suggested change

	def _on_cohort_behavioral_cache_invalidated(team_id: int) -> None:
	def _safe_invalidate_team_behavioral_cohort_cache(team_id: int) -> None:

[haacked]

suggestion: Widening the exception catch widens what gets retried, but retry only makes sense for connection failures. For dropped connections nothing was written, so re-running save_fn() is safe. For other exceptions like Redis or broker failures, the row is already committed by the time the receiver fails (Django sends post_save after the UPDATE). The retry then unnecessarily closes a healthy connection and re-runs the whole save, re-firing every post_save receiver on the cohort, including the still-unguarded update_team_flags_cache.delay(...) at local_evaluation.py:958.

The blast radius is small today, which is why this is a suggestion. CONN_MAX_AGE is 0, so the reconnect costs a single connection attempt, and the update_fields at cohort.py:736 omit filters, so no duplicate backfill run can be enqueued.

This isn't the approach you already declined. Every exception still gets swallowed and captured, so error tracking doesn't become less visible. The reconnect-and-retry logic narrows specifically to when it makes sense: failures that a reconnect can fix. It requires InterfaceError and OperationalError from django.db at util.py:11.

     try:
-        save_fn()
-    except Exception:
-        connections[DEFAULT_DB_ALIAS].close()  # next query opens a fresh connection
-        try:
+        try:
             save_fn()
-        except Exception as retry_error:
-            # A swallowed retry means the cohort is stuck with is_calculating=True and no bookkeeping
-            # recorded. Surface it to error tracking, matching how other swallowed cohort-calculation
-            # errors are captured, so it alerts rather than only living in structured logs.
-            logger.warning("cohort_recalc_recovery_save_failed", cohort_id=cohort_id, team_id=team_id, exc_info=True)
-            capture_exception(retry_error, additional_properties={"cohort_id": cohort_id, "team_id": team_id})
+        except (InterfaceError, OperationalError):
+            connections[DEFAULT_DB_ALIAS].close()  # next query opens a fresh connection
+            save_fn()
+    except Exception as error:
+        # A swallowed failure means the cohort is stuck with is_calculating=True and no bookkeeping
+        # recorded. Surface it to error tracking, matching how other swallowed cohort-calculation
+        # errors are captured, so it alerts rather than only living in structured logs.
+        logger.warning("cohort_recalc_recovery_save_failed", cohort_id=cohort_id, team_id=team_id, exc_info=True)
+        capture_exception(error, additional_properties={"cohort_id": cohort_id, "team_id": team_id})

Update the non-connection error case in test_util.py:1398: the reconnect no longer runs, so the assertions become assert mock_reset.call_count == 1 and close.assert_not_called(). The connection_error case keeps its current assertions.