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

• New Features

  • Added persistent cached count columns for collections, deployments, events and images plus transaction-aware async recompute to keep them up to date.

• Improvements

  • Simplified collection list queries to rely on stored counts and prefetching for more predictable performance.
  • Added an integrity check and repair flow to detect and reconcile stale cached counts.

• Tests

  • Expanded tests to validate count consistency across M2M, detection changes, and bulk operations.

[netlify]

✅ Deploy Preview for antenna-preview ready!

Name	Link
🔨 Latest commit	9ac2ce1
🔍 Latest deploy log	https://app.netlify.com/projects/antenna-preview/deploys/6a06c92eb7c50b0008eaea85
😎 Deploy Preview	https://deploy-preview-1301--antenna-preview.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.
Lighthouse	1 paths audited Performance: 65 (no change from production) Accessibility: 89 (no change from production) Best Practices: 92 (🔴 down 8 from production) SEO: 92 (no change from production) PWA: 80 (no change from production) View the detailed breakdown and full score reports
🤖 Make changes	Run an agent on this branch

---

To edit notification comments on pull requests, go to your Netlify project configuration.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds a CachedCountField marker and a per-transaction dedup scheduler that enqueues Celery recompute tasks; migrates and backfills denormalized count columns, implements model/queryset update_cached_counts, wires signals and tasks, adds reconciliation checks, and expands tests for correctness and query behavior.

Changes

Denormalized Count Field Infrastructure

Layer / File(s)	Summary
CachedCountField & scheduler ami/base/models.py, ami/base/cached_counts.py	Adds CachedCountField and a per-DB-connection schedule_recompute with on_commit flush that enqueues one recompute_cached_counts_task per unique (model_label, pk).
Database schema and backfill migrations ami/main/migrations/0085_denormalize_sourceimagecollection_counts.py, ami/main/migrations/0086_use_cached_count_field.py	Adds three denormalized count columns to SourceImageCollection and backfills them with a grouped UPDATE; then converts cached count fields across models to CachedCountField.
Model conversions & SourceImageCollection computation ami/main/models.py	Converts Deployment/Event/SourceImage count fields to CachedCountField; removes with_source_images_* queryset helpers; adds source_images_* cached fields on SourceImageCollection; implements get_source_image_counts() and update_calculated_fields(save=...).
Signals, ML pipeline, and API view ami/main/signals.py, ami/ml/models/pipeline.py, ami/main/api/views.py	Adds m2m_changed and Detection post_save/post_delete handlers that call .update_cached_counts(); refreshes collections after bulk Detection creation; simplifies SourceImageCollectionViewSet.queryset.
Celery tasks & reconciliation checks ami/main/tasks.py, ami/main/checks/cached_counts.py	Adds recompute_cached_counts_task and reconcile_cached_counts_task; implements discovery, stale-finding, and reconcile_cached_counts which can repair synchronously via update_cached_counts(run_async=False).
Tests ami/main/tests.py	Adds transaction-aware tests for list query counts, denormalization correctness under M2M/Detection changes, and cached-count reconciliation; adjusts query-count assertions and imports.
Design spec docs/superpowers/specs/2026-05-14-cached-counts-update-method-design.md	Design document describing scheduler, update_cached_counts API, call-site changes, risks, migration path, and test plan (including dedup-focused test).

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related issues

• Dynamic capture sets: save filter queries as reusable capture sets (supersedes #907, #487) #1253 — touches SourceImageCollection count fields and related queryset/count helper logic; this PR migrates those counts to denormalized CachedCountField storage and removes the annotation helpers.

Possibly related PRs

• RolnickLab/antenna#1300: Prior rewrite of with_source_images_*_count annotation logic; this PR replaces that approach with denormalized cached fields.
• RolnickLab/antenna#1093: Adds/uses the SourceImageCollection "processed" count; this PR migrates that field to denormalized storage and adjusts where it's computed/used.

Suggested labels

backend

Suggested reviewers

• annavik

Poem

"🐰 I hop through counts both new and old,
Cached fields snug, their stories told,
Signals nudge and tasks dedupe,
Tests confirm we didn't goof,
Hooray — small carrots for the code!"

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 43.14% 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 objective of this PR - optimizing the performance of the SourceImageCollection list view by denormalizing cached counts and introducing CachedCountField markers.
Description check	✅ Passed	The PR description comprehensively covers all required template sections: summary, detailed list of changes, denormalization/CachedCountField details with technical depth, rationale, test coverage, and rollout/deployment notes. All major sections from the template are addressed.
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 ready!

Name	Link
🔨 Latest commit	9ac2ce1
🔍 Latest deploy log	https://app.netlify.com/projects/antenna-ssec/deploys/6a06c92e77c58b00085a5dfd
😎 Deploy Preview	https://deploy-preview-1301--antenna-ssec.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.
🤖 Make changes	Run an agent on this branch

---

To edit notification comments on pull requests, go to your Netlify project configuration.

[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).