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

potpie/context-engine/adapters/inbound/mcp/server.py

@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 import logging
-from datetime import datetime
+from datetime import datetime, timezone
 from functools import lru_cache
 from typing import Any
 
@@ -54,7 +54,13 @@ def _parse_as_of_iso(value: str | None) -> datetime | None:
     s = value.strip()
     if s.endswith("Z"):
         s = s[:-1] + "+00:00"
-    return datetime.fromisoformat(s)
+    parsed = datetime.fromisoformat(s)
+    if parsed.tzinfo is None:
+        # Treat naive timestamps (e.g. "2024-06-01") as UTC. A naive datetime
+        # would otherwise flow into recency scoring as `now` and crash against
+        # tz-normalized claim times (naive minus aware TypeError).
+        parsed = parsed.replace(tzinfo=timezone.utc)
+    return parsed
 
 
 def _scope(**fields: Any) -> dict[str, Any]:

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_mcp_as_of_parsing.py

@@ -0,0 +1,55 @@
+"""MCP ``as_of`` parsing: every accepted form must yield a tz-aware datetime.
+
+A naive datetime here becomes the ranker's ``now`` (via ``ReadRequest.as_of``
+-> ``TaskContext.now``) and crashed recency scoring against tz-normalized
+claim times with ``TypeError: can't subtract offset-naive and offset-aware
+datetimes``.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+import pytest
+
+from adapters.inbound.mcp.server import _parse_as_of_iso
+
+pytestmark = pytest.mark.unit
+
+
+def test_none_and_blank_return_none() -> None:
+    assert _parse_as_of_iso(None) is None
+    assert _parse_as_of_iso("") is None
+    assert _parse_as_of_iso("   ") is None
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        "2024-06-01",
+        "2024-06-01T10:30:00",
+        "2024-06-01T10:30:00Z",
+        "2024-06-01T10:30:00+05:30",
+    ],
+)
+def test_parsed_as_of_is_always_tz_aware(value: str) -> None:
+    parsed = _parse_as_of_iso(value)
+    assert parsed is not None
+    assert parsed.tzinfo is not None
+
+
+def test_naive_input_interpreted_as_utc() -> None:
+    parsed = _parse_as_of_iso("2024-06-01T10:30:00")
+    assert parsed == datetime(2024, 6, 1, 10, 30, tzinfo=timezone.utc)
+
+
+def test_explicit_offset_preserved() -> None:
+    parsed = _parse_as_of_iso("2024-06-01T10:30:00+05:30")
+    assert parsed is not None
+    assert parsed.utcoffset() is not None
+    assert parsed.utcoffset().total_seconds() == 5.5 * 3600
+
+
+def test_z_suffix_is_utc() -> None:
+    parsed = _parse_as_of_iso("2024-06-01T10:30:00Z")
+    assert parsed == datetime(2024, 6, 1, 10, 30, tzinfo=timezone.utc)

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"]