Speed up the capture set list view

[mihow]

Summary

Two changes, both targeting SourceImageCollectionViewSet.list query cost on large M2M tables:

1. Denormalize the 3 source-image counts as cached columns on SourceImageCollection. Replaces 3 per-row Subquery(Count(...)) annotations with a plain column SELECT; O(1) read regardless of M2M size.
2. Introduce CachedCountField in ami/base/models.py — a marker subclass of IntegerField — and apply it to all 12 existing cached count columns project-wide. Discoverable via Model._meta.get_fields() + isinstance(f, CachedCountField) so future cross-cutting tasks (admin display, refresh jobs, list-endpoint defer) can enumerate them.

Stacked on #1300 (select_related on captures list). Base branch is perf/list-endpoint-subquery-rewrite; will retarget to main once #1300 merges.

Denormalization details

• 3 new columns on SourceImageCollection: source_images_count, source_images_with_detections_count, source_images_processed_count. Schema migration 0085_denormalize_sourceimagecollection_counts creates them and backfills via a single GROUP BY + FILTER scan of the M2M (computes all three counts in one pass). atomic=False so the UPDATE can chunk on prod-sized data.
• SourceImageCollection.get_source_image_counts() returns the 3 counts as a dict from a single aggregate query (no writes — pure compute, testable in isolation). update_calculated_fields(save=True) is the side-effecting wrapper that also writes the columns.
• Signal handlers in ami/main/signals.py:
  • m2m_changed on SourceImageCollection.images.through for post_add, post_remove, post_clear.
  • Detection post_save / post_delete recompute for any collection containing the affected source_image.
• pipeline.save_results() explicitly recomputes for affected collections after its bulk_create batch, since bulk_create skips signals.
• Drops the 3 with_*_count() annotation methods on SourceImageCollectionQuerySet and removes them from the viewset chain.

CachedCountField details

• New class in ami/base/models.py — class CachedCountField(models.IntegerField). Pure marker; no field-level behavior change.
• Applied to: Deployment.{events,occurrences,captures,detections,taxa}_count, Event.{captures,detections,occurrences}_count, SourceImage.detections_count, SourceImageCollection.{source_images,source_images_with_detections,source_images_processed}_count.
• Migration 0086_use_cached_count_field is recorded but the generated SQL is (no-op) for every AlterField — the column type is unchanged. The migration only updates Django's recorded model state so future operations see the right field class.
• Kept separate from the 0085 backfill migration so the schema-and-data migration is auditable on its own.

Why denormalize rather than the subquery rewrite from the earlier draft of #1300

The collections list UI does not paginate. The paginator COUNT(*) win from rewriting the annotations as correlated subqueries (the original #1300 approach) is invisible to UX. The remaining cost was 3 per-row count subqueries against the M2M + main_detection; on a collection holding 500k+ images, each request still spent hundreds of ms per subquery.

Bench data from the earlier draft also showed the subquery rewrite was slower than the original Count(images, filter=..., distinct=True) on small projects and roughly equal on the heaviest project. Cached columns are an unambiguous win across all project sizes.

Precedent: Deployment already had 5 cached counts (events_count, occurrences_count, captures_count, detections_count, taxa_count). The CachedCountField marker formalizes that pattern.

Recompute granularity tradeoff

Per-Detection signal recompute fires on every individual Detection.save() / Detection.delete(). For ML pipeline writes this is a non-issue: Detection.objects.bulk_create() skips signals, and pipeline.save_results() batches one recompute per job. The per-row signal handles non-bulk paths (manual workflows, retroactive ML, admin actions).

Tests

• New TestSourceImageCollectionCountsDenormalize (7 cases) — image add/remove, Detection create/delete, null-bbox processed-but-not-with-detections case, update_calculated_fields() recovery from drift, pure-compute get_source_image_counts() returns dict without writes. A tautological test_initial_counts_zero was dropped during review — it asserted Django's default=0, not our code.
• New TestSourceImageCollectionListQueryCount (3 cases) — page-size scaling, with_counts=true, ordering=-source_images_count (sort hits the cached column directly, no extra subquery).
• Full ami.main + ami.ml + ami.jobs suites pass on local.

Rollout

• 0085 creates columns with default=0 and backfills in one deploy. Backfill is the slow step on prod-sized data — measure on staging arbutus-2026 first.
• 0086 is no-op SQL.
• Signals fire on new writes immediately after deploy. The viewset switches from annotation to column reads in the same commit, so the deploy must include both the migration and the code.

Test plan

• Local: full ami.main + ami.ml + ami.jobs test suites
• Staging arbutus-2026: measure 0085 backfill duration on production-shape M2M
• Staging: NR check that /api/v2/captures/collections/ p99 drops vs production

Summary by CodeRabbit

Release Notes

• New Features

  • Collection counts (total images, images with detections, processed images) are now denormalized for improved query performance.
  • Cached counts across collections, deployments, events, and source images automatically stay in sync when related data changes.
  • Added methods to manually refresh cached counts when needed.

• Chores

  • Database migrations applied to add denormalized count fields and backfill existing data.

[netlify]

✅ Deploy Preview for antenna-preview canceled.

Name	Link
🔨 Latest commit	2533857
🔍 Latest deploy log	https://app.netlify.com/projects/antenna-preview/deploys/6a0cf516152163000881e4ca

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces a denormalized cached-count field infrastructure: a new CachedCountField marker type, a per-connection deduplicating recompute scheduler with transaction-aware commit hooks, database migrations to add/convert count columns, signal wiring to maintain counts on M2M and Detection changes, Celery task support for async recomputation, and comprehensive tests validating consistency and query performance.

Changes

Cached Count Denormalization and Async Recomputation

Layer / File(s)	Summary
CachedCountField marker and recompute scheduler ami/base/models.py, ami/base/cached_counts.py (new)	Defines CachedCountField as an IntegerField subclass marker; implements transaction-aware per-connection dedup scheduler that queues (model_label, pk) pairs and dispatches one Celery task per unique row on commit; adds update_cached_counts(run_async=True) methods to BaseModel and BaseQuerySet supporting sync (inline) and async (scheduled) recomputation modes.
Database schema, backfill, and field conversions ami/main/migrations/0085_*, ami/main/migrations/0086_*, ami/main/migrations/0087_*	Adds three denormalized count columns to SourceImageCollection with default=0; backfills them via idempotent non-transactional SQL UPDATE with Exists subqueries and bbox validity checks; converts existing count fields in Deployment, Event, SourceImage, and SourceImageCollection to use CachedCountField.
Model field conversions and SourceImageCollection computation ami/main/models.py	Changes count fields in Deployment (5 fields), Event (3 fields), and SourceImage (1 field) to CachedCountField; removes SourceImageCollectionQuerySet annotation helpers (with_source_images_*_count); adds three CachedCountField columns to SourceImageCollection; implements get_source_image_counts() aggregate query using Exists and update_calculated_fields(save=False|True) to recompute and optionally persist counts.
Signal handlers, tasks, and API updates ami/main/signals.py, ami/main/tasks.py, ami/main/api/views.py, ami/main/api/serializers.py, ami/ml/models/pipeline.py	Wires m2m_changed handler on SourceImageCollection.images and post_save/post_delete handlers on Detection to call update_cached_counts(); adds recompute_cached_counts_task Celery task; updates refresh_project_cached_counts docstring; handles bulk Detection creates in ML pipeline by refreshing affected collections; simplifies SourceImageCollectionViewSet.queryset to only prefetch jobs; marks denormalized counts read-only in serializer.
Tests for query counts and denormalization correctness ami/main/tests.py	Tightens TestSourceImageListQueryCount query-count assertions; adds TestSourceImageCollectionListQueryCount measuring cold query counts for captures/collections list across pagination limits with count ordering; adds TestSourceImageCollectionCountsDenormalize using TransactionTestCase to validate count consistency under M2M updates, Detection creates/deletes with bbox semantics, and explicit recomputation calls; verifies get_source_image_counts() returns computed values without persisting.
Design specification and follow-up planning docs/superpowers/specs/2026-05-14-cached-counts-update-method-design.md, docs/claude/planning/cached-counts-reconcile-followup.md	Documents the refactor goals, scheduler semantics, update_cached_counts API, call-site changes, risks, and migration path; includes follow-up planning for reconciliation restoration with safer task defaults, per-project periodic scheduling, and set-based diff strategy for detecting stale counts.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related issues

• RolnickLab/antenna#1253: Proposes denormalization and caching strategy for SourceImageCollection counts, which this PR directly implements with CachedCountField, denormalized columns, and async recomputation wiring.

Possibly related PRs

• RolnickLab/antenna#1300: Modifies SourceImageCollection queryset annotation helpers and list tests that rely on those annotated counts; this PR removes those helpers and shifts to denormalized cached columns with explicit recomputation.
• RolnickLab/antenna#1093: Adds with_source_images_processed_count() queryset helper and NULL_DETECTIONS_FILTER for the processed count metric; this PR replaces that annotation-based approach with a persisted CachedCountField column and recomputation logic.

Suggested labels

backend

Suggested reviewers

• annavik

Poem

🐰 A count denormalized with care,
Dedup'd and scheduled in the air,
Signals fire as data flows,
Tasks dispatch when commit knows.
Cached and consistent, fast and fair!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 44.68% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Speed up the capture set list view' accurately describes the main performance objective of the changeset—optimizing the SourceImageCollection list view by denormalizing cached counts and introducing CachedCountField.
Description check	✅ Passed	The PR description is comprehensive and well-structured, covering summary, detailed denormalization and CachedCountField details, rationale, test plan, and rollout notes. All required sections from the template are present and substantively filled.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch perf/sourceimagecollection-cached-counts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[mihow]

How can we indicate which fields are cached counts? field properties? prefix? I believe we have a cached_ prefix on other models, but then it makes the method replacements not 1:1. maybe thats okay. open to ideas.

[mihow]

Claude says: Checked — no cached_ prefix exists on any model field in the codebase. All denormalized count fields use the bare <thing>_count pattern: Deployment.events_count, Deployment.captures_count, Event.captures_count, SourceImage.detections_count, etc. The new fields match this convention.

I'd lean against adding a cached_ prefix just to this PR's fields — it would create a split where some counts are prefixed and most aren't, and method names (occurrences_count, taxa_count on the same model) wouldn't be 1:1 anymore as you noted. If we want the prefix project-wide, that's a separate refactor PR covering all the count fields.

Happy to do either — flag if you'd prefer the new fields prefixed and we adopt project-wide later, or leave matching the existing convention.

[mihow]

Is there a simple way to mark/tag/categorize/group model fields? To indicate these are cached counts that could be potentially stale or empty at any time? declare in a list on the model? add a model field property, etc?

[mihow]

Claude says: PR body updated. Two new pieces since the initial draft:

1. CachedCountField marker class added to ami/base/models.py and applied to all 12 existing cached-count columns project-wide (Deployment×5, Event×3, SourceImage×1, SourceImageCollection×3). Migration 0086_use_cached_count_field is recorded but the generated SQL is no-op for every AlterField — column type unchanged. Kept separate from the 0085 backfill migration on purpose.
2. Dropped a tautological test_initial_counts_zero (asserted Django's default=0, not our code) and rewrote a handful of em-dash fragment comments added during earlier review rounds.

@coderabbitai full review please. Base branch is perf/list-endpoint-subquery-rewrite (sister PR #1300, which strips down to a single select_related line). Will retarget to main once #1300 merges.

[coderabbitai]

@mihow Sure, I'll kick off a full review of this PR now, taking into account the updated PR body and the two new additions (the CachedCountField marker class with its no-op migration 0086, and the dropped tautological test).

✅ Actions performed

Full review triggered.

[mihow]

@mohamedelabbas1996 I think we discussed a field type for cached counts last year. it finally happened!

[netlify]

✅ Deploy Preview for antenna-ssec canceled.

Name	Link
🔨 Latest commit	2533857
🔍 Latest deploy log	https://app.netlify.com/projects/antenna-ssec/deploys/6a0cf5167a20ec000867f7c6

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@ami/main/signals.py`:
- Around line 210-211: The signal currently calls
instance.update_calculated_fields(save=True) synchronously for actions
("post_add", "post_remove", "post_clear"), which can persist stale counts under
concurrent transactions; change these immediate calls to run after the DB
transaction commits by wrapping the call in django.db.transaction.on_commit
(e.g., transaction.on_commit(lambda:
instance.update_calculated_fields(save=True))). Apply the same replacement for
the other occurrence around lines referencing the same actions so recomputation
is deferred until successful commit.

In `@ami/main/tests.py`:
- Line 3386: Test fixtures are using normalized [0–1] bbox values but this repo
expects pixel-space coordinates; update the Detection/BoundingBox fixtures
(e.g., the Detection instances setting bbox=[0.1,0.1,0.2,0.2] and the other
occurrences at the mentioned sites) to use absolute pixel coordinates consistent
with the image dimensions used in these tests. Locate the test rows creating
Detection objects (look for Detection(..., bbox=...) and any direct
BoundingBox(...) usages) and replace the normalized floats with the
corresponding pixel values (or explicit integers) so the fixtures represent true
pixel-space bboxes across the cases mentioned (lines ~3386, ~3398, ~3424,
~3443).

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 1e0f8a9e-539a-4e36-8acc-5da8d6ba207a

📥 Commits

Reviewing files that changed from the base of the PR and between 0dedd4c and 2725599.

📒 Files selected for processing (8)

• ami/base/models.py
• ami/main/api/views.py
• ami/main/migrations/0085_denormalize_sourceimagecollection_counts.py
• ami/main/migrations/0086_use_cached_count_field.py
• ami/main/models.py
• ami/main/signals.py
• ami/main/tests.py
• ami/ml/models/pipeline.py

[Copilot]

Pull request overview

This PR speeds up capture-set list queries by replacing per-row source-image count annotations with denormalized cached columns and adding infrastructure to discover/reconcile cached count drift.

Changes:

• Adds cached count fields and migrations for SourceImageCollection, plus CachedCountField marker usage.
• Updates list/query logic, signals, tasks, and pipeline bulk-save handling to refresh cached collection counts.
• Adds integrity-check helpers and regression/query-count tests.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
ami/base/models.py	Adds CachedCountField marker class.
ami/main/models.py	Converts cached count fields and adds SourceImageCollection count recomputation.
ami/main/api/views.py	Removes source-image count annotations from collection queryset.
ami/main/signals.py	Adds M2M/detection signal handlers for collection count refresh.
ami/main/tasks.py	Adds collection refresh and cached-count reconciliation Celery tasks.
ami/main/checks/cached_counts.py	Adds cached-count discovery, drift detection, and reconciliation helpers.
ami/ml/models/pipeline.py	Refreshes affected collection counts after bulk-created detections.
ami/main/tests.py	Adds denormalized count, integrity-check, and list query-count tests.
ami/main/migrations/0085_denormalize_sourceimagecollection_counts.py	Adds/backfills SourceImageCollection cached count columns.
ami/main/migrations/0086_use_cached_count_field.py	Updates migration state to use CachedCountField.

Comments suppressed due to low confidence (3)

ami/main/signals.py:268

• This only refreshes collections for the detection's current source_image_id. Because Detection.source_image is writable through the serializer/API, moving a detection to another capture leaves collections containing the old source image with stale processed/with-detections counts; the old source image id needs to be captured before save and refreshed too when it changes.

@receiver(post_save, sender=Detection)
@receiver(post_delete, sender=Detection)
def update_collection_counts_on_detection_change(sender, instance, **kwargs):
    """Schedule a collection-counts refresh for every collection containing the affected SourceImage.

    The dedup + on_commit indirection means even tight per-row Detection write
    loops fan out to at most one task per affected collection. ``bulk_create``
    / ``bulk_update`` skip signals entirely — those rely on the periodic
    reconciliation task to repair drift.
    """
    if not instance.source_image_id:
        return
    _schedule_collection_refresh_for_source_image(instance.source_image_id)

ami/main/signals.py:253

• This does not cover SourceImage deletion: Django removes auto-created M2M through rows during cascade delete without sending m2m_changed, and the detection post_delete path can no longer discover the collections once the image membership is gone. Deleting captures will therefore leave source_images_count (and derived processed/detection counts) stale unless affected collection ids are captured before the image is deleted.

@receiver(m2m_changed, sender=SourceImageCollection.images.through)
def update_collection_counts_on_m2m(sender, instance, action, **kwargs):
    """Recompute denormalized counts when images are added to or removed from a collection."""
    if action in ("post_add", "post_remove", "post_clear"):
        collection_pk = instance.pk
        transaction.on_commit(lambda: refresh_collection_cached_counts.delay(collection_pk))

ami/main/signals.py:253

• The handler assumes instance is always a SourceImageCollection, but m2m_changed also fires for the reverse manager (source_image.collections.add/remove(...)) with instance as a SourceImage and the collection ids in pk_set. In that path this schedules a refresh for the source-image id instead of the affected collection ids, leaving the real collection counts stale.

@receiver(m2m_changed, sender=SourceImageCollection.images.through)
def update_collection_counts_on_m2m(sender, instance, action, **kwargs):
    """Recompute denormalized counts when images are added to or removed from a collection."""
    if action in ("post_add", "post_remove", "post_clear"):
        collection_pk = instance.pk
        transaction.on_commit(lambda: refresh_collection_cached_counts.delay(collection_pk))

---

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

[mihow]

Great catch

[mihow]

Claude says: Self-suggestion for follow-up PR (or fold into this one): simplify the cached-counts plumbing introduced here. Current shape works but has three per-model pieces (dedup state attr + flush helper + per-model Celery task). Next cached-count addition repeats all three.

Proposed consolidation

One free helper + queryset/instance methods on BaseModel / BaseQuerySet. Fits the established run_async= flag pattern (Deployment.save(regroup_async=), Job.retry(async_task=), process_single_source_image(run_async=)).

# ami/base/cached_counts.py
_PENDING_ATTR = "_pending_cached_count_recomputes"

def _schedule_recompute(model, pk):
    """Per-connection dedup + transaction.on_commit dispatch.

    Outside an atomic block, on_commit fires synchronously at registration
    time, so we set the pending set BEFORE registering the flush callback.
    """
    pending = getattr(connection, _PENDING_ATTR, None)
    is_new = pending is None
    if is_new:
        pending = set()
        setattr(connection, _PENDING_ATTR, pending)
    pending.add((model._meta.label, pk))
    if is_new:
        transaction.on_commit(_flush)

def _flush():
    pending = getattr(connection, _PENDING_ATTR, set())
    try: delattr(connection, _PENDING_ATTR)
    except AttributeError: pass
    for label, pk in pending:
        recompute_cached_counts_task.delay(label, pk)

@shared_task(ignore_result=True)
def recompute_cached_counts_task(model_label, pk):
    model = apps.get_model(model_label)
    try: instance = model.objects.get(pk=pk)
    except model.DoesNotExist: return
    instance.update_calculated_fields(save=True)

# ami/base/models.py
class BaseModel(models.Model):
    def update_cached_counts(self, run_async: bool = True) -> None:
        if run_async:
            _schedule_recompute(type(self), self.pk)
            return
        self.update_calculated_fields(save=True)

class BaseQuerySet(QuerySet):
    def update_cached_counts(self, run_async: bool = True) -> None:
        for pk in self.values_list("pk", flat=True):
            if run_async:
                _schedule_recompute(self.model, pk)
            else:
                self.model.objects.get(pk=pk).update_calculated_fields(save=True)

Call sites collapse

Pipeline save_results:

SourceImageCollection.objects.filter(
    images__id__in=source_image_ids
).update_cached_counts(run_async=False)

Signals:

@receiver(post_save, sender=Detection)
@receiver(post_delete, sender=Detection)
def update_collection_counts_on_detection_change(sender, instance, **kwargs):
    if not instance.source_image_id:
        return
    SourceImageCollection.objects.filter(
        images__id=instance.source_image_id
    ).update_cached_counts()  # async by default

Reconcile loop in ami/main/checks/cached_counts.py swaps update_calculated_fields(save=True) → update_cached_counts(run_async=False).

What goes away

• _PENDING_SOURCE_IMAGE_IDS_ATTR + _flush_pending_collection_refreshes + _schedule_collection_refresh_for_source_image in signals.py (per-model state attrs disappear)
• refresh_collection_cached_counts(collection_id) task in ami/main/tasks.py (replaced by generic recompute_cached_counts_task(model_label, pk))
• Per-model signal handlers that need to know which Celery task to call

What stays

• CachedCountField marker (still drives reconcile discovery)
• Per-model update_calculated_fields(save=True) (recompute logic lives where the model lives)
• Periodic reconcile_cached_counts_task (safety net for bulk-write paths that bypass signals)
• .filter(pk=).update() vs .save() shape divergence in update_calculated_fields — documented inline, leaving alone

Next cached-count addition

Three things instead of six:

1. Field: foo_count = CachedCountField()
2. Recompute body in update_calculated_fields(save)
3. One signal handler per write source, each ending in .update_cached_counts()

Reconcile + drift detection are free.

What I'd NOT do

• Declarative invalidate_on=[...] on the field — hides signal wiring, adds magic
• Drop signals entirely + rely on TTL reconciliation — breaks interactive freshness (just-uploaded image not showing in collection)
• Lazy recompute-on-read — read path doing writes is its own footgun

Happy to land as a follow-up PR (cleaner review scope) or fold into this branch if reviewers prefer the unified pattern in one place.

[mihow]

Claude says: Landed in 9ac2ce17. update_cached_counts(run_async=True) lives on BaseModel and BaseQuerySet; backing module is ami/base/cached_counts.py. Detection signal handler is one line. refresh_collection_cached_counts task is gone; refresh_project_cached_counts body unchanged from this PR (kept the bulk-subquery UPDATE for SourceImage.detections_count).

Test isolation note worth flagging: the per-connection dedup set survives TestCase rollback because connection attribute state is not reset between tests, but the registered on_commit callbacks are discarded. Fix is to register on_commit on every schedule_recompute call and make _flush_pending_recomputes idempotent on the drain side (first call drains, subsequent calls see an empty set and no-op). That keeps the dedup benefit (10k Detection writes still collapse to one task per collection) without the rollback leak.

Reconcile loop in ami/main/checks/cached_counts.py switched to update_cached_counts(run_async=False) for stylistic unity — same behavior since reconcile is already in a worker.

Design doc + trade-offs in docs/superpowers/specs/2026-05-14-cached-counts-update-method-design.md (committed in 451c8be).

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

ami/base/models.py (1)

74-80: ⚡ Quick win

Optimize the sync path to avoid N+1 queries.

The sync path (lines 79-80) loads instances one at a time with self.model.objects.get(pk=pk), creating N+1 queries. Since the async path only needs PKs but the sync path needs full instances, branch on run_async and use self.iterator() for the sync path to load instances efficiently in a single query.

⚡ Proposed refactor

-        model_label = self.model._meta.label
-        for pk in self.values_list("pk", flat=True):
-            if run_async:
+        if run_async:
+            model_label = self.model._meta.label
+            for pk in self.values_list("pk", flat=True):
                 schedule_recompute(model_label, pk)
-            else:
-                instance = self.model.objects.get(pk=pk)
+        else:
+            for instance in self.iterator():
                 instance.update_calculated_fields(save=True)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ami/base/models.py` around lines 74 - 80, The current sync path does N+1
queries by calling self.model.objects.get(pk=pk) for each pk; change the logic
to branch on run_async: if run_async keep the existing loop over
self.values_list("pk", flat=True) and call schedule_recompute(model_label, pk),
but for the sync path iterate the queryset itself using self.iterator() (or
another queryset iterator) to load instances efficiently and call
instance.update_calculated_fields(save=True) on each instance; reference
model_label, schedule_recompute, run_async, self.values_list("pk", flat=True),
self.iterator(), and update_calculated_fields when applying the change.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@ami/main/tasks.py`:
- Line 11: The task function recompute_cached_counts_task currently types pk as
int but should accept any primary key type to match schedule_recompute in
ami/base/cached_counts.py; change the signature to use pk: Any and add "from
typing import Any" at the top of ami/main/tasks.py, then pass pk through to
schedule_recompute unchanged so type checkers no longer flag the mismatch (refer
to recompute_cached_counts_task and schedule_recompute).

---

Nitpick comments:
In `@ami/base/models.py`:
- Around line 74-80: The current sync path does N+1 queries by calling
self.model.objects.get(pk=pk) for each pk; change the logic to branch on
run_async: if run_async keep the existing loop over self.values_list("pk",
flat=True) and call schedule_recompute(model_label, pk), but for the sync path
iterate the queryset itself using self.iterator() (or another queryset iterator)
to load instances efficiently and call
instance.update_calculated_fields(save=True) on each instance; reference
model_label, schedule_recompute, run_async, self.values_list("pk", flat=True),
self.iterator(), and update_calculated_fields when applying the change.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: e47d60b5-368e-40bd-a816-14549cda47b0

📥 Commits

Reviewing files that changed from the base of the PR and between 2725599 and 9ac2ce1.

📒 Files selected for processing (8)

• ami/base/cached_counts.py
• ami/base/models.py
• ami/main/checks/cached_counts.py
• ami/main/models.py
• ami/main/signals.py
• ami/main/tasks.py
• ami/main/tests.py
• docs/superpowers/specs/2026-05-14-cached-counts-update-method-design.md

✅ Files skipped from review due to trivial changes (1)

• docs/superpowers/specs/2026-05-14-cached-counts-update-method-design.md

🚧 Files skipped from review as they are similar to previous changes (2)

• ami/main/signals.py
• ami/main/models.py

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Widen the pk type annotation to support non-integer primary keys.

The signature declares pk: int, but schedule_recompute in ami/base/cached_counts.py accepts pk: Any (line 27), and Django models can use UUIDField or other types as primary keys. While Celery serialization handles this at runtime, type checkers will flag the mismatch.

🔧 Proposed fix

-def recompute_cached_counts_task(model_label: str, pk: int) -> None:
+def recompute_cached_counts_task(model_label: str, pk: Any) -> None:

Add the import at the top if not already present:

+from typing import Any

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	def recompute_cached_counts_task(model_label: str, pk: int) -> None:
	from typing import Any
	def recompute_cached_counts_task(model_label: str, pk: Any) -> None:

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ami/main/tasks.py` at line 11, The task function recompute_cached_counts_task
currently types pk as int but should accept any primary key type to match
schedule_recompute in ami/base/cached_counts.py; change the signature to use pk:
Any and add "from typing import Any" at the top of ami/main/tasks.py, then pass
pk through to schedule_recompute unchanged so type checkers no longer flag the
mismatch (refer to recompute_cached_counts_task and schedule_recompute).

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (3)

ami/main/tests.py (3)

3444-3470: ⚡ Quick win

Assert all cached columns stay untouched in the no-write test.

This currently proves get_source_image_counts() leaves source_images_count alone, but it would still pass if the method accidentally persisted source_images_processed_count or source_images_with_detections_count.

Suggested assertion expansion

         # Confirm the DB row was not updated.
         self.collection.refresh_from_db()
         self.assertEqual(self.collection.source_images_count, 0)
+        self.assertEqual(self.collection.source_images_processed_count, 0)
+        self.assertEqual(self.collection.source_images_with_detections_count, 0)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ami/main/tests.py` around lines 3444 - 3470, The test currently only asserts
that source_images_count is unchanged after calling get_source_image_counts();
extend it to also assert that the other cached DB columns
(source_images_processed_count and source_images_with_detections_count) remain
untouched in the no-write case: after calling
self.collection.get_source_image_counts() and refreshing from DB via
self.collection.refresh_from_db(), add assertions that
self.collection.source_images_processed_count == 0 and
self.collection.source_images_with_detections_count == 0 so the test fails if
get_source_image_counts() erroneously persists those fields on
SourceImageCollection.

---

3376-3385: ⚡ Quick win

Add a clear() regression case.

This suite exercises add/remove, but the PR also relies on the images.clear() maintenance path. A small test that clears the collection and expects all three cached count columns to drop back to zero would close that last signal path.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ami/main/tests.py` around lines 3376 - 3385, Add a new test method (e.g.,
test_count_resets_on_image_clear) that sets up the collection with
self.collection.images.set(self.images), calls self.collection.images.clear(),
refreshes state with self._refresh(), and then asserts that the three cached
count columns on the collection (including source_images_count) have dropped to
zero; reference collection.images.clear() for the action and
collection.source_images_count (and the other two cached count attributes on the
collection) for the assertions.

---

3295-3305: ⚡ Quick win

Make the ordering path observable.

Every seeded collection gets the same image set here, so ordering=-source_images_count can be ignored and this test still stays green. Seed a few distinct cached counts and assert the returned rows are actually sorted before checking query count.

Also applies to: 3342-3349

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ami/main/tests.py` around lines 3295 - 3305, The test seeds 30 collections
all with the exact same image set so the ordering-by -source_images_count is
invisible; change the seeding in the SourceImageCollection creation (in the
block that creates "qcount-collection-{i}" and the similar block at 3342-3349)
to produce several distinct cached counts (e.g., create collections with varying
numbers of images or explicitly set the cached count field on
SourceImageCollection after creation), then query the view and assert that the
returned rows are actually sorted by source_images_count (check the sequence of
returned collection IDs or counts) before performing the query-count assertion;
update the creation loop for SourceImageCollection and add an assertion
verifying ordering so the test will fail if ordering is ignored.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/claude/planning/cached-counts-reconcile-followup.md`:
- Line 8: The PR reference at the start of the line " `#1301` ships the perf win
(denormalized columns, signals, ML-path recompute," is being parsed as a
Markdown heading; update that line to escape the leading hash or wrap the PR
number in inline code so it renders as plain text (e.g., replace "`#1301`" with
"\`#1301`" or "`#1301`") ensuring the rest of the sentence remains unchanged and
the file docs/claude/planning/cached-counts-reconcile-followup.md still conveys
the same message.

---

Nitpick comments:
In `@ami/main/tests.py`:
- Around line 3444-3470: The test currently only asserts that
source_images_count is unchanged after calling get_source_image_counts(); extend
it to also assert that the other cached DB columns
(source_images_processed_count and source_images_with_detections_count) remain
untouched in the no-write case: after calling
self.collection.get_source_image_counts() and refreshing from DB via
self.collection.refresh_from_db(), add assertions that
self.collection.source_images_processed_count == 0 and
self.collection.source_images_with_detections_count == 0 so the test fails if
get_source_image_counts() erroneously persists those fields on
SourceImageCollection.
- Around line 3376-3385: Add a new test method (e.g.,
test_count_resets_on_image_clear) that sets up the collection with
self.collection.images.set(self.images), calls self.collection.images.clear(),
refreshes state with self._refresh(), and then asserts that the three cached
count columns on the collection (including source_images_count) have dropped to
zero; reference collection.images.clear() for the action and
collection.source_images_count (and the other two cached count attributes on the
collection) for the assertions.
- Around line 3295-3305: The test seeds 30 collections all with the exact same
image set so the ordering-by -source_images_count is invisible; change the
seeding in the SourceImageCollection creation (in the block that creates
"qcount-collection-{i}" and the similar block at 3342-3349) to produce several
distinct cached counts (e.g., create collections with varying numbers of images
or explicitly set the cached count field on SourceImageCollection after
creation), then query the view and assert that the returned rows are actually
sorted by source_images_count (check the sequence of returned collection IDs or
counts) before performing the query-count assertion; update the creation loop
for SourceImageCollection and add an assertion verifying ordering so the test
will fail if ordering is ignored.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: c4041159-0f7e-4adf-86e1-2c798d80337d

📥 Commits

Reviewing files that changed from the base of the PR and between 9ac2ce1 and 2533857.

📒 Files selected for processing (9)

• ami/base/models.py
• ami/main/api/serializers.py
• ami/main/migrations/0085_denormalize_sourceimagecollection_counts.py
• ami/main/migrations/0086_backfill_sourceimagecollection_counts.py
• ami/main/migrations/0087_use_cached_count_field.py
• ami/main/signals.py
• ami/main/tasks.py
• ami/main/tests.py
• docs/claude/planning/cached-counts-reconcile-followup.md

🚧 Files skipped from review as they are similar to previous changes (1)

• ami/base/models.py

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Escape PR reference to avoid accidental/malformed heading rendering.

Line 8 starts with #1301, which markdown parsers/lint treat as an ATX heading token pattern. Use an escaped hash or inline code so this stays a plain PR reference.

Suggested fix

-#1301 ships the perf win (denormalized columns, signals, ML-path recompute,
+\`#1301` ships the perf win (denormalized columns, signals, ML-path recompute,

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	#1301 ships the perf win (denormalized columns, signals, ML-path recompute,
	\`#1301` ships the perf win (denormalized columns, signals, ML-path recompute,

🧰 Tools

🪛 markdownlint-cli2 (0.22.1)

[warning] 8-8: No space after hash on atx style heading

(MD018, no-missing-space-atx)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/claude/planning/cached-counts-reconcile-followup.md` at line 8, The PR
reference at the start of the line " `#1301` ships the perf win (denormalized
columns, signals, ML-path recompute," is being parsed as a Markdown heading;
update that line to escape the leading hash or wrap the PR number in inline code
so it renders as plain text (e.g., replace "`#1301`" with "\`#1301`" or "`#1301`")
ensuring the rest of the sentence remains unchanged and the file
docs/claude/planning/cached-counts-reconcile-followup.md still conveys the same
message.