From c1b4578e7d400ac98bf9c68dc277b7c14364a031 Mon Sep 17 00:00:00 2001 From: Aditya painuli Date: Mon, 13 Jul 2026 15:15:20 +0530 Subject: [PATCH] Fix: treat naive as_of/now datetimes as UTC in recency scoring (#1008) --- .../adapters/inbound/mcp/server.py | 10 +++- potpie/context-engine/domain/ranking.py | 19 ++++++- .../tests/unit/test_mcp_as_of_parsing.py | 55 +++++++++++++++++++ .../context-engine/tests/unit/test_ranking.py | 30 ++++++++++ 4 files changed, 109 insertions(+), 5 deletions(-) create mode 100644 potpie/context-engine/tests/unit/test_mcp_as_of_parsing.py diff --git a/potpie/context-engine/adapters/inbound/mcp/server.py b/potpie/context-engine/adapters/inbound/mcp/server.py index 86f6597be..ee5dfc2dd 100644 --- a/potpie/context-engine/adapters/inbound/mcp/server.py +++ b/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]: diff --git a/potpie/context-engine/domain/ranking.py b/potpie/context-engine/domain/ranking.py index b03e2558d..a96375352 100644 --- a/potpie/context-engine/domain/ranking.py +++ b/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 diff --git a/potpie/context-engine/tests/unit/test_mcp_as_of_parsing.py b/potpie/context-engine/tests/unit/test_mcp_as_of_parsing.py new file mode 100644 index 000000000..af171420b --- /dev/null +++ b/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) diff --git a/potpie/context-engine/tests/unit/test_ranking.py b/potpie/context-engine/tests/unit/test_ranking.py index 4229af5bc..9005fd2db 100644 --- a/potpie/context-engine/tests/unit/test_ranking.py +++ b/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"]