fix: enforce search_web's cutoff date from the harness, not the LLM

aieng-forecasting/aieng/forecasting/methods/agentic/adk_runner.py

@@ -182,7 +182,13 @@ def runner(self) -> InMemoryRunner:
         """Underlying ADK runner (session, artifact, memory services)."""
         return self._runner
 
-    async def _resolve_session_id(self, user_id: str | None, session_id: str | None) -> str:
+    async def _resolve_session_id(
+        self,
+        user_id: str | None,
+        session_id: str | None,
+        *,
+        initial_state: dict[str, Any] | None = None,
+    ) -> str:
         """Return the ADK session id to use for a single turn.
 
         Parameters
@@ -192,6 +198,11 @@ async def _resolve_session_id(self, user_id: str | None, session_id: str | None)
         session_id : str or None
             Explicit session id from the caller.  ``None`` triggers sticky-session
             lookup or new-session creation depending on ``fresh_session_per_message``.
+        initial_state : dict[str, Any] or None
+            Seeded into a newly-created session's state. Only takes effect when
+            this call actually creates a session — has no effect when an
+            existing sticky session (``fresh_session_per_message=False``) is
+            reused, since session state can only be seeded at creation.
 
         Returns
         -------
@@ -205,6 +216,7 @@ async def _resolve_session_id(self, user_id: str | None, session_id: str | None)
             new_session = await self._runner.session_service.create_session(
                 app_name=self.config.app_name,
                 user_id=user_id,
+                state=initial_state,
             )
             sid = new_session.id
         elif session_id is not None:
@@ -216,6 +228,7 @@ async def _resolve_session_id(self, user_id: str | None, session_id: str | None)
             new_session = await self._runner.session_service.create_session(
                 app_name=self.config.app_name,
                 user_id=user_id,
+                state=initial_state,
             )
             sid = new_session.id
             self._conversation_session_by_user[user_id] = sid
@@ -229,6 +242,7 @@ async def run_text_async(
         user_id: str | None = None,
         session_id: str | None = None,
         run_config: RunConfig | None = None,
+        initial_state: dict[str, Any] | None = None,
     ) -> str:
         """Run one user turn; return the first final model text or an empty string.
 
@@ -247,6 +261,12 @@ async def run_text_async(
         run_config : RunConfig | None, optional
             The run configuration to use for the run. If not provided, the default
             run configuration is used.
+        initial_state : dict[str, Any] | None, optional
+            Seeded into the session's state when this call creates a new
+            session (see :meth:`_resolve_session_id`). Use this to pass
+            harness-controlled values (e.g. a forecast's ``as_of`` date) that
+            tools can read via ``ToolContext.state`` without the LLM being
+            able to see or influence them.
 
         Returns
         -------
@@ -270,7 +290,7 @@ async def run_text_async(
 
         user_id = user_id or self.config.default_user_id
 
-        session_id = await self._resolve_session_id(user_id, session_id)
+        session_id = await self._resolve_session_id(user_id, session_id, initial_state=initial_state)
 
         content = genai_types.Content(role="user", parts=[genai_types.Part(text=prompt)])
 

aieng-forecasting/aieng/forecasting/methods/agentic/agent_factory.py

@@ -79,6 +79,13 @@ def filter(self, record: logging.LogRecord) -> bool:
 # final text, giving the predictor the structured JSON it expects.
 SMR_STATE_KEY = "__smr_output__"
 
+# Session-state key AgentPredictor seeds with the current prediction's as_of
+# date before each run (see AdkTextRunner.run_text_async's initial_state).
+# search_web reads this via its ADK-injected ToolContext as the authoritative
+# cutoff — unlike the LLM-supplied cutoff_date argument, the model can never
+# see or influence this key, so it can't be silently omitted or spoofed.
+AS_OF_STATE_KEY = "__as_of__"
+
 
 def _build_set_model_response_tool() -> FunctionTool:
     """Return a proxy-compatible ``set_model_response`` shim.
@@ -390,7 +397,7 @@ async def _do_search(user_content: str) -> tuple[str, list[str]]:
         ]
         return content, sources
 
-    async def search_web(query: str, cutoff_date: str | None = None) -> str:
+    async def search_web(query: str, cutoff_date: str | None = None, tool_context: ToolContext | None = None) -> str:
         """Search the web and return a grounded summary with source URLs.
 
         Args:
@@ -404,22 +411,45 @@ async def search_web(query: str, cutoff_date: str | None = None) -> str:
             When cutoff verification is enabled and cannot be satisfied
             within the attempt budget, returns a ``[SEARCH_VERIFICATION_FAILED]``
             sentinel instead of unverified content.
+
+        Notes
+        -----
+        ``tool_context`` is not part of the LLM-visible tool schema — ADK
+        injects it automatically because of its type annotation. When
+        :class:`AgentPredictor` runs this tool, it has already seeded the
+        session with the current prediction's ``as_of`` date under
+        :data:`AS_OF_STATE_KEY`; that harness-controlled value is the
+        authoritative cutoff whenever present, since (unlike ``cutoff_date``)
+        the calling LLM cannot see, omit, or alter it. This closes a bypass
+        where the model simply didn't pass ``cutoff_date`` and both the soft
+        cutoff instruction and the verifier below were silently skipped.
         """
-        needs_verification = bool(cutoff_date and config.enforce_cutoff)
+        harness_as_of = tool_context.state.get(AS_OF_STATE_KEY) if tool_context is not None else None
+        if harness_as_of and cutoff_date and harness_as_of != cutoff_date:
+            logger.warning(
+                "search_web: cutoff_date=%r disagrees with harness as_of=%r; using the harness value.",
+                cutoff_date,
+                harness_as_of,
+            )
+        effective_cutoff = harness_as_of or cutoff_date
+
+        needs_verification = bool(effective_cutoff and config.enforce_cutoff)
         if not needs_verification:
             content, sources = await _do_search(query)
             return _format_result(content, sources)
 
         negative_guidance = ""
         for attempt in range(1, config.verifier_max_attempts + 1):
-            user_content = query + f"\n\nOnly include and cite information published strictly before {cutoff_date}."
+            user_content = (
+                query + f"\n\nOnly include and cite information published strictly before {effective_cutoff}."
+            )
             if negative_guidance:
                 user_content += f"\n\n{negative_guidance}"
             content, sources = await _do_search(user_content)
             verdict = await _verify_no_leakage(
                 text=content,
                 query=query,
-                cutoff_date=cutoff_date,  # type: ignore[arg-type]
+                cutoff_date=effective_cutoff,  # type: ignore[arg-type]
                 verifier_model=config.verifier_model,
                 openai_base_url=openai_base_url,
                 openai_api_key=openai_api_key,
@@ -437,14 +467,14 @@ async def search_web(query: str, cutoff_date: str | None = None) -> str:
             logger.warning("search_web attempt %d flagged %d claim(s); retrying.", attempt, len(verdict.flagged_claims))
             negative_guidance = (
                 f"Your previous search result may have included information published on or after "
-                f"{cutoff_date}. Do not repeat or rely on these claims:\n- "
+                f"{effective_cutoff}. Do not repeat or rely on these claims:\n- "
                 + "\n- ".join(verdict.flagged_claims or ["(unspecified — be more conservative)"])
             )
 
         logger.error("search_web exhausted %d attempts without a verified clean result.", config.verifier_max_attempts)
         return (
             f"[SEARCH_VERIFICATION_FAILED] Could not verify search results as free of information "
-            f"published on or after {cutoff_date} after {config.verifier_max_attempts} attempts. "
+            f"published on or after {effective_cutoff} after {config.verifier_max_attempts} attempts. "
             "Treat this as no verified news context being available for this query."
         )
 

aieng-forecasting/aieng/forecasting/methods/agentic/predictor.py

@@ -29,7 +29,7 @@
 from aieng.forecasting.evaluation.predictor import Predictor
 from aieng.forecasting.evaluation.task import ForecastingTask
 from aieng.forecasting.methods.agentic.adk_runner import AdkTextRunner, AdkTextRunnerConfig
-from aieng.forecasting.methods.agentic.agent_factory import AgentConfig, build_adk_agent
+from aieng.forecasting.methods.agentic.agent_factory import AS_OF_STATE_KEY, AgentConfig, build_adk_agent
 from aieng.forecasting.methods.agentic.outputs import AgentForecastOutput
 from aieng.forecasting.methods.llm_processes._client import strip_markdown_fence, trace_url_for
 from google.adk.agents.base_agent import BaseAgent
@@ -277,7 +277,11 @@ def predict(self, task: ForecastingTask, context: ForecastContext) -> list[Predi
             validation errors on the agent's JSON are not swallowed.
         """
         prompt = self.prompt_builder(task=task, context=context)
-        output_str = _run_coroutine_sync(self._runner.run_text_async(prompt))
+        # Seed the harness-controlled as_of into the ADK session before the run,
+        # so search_web can enforce it via ToolContext.state regardless of
+        # whether the LLM remembers to pass a matching cutoff_date argument.
+        initial_state = {AS_OF_STATE_KEY: str(context.as_of)[:10]}
+        output_str = _run_coroutine_sync(self._runner.run_text_async(prompt, initial_state=initial_state))
 
         # Normalise: strip markdown fences before validation so any model can
         # be swapped in without breaking the parse layer.

aieng-forecasting/tests/aieng/forecasting/methods/agentic/test_adk_runner.py

@@ -147,6 +147,37 @@ async def test_caller_supplied_session_id_is_ignored(self, patch_runner_cls, moc
 
         assert patch_runner_cls.run_async.call_args.kwargs["session_id"] == "fresh-sid"
 
+    async def test_initial_state_passed_through_to_create_session(self, patch_runner_cls, mock_agent) -> None:
+        """initial_state reaches session_service.create_session's state kwarg.
+
+        This is what lets AgentPredictor seed a harness-controlled as_of value
+        that search_web can enforce via ToolContext.state, regardless of
+        whether the LLM remembers to pass a matching cutoff_date argument.
+        """
+        patch_runner_cls.session_service.create_session.return_value = _session("s1")
+        patch_runner_cls.run_async.return_value = _stream(_final_event("ok"))
+        runner = AdkTextRunner(
+            mock_agent,
+            config=AdkTextRunnerConfig(app_name="app", fresh_session_per_message=True),
+        )
+
+        await runner.run_text_async("hello", user_id="alice", initial_state={"__as_of__": "2024-01-15"})
+
+        assert patch_runner_cls.session_service.create_session.call_args.kwargs["state"] == {"__as_of__": "2024-01-15"}
+
+    async def test_no_initial_state_passes_none_to_create_session(self, patch_runner_cls, mock_agent) -> None:
+        """Omitting initial_state passes state=None (unchanged default behavior)."""
+        patch_runner_cls.session_service.create_session.return_value = _session("s1")
+        patch_runner_cls.run_async.return_value = _stream(_final_event("ok"))
+        runner = AdkTextRunner(
+            mock_agent,
+            config=AdkTextRunnerConfig(app_name="app", fresh_session_per_message=True),
+        )
+
+        await runner.run_text_async("hello", user_id="alice")
+
+        assert patch_runner_cls.session_service.create_session.call_args.kwargs["state"] is None
+
 
 # ---------------------------------------------------------------------------
 # Session resolution — fresh_session_per_message=False (sticky)

aieng-forecasting/tests/aieng/forecasting/methods/agentic/test_agent_factory.py

@@ -2,11 +2,14 @@
 
 import inspect
 import json
+import logging
 from pathlib import Path
+from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from aieng.forecasting.methods.agentic.agent_factory import (
+    AS_OF_STATE_KEY,
     AgentConfig,
     CodeExecutionConfig,
     ContextRetrievalConfig,
@@ -264,6 +267,8 @@ def test_returns_callable_with_expected_signature(self) -> None:
         assert "query" in sig.parameters
         assert "cutoff_date" in sig.parameters
         assert sig.parameters["cutoff_date"].default is None
+        assert "tool_context" in sig.parameters
+        assert sig.parameters["tool_context"].default is None
 
     @pytest.mark.asyncio
     async def test_cutoff_date_appended_when_enforce_cutoff_true(self) -> None:
@@ -559,3 +564,73 @@ async def _fake_acompletion(**kwargs):  # type: ignore[override]
 
         assert len(calls) == config.verifier_max_attempts * 2
         assert result.startswith("[SEARCH_VERIFICATION_FAILED]")
+
+    @pytest.mark.asyncio
+    async def test_harness_as_of_triggers_verification_even_without_cutoff_date(self) -> None:
+        """Verification runs from harness as_of alone when the LLM omits cutoff_date.
+
+        This is the production bypass this fix closes: 11 of 14 search_web
+        calls in a real trace omitted cutoff_date, silently skipping the guard.
+        """
+        config = ContextRetrievalConfig(enabled=True, instruction="Search assistant.")
+        tool = _build_search_tool(config, openai_base_url="https://proxy.example.com/v1", openai_api_key="test-key")
+        fake_tool_context = SimpleNamespace(state={AS_OF_STATE_KEY: "2024-01-15"})
+        calls: list[dict] = []
+
+        async def _fake_acompletion(**kwargs):  # type: ignore[override]
+            calls.append(kwargs)
+            if kwargs["model"] == f"openai/{config.verifier_model}":
+                return self._verify_response(clean=True, confidence=9, filtered_text="Clean summary.")
+            return self._search_response("Raw summary with a source.")
+
+        with patch("litellm.acompletion", new=AsyncMock(side_effect=_fake_acompletion)):
+            result = await tool(query="WTI price", cutoff_date=None, tool_context=fake_tool_context)
+
+        assert len(calls) == 2
+        assert result == "Clean summary."
+        search_user_msg = next(m for m in calls[0]["messages"] if m["role"] == "user")
+        assert "2024-01-15" in search_user_msg["content"]
+
+    @pytest.mark.asyncio
+    async def test_harness_as_of_overrides_disagreeing_cutoff_date(self, caplog: pytest.LogCaptureFixture) -> None:
+        """The harness as_of wins over a disagreeing cutoff_date, and logs a warning."""
+        config = ContextRetrievalConfig(enabled=True, instruction="Search assistant.")
+        tool = _build_search_tool(config, openai_base_url="https://proxy.example.com/v1", openai_api_key="test-key")
+        fake_tool_context = SimpleNamespace(state={AS_OF_STATE_KEY: "2024-01-15"})
+        calls: list[dict] = []
+
+        async def _fake_acompletion(**kwargs):  # type: ignore[override]
+            calls.append(kwargs)
+            if kwargs["model"] == f"openai/{config.verifier_model}":
+                return self._verify_response(clean=True, confidence=9, filtered_text="Clean summary.")
+            return self._search_response("Raw summary with a source.")
+
+        with (
+            caplog.at_level(logging.WARNING),
+            patch("litellm.acompletion", new=AsyncMock(side_effect=_fake_acompletion)),
+        ):
+            await tool(query="WTI price", cutoff_date="2026-06-01", tool_context=fake_tool_context)
+
+        search_user_msg = next(m for m in calls[0]["messages"] if m["role"] == "user")
+        assert "2024-01-15" in search_user_msg["content"]
+        assert "2026-06-01" not in search_user_msg["content"]
+        assert any("disagrees with harness as_of" in r.message for r in caplog.records)
+
+    @pytest.mark.asyncio
+    async def test_no_tool_context_falls_back_to_cutoff_date(self) -> None:
+        """With no tool_context (e.g. interactive use), cutoff_date still works."""
+        config = ContextRetrievalConfig(enabled=True, instruction="Search assistant.")
+        tool = _build_search_tool(config, openai_base_url="https://proxy.example.com/v1", openai_api_key="test-key")
+        calls: list[dict] = []
+
+        async def _fake_acompletion(**kwargs):  # type: ignore[override]
+            calls.append(kwargs)
+            if kwargs["model"] == f"openai/{config.verifier_model}":
+                return self._verify_response(clean=True, confidence=9, filtered_text="Clean summary.")
+            return self._search_response("Raw summary with a source.")
+
+        with patch("litellm.acompletion", new=AsyncMock(side_effect=_fake_acompletion)):
+            result = await tool(query="WTI price", cutoff_date="2024-01-15")
+
+        assert len(calls) == 2
+        assert result == "Clean summary."

aieng-forecasting/tests/aieng/forecasting/methods/agentic/test_predictor.py

@@ -26,7 +26,7 @@
 from aieng.forecasting.data.context import ForecastContext
 from aieng.forecasting.evaluation.prediction import STANDARD_QUANTILES, Prediction
 from aieng.forecasting.evaluation.task import ForecastingTask
-from aieng.forecasting.methods.agentic.agent_factory import AgentConfig
+from aieng.forecasting.methods.agentic.agent_factory import AS_OF_STATE_KEY, AgentConfig
 from aieng.forecasting.methods.agentic.outputs import AgentForecastOutput, ContinuousAgentForecastOutput
 from aieng.forecasting.methods.agentic.predictor import AgentPredictor
 from pydantic import ValidationError
@@ -57,14 +57,16 @@ def __init__(
         self._agent.model = model
         # No tracing in unit tests, so no trace is captured.
         self.last_trace_id: str | None = None
+        self.run_text_async_calls: list[dict[str, Any]] = []
 
     @property
     def agent(self) -> Any:
         """Return the stub agent so the predictor can read ``name``/``model``."""
         return self._agent
 
-    async def run_text_async(self, prompt: str, **_: Any) -> str:
-        """Return the canned response regardless of prompt."""
+    async def run_text_async(self, prompt: str, **kwargs: Any) -> str:
+        """Record the call and return the canned response regardless of prompt."""
+        self.run_text_async_calls.append({"prompt": prompt, **kwargs})
         return self._response
 
 
@@ -213,6 +215,22 @@ def test_prompt_builder_is_invoked_with_task_and_context(self) -> None:
 
         builder.assert_called_once_with(task=task, context=context)
 
+    def test_seeds_harness_as_of_into_run_text_async(self) -> None:
+        """predict() seeds the session with context.as_of so search_web can enforce it.
+
+        This is what closes the bypass where an LLM omits cutoff_date on a
+        search_web call: the harness-controlled as_of is injected regardless
+        of what (if anything) the model passes.
+        """
+        predictor, _ = _make_predictor(response=_output_json([1]))
+        context = _context()  # as_of = datetime(2024, 1, 1)
+
+        predictor.predict(_task([1]), context)
+
+        calls = predictor._runner.run_text_async_calls  # type: ignore[attr-defined]
+        assert len(calls) == 1
+        assert calls[0]["initial_state"] == {AS_OF_STATE_KEY: "2024-01-01"}
+
     def test_fenced_json_is_accepted_and_converts_to_predictions(self) -> None:
         """JSON wrapped in a ```json ... ``` fence still produces valid predictions.