Fix: treat naive as_of/now datetimes as UTC in recency scoring

potpie/context-engine/src/potpie_context_engine/domain/ranking.py

@@ -117,7 +117,7 @@ def rank(
         ``breakdown`` so readers can surface "why this ranked high" if
         the agent asks.
         """
-        now = context.now or datetime.now(tz=timezone.utc)
+        now = _ensure_utc(context.now) or datetime.now(tz=timezone.utc)
         ranked: list[RankedItem] = []
         for cand in candidates:
             breakdown = self._score_one(cand, now=now, context=context)
@@ -176,6 +176,20 @@ def _combine(self, breakdown: Mapping[str, float]) -> float:
 # ---------------------------------------------------------------------------
 
 
+def _ensure_utc(value: datetime | None) -> datetime | None:
+    """Treat a naive datetime as UTC so aware/naive inputs never mix.
+
+    ``TaskContext.now`` is caller-supplied (e.g. an agent's ``as_of``); a naive
+    value would otherwise raise ``TypeError`` when subtracted from the
+    tz-normalized ``valid_at`` in :func:`_recency_score`.
+    """
+    if value is None:
+        return None
+    if value.tzinfo is None:
+        return value.replace(tzinfo=timezone.utc)
+    return value
+
+
 def _clamp(value: float | None, *, default: float) -> float:
     if value is None:
         return default
@@ -198,8 +212,7 @@ def _recency_score(
     """Exponential decay; freshness preference shifts the half-life."""
     if valid_at is None:
         return 0.5
-    if valid_at.tzinfo is None:
-        valid_at = valid_at.replace(tzinfo=timezone.utc)
+    valid_at = _ensure_utc(valid_at)
     age = max(now - valid_at, timedelta(0))
 
     effective_half_life = half_life

potpie/context-engine/tests/unit/test_ranking.py

@@ -249,3 +249,33 @@ def test_truncate_zero_returns_empty(self) -> None:
         service = RankingService()
         ranked = service.rank([_make_candidate(key="a")], _ctx())
         assert truncate(ranked, max_items=0) == []
+
+
+class TestNaiveDatetimeHandling:
+    """A naive ``TaskContext.now`` (e.g. an agent's date-only ``as_of``) must
+    rank instead of raising naive-minus-aware ``TypeError``."""
+
+    @staticmethod
+    def _naive_ctx(now: datetime) -> TaskContext:
+        return TaskContext(pot_id="pot-1", now=now)
+
+    def test_naive_now_with_aware_valid_at_does_not_crash(self) -> None:
+        service = RankingService()
+        cand = _make_candidate(key="a", valid_at=_NOW - timedelta(days=3))
+        ranked = service.rank([cand], self._naive_ctx(datetime(2026, 5, 20)))
+        assert len(ranked) == 1
+
+    def test_naive_now_with_naive_valid_at_does_not_crash(self) -> None:
+        service = RankingService()
+        cand = _make_candidate(key="a", valid_at=datetime(2026, 5, 17))
+        ranked = service.rank([cand], self._naive_ctx(datetime(2026, 5, 20)))
+        assert len(ranked) == 1
+
+    def test_naive_now_scores_as_utc(self) -> None:
+        """A naive now is interpreted as UTC: same instant, same score."""
+        service = RankingService()
+        cand = _make_candidate(key="a", valid_at=_NOW - timedelta(days=3))
+        aware = service.rank([cand], self._naive_ctx(_NOW))[0]
+        naive = service.rank([cand], self._naive_ctx(_NOW.replace(tzinfo=None)))[0]
+        assert naive.score == aware.score
+        assert naive.breakdown["recency"] == aware.breakdown["recency"]