Extend preprocessing/postprocessing hooks with Session and Prompt access

- preprocess_request now receives Session for context-aware preprocessing
- postprocess_response now receives both Prompt and Session
- Preprocessing moved inside prepare() for proper request transformation
- Use instance variables to pass context from prepare() to postprocess_response()

Author: claude

src/trivia_agent/worker.py

@@ -269,29 +269,45 @@ def __init__(
         self._workspace_dir = workspace_dir or DEFAULT_WORKSPACE_DIR
         self._base_template = build_prompt_template()
         self._overrides_store = overrides_store
+        self._last_prompt: Prompt[TriviaResponse] | None = None
+        self._last_session: Session | None = None
 
-    def preprocess_request(self, request: TriviaRequest) -> TriviaRequest:
+    def preprocess_request(
+        self,
+        request: TriviaRequest,
+        session: Session,
+    ) -> TriviaRequest:
         """Transform request before agent processing.
 
         Override this method in subclasses to implement custom preprocessing
-        logic such as validation, normalization, or enrichment.
+        logic such as validation, normalization, or enrichment. Has access to
+        the session for context-aware preprocessing.
 
         Args:
             request: The incoming TriviaRequest.
+            session: The session for this execution.
 
         Returns:
             TriviaRequest: The preprocessed request.
         """
         return request
 
-    def postprocess_response(self, response: TriviaResponse) -> TriviaResponse:
+    def postprocess_response(
+        self,
+        response: TriviaResponse,
+        prompt: Prompt[TriviaResponse],
+        session: Session,
+    ) -> TriviaResponse:
         """Transform response before returning to caller.
 
         Override this method in subclasses to implement custom postprocessing
-        logic such as formatting, cleanup, or validation.
+        logic such as formatting, cleanup, or validation. Has access to the
+        prompt and session for context-aware postprocessing.
 
         Args:
             response: The TriviaResponse from the agent.
+            prompt: The prompt used for this execution.
+            session: The session used for this execution.
 
         Returns:
             TriviaResponse: The postprocessed response.
@@ -310,38 +326,20 @@ def prepare(
         for isolation, builds the complete PromptTemplate with workspace section,
         binds request parameters, and optionally applies experiment overrides.
 
-        Note: Request preprocessing is applied in the execute() method before
-        prepare() is called. This ensures the preprocessed request is used
-        consistently throughout the execution flow.
-
-        This method demonstrates key WINK patterns:
-
-        - **Session per request**: Each request gets its own Session for proper
-          isolation between concurrent requests
-        - **Dynamic workspace section**: ClaudeAgentWorkspaceSection is created
-          here (not in build_prompt_template) because it needs the Session
-        - **Parameter binding**: Binds QuestionParams with the user's question
-          and EmptyParams for parameterless sections
-        - **Experiment support**: Uses experiment.overrides_tag to select prompt
-          variants for A/B testing; defaults to "latest"
-        - **Override seeding**: Automatically seeds the overrides store with
-          current prompt state, creating editable files for customization
+        Calls preprocess_request() to allow request transformation before binding.
 
         Args:
-            request: TriviaRequest containing the question field. The question
-                is bound to QuestionSection via QuestionParams.
-            experiment: Optional Experiment instance for evaluation runs. When
-                provided, uses experiment.overrides_tag to select prompt variant.
-                Pass None for production requests.
+            request: TriviaRequest containing the question field.
+            experiment: Optional Experiment instance for evaluation runs.
 
         Returns:
-            tuple[Prompt[TriviaResponse], Session]: A 2-tuple containing:
-                - Prompt: Fully configured prompt with all sections and bound
-                  parameters, ready for adapter.run()
-                - Session: Fresh session instance for this request's execution
+            tuple[Prompt[TriviaResponse], Session]: Prompt and session for execution.
         """
         session = Session()
 
+        # Allow subclasses to preprocess the request
+        request = self.preprocess_request(request, session)
+
         # Create workspace section with seeded files
         # This needs to be per-request because it references the session
         workspace_section = create_workspace_section(
@@ -383,6 +381,10 @@ def prepare(
         prompt.bind(QuestionParams(question=request.question))
         prompt.bind(EmptyParams())  # For sections without params (GameRules, Hints)
 
+        # Store for postprocess_response access
+        self._last_prompt = prompt
+        self._last_session = session
+
         return prompt, session
 
     def execute(
@@ -397,8 +399,8 @@ def execute(
     ) -> tuple[PromptResponse[TriviaResponse], Session]:
         """Execute a trivia request with preprocessing and postprocessing.
 
-        Overrides the parent MainLoop.execute() to apply preprocess_request()
-        before execution and postprocess_response() after execution.
+        Preprocessing happens in prepare() via preprocess_request().
+        Postprocessing happens after execution via postprocess_response().
 
         Args:
             request: TriviaRequest containing the question to process.
@@ -411,12 +413,9 @@ def execute(
         Returns:
             tuple[PromptResponse[TriviaResponse], Session]: Response and session.
         """
-        # Apply preprocessing to the request
-        preprocessed_request = self.preprocess_request(request)
-
-        # Execute with the preprocessed request
+        # Execute (prepare() is called internally by parent, which calls preprocess_request)
         prompt_response, session = super().execute(
-            preprocessed_request,
+            request,
             budget=budget,
             deadline=deadline,
             resources=resources,
@@ -426,8 +425,10 @@ def execute(
 
         # Apply postprocessing to the response output (if present)
         output = prompt_response.output
-        if output is not None:
-            postprocessed_output = self.postprocess_response(output)
+        if output is not None and self._last_prompt is not None:
+            postprocessed_output = self.postprocess_response(
+                output, self._last_prompt, self._last_session or session
+            )
             return prompt_response.update(output=postprocessed_output), session  # type: ignore[return-value]
 
         return prompt_response, session

tests/trivia_agent/test_worker.py

@@ -246,6 +246,8 @@ def test_preprocess_request_returns_unchanged_by_default(
         fake_mailboxes: TriviaMailboxes,
     ) -> None:
         """Test that preprocess_request returns request unchanged by default."""
+        from weakincentives.runtime import Session
+
         mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
 
         loop = TriviaAgentLoop(
@@ -254,7 +256,8 @@ def test_preprocess_request_returns_unchanged_by_default(
         )
 
         request = TriviaRequest(question="test question")
-        result = loop.preprocess_request(request)
+        session = Session()
+        result = loop.preprocess_request(request, session)
 
         assert result is request
 
@@ -263,9 +266,10 @@ def test_preprocess_request_can_be_overridden(
         fake_mailboxes: TriviaMailboxes,
     ) -> None:
         """Test that preprocess_request can be overridden in subclass."""
+        from weakincentives.runtime import Session
 
         class CustomLoop(TriviaAgentLoop):
-            def preprocess_request(self, request: TriviaRequest) -> TriviaRequest:
+            def preprocess_request(self, request: TriviaRequest, session: Session) -> TriviaRequest:
                 return TriviaRequest(question=request.question.upper())
 
         mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
@@ -276,7 +280,8 @@ def preprocess_request(self, request: TriviaRequest) -> TriviaRequest:
         )
 
         request = TriviaRequest(question="hello world")
-        result = loop.preprocess_request(request)
+        session = Session()
+        result = loop.preprocess_request(request, session)
 
         assert result.question == "HELLO WORLD"
 
@@ -289,6 +294,8 @@ def test_postprocess_response_returns_unchanged_by_default(
         fake_mailboxes: TriviaMailboxes,
     ) -> None:
         """Test that postprocess_response returns response unchanged by default."""
+        from weakincentives.runtime import Session
+
         mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
 
         loop = TriviaAgentLoop(
@@ -297,7 +304,9 @@ def test_postprocess_response_returns_unchanged_by_default(
         )
 
         response = TriviaResponse(answer="42")
-        result = loop.postprocess_response(response)
+        mock_prompt = MagicMock()
+        session = Session()
+        result = loop.postprocess_response(response, mock_prompt, session)
 
         assert result is response
 
@@ -306,9 +315,16 @@ def test_postprocess_response_can_be_overridden(
         fake_mailboxes: TriviaMailboxes,
     ) -> None:
         """Test that postprocess_response can be overridden in subclass."""
+        from weakincentives import Prompt
+        from weakincentives.runtime import Session
 
         class CustomLoop(TriviaAgentLoop):
-            def postprocess_response(self, response: TriviaResponse) -> TriviaResponse:
+            def postprocess_response(
+                self,
+                response: TriviaResponse,
+                prompt: Prompt[TriviaResponse],
+                session: Session,
+            ) -> TriviaResponse:
                 return TriviaResponse(answer=f"Answer: {response.answer}")
 
         mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
@@ -319,23 +335,28 @@ def postprocess_response(self, response: TriviaResponse) -> TriviaResponse:
         )
 
         response = TriviaResponse(answer="42")
-        result = loop.postprocess_response(response)
+        mock_prompt = MagicMock()
+        session = Session()
+        result = loop.postprocess_response(response, mock_prompt, session)
 
         assert result.answer == "Answer: 42"
 
 
 class TestTriviaAgentLoopExecute:
     """Tests for TriviaAgentLoop.execute() with preprocessing/postprocessing."""
 
-    def test_execute_calls_preprocess_request(
+    def test_execute_calls_preprocess_request_via_prepare(
         self,
         fake_mailboxes: TriviaMailboxes,
     ) -> None:
-        """Test that execute() calls preprocess_request."""
-        from weakincentives.runtime import MainLoop
+        """Test that prepare() calls preprocess_request during execution."""
+        from weakincentives.runtime import Session
+
+        preprocess_calls: list[TriviaRequest] = []
 
         class CustomLoop(TriviaAgentLoop):
-            def preprocess_request(self, request: TriviaRequest) -> TriviaRequest:
+            def preprocess_request(self, request: TriviaRequest, session: Session) -> TriviaRequest:
+                preprocess_calls.append(request)
                 return TriviaRequest(question=request.question.strip())
 
         mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
@@ -345,32 +366,33 @@ def preprocess_request(self, request: TriviaRequest) -> TriviaRequest:
             requests=fake_mailboxes.requests,
         )
 
-        mock_response = MagicMock()
-        mock_response.output = TriviaResponse(answer="42")
-        mock_session = MagicMock()
-
-        captured_requests: list[TriviaRequest] = []
-
-        def capture_execute(self_arg, request, **kwargs):
-            captured_requests.append(request)
-            return (mock_response, mock_session)
-
-        with patch.object(MainLoop, "execute", capture_execute):
-            request = TriviaRequest(question="  What is the answer?  ")
-            loop.execute(request)
+        request = TriviaRequest(question="  What is the answer?  ")
+        prompt, session = loop.prepare(request)
 
-            assert len(captured_requests) == 1
-            assert captured_requests[0].question == "What is the answer?"
+        assert len(preprocess_calls) == 1
+        assert preprocess_calls[0].question == "  What is the answer?  "
+        # The prompt should have the preprocessed question bound
+        rendered = str(prompt.render())
+        assert "What is the answer?" in rendered
 
     def test_execute_calls_postprocess_response(
         self,
         fake_mailboxes: TriviaMailboxes,
     ) -> None:
         """Test that execute() calls postprocess_response."""
-        from weakincentives.runtime import MainLoop
+        from weakincentives import Prompt
+        from weakincentives.runtime import MainLoop, Session
+
+        postprocess_calls: list[tuple[TriviaResponse, bool, bool]] = []
 
         class CustomLoop(TriviaAgentLoop):
-            def postprocess_response(self, response: TriviaResponse) -> TriviaResponse:
+            def postprocess_response(
+                self,
+                response: TriviaResponse,
+                prompt: Prompt[TriviaResponse],
+                session: Session,
+            ) -> TriviaResponse:
+                postprocess_calls.append((response, prompt is not None, session is not None))
                 return TriviaResponse(answer=response.answer.strip())
 
         mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
@@ -385,13 +407,18 @@ def postprocess_response(self, response: TriviaResponse) -> TriviaResponse:
         mock_response.update = MagicMock(return_value=mock_response)
         mock_session = MagicMock()
 
+        # Call prepare() first to set up _last_prompt/_last_session
+        request = TriviaRequest(question="What is the answer?")
+        loop.prepare(request)
+
         with patch.object(MainLoop, "execute", return_value=(mock_response, mock_session)):
-            request = TriviaRequest(question="What is the answer?")
             loop.execute(request)
 
+            assert len(postprocess_calls) == 1
+            assert postprocess_calls[0][0].answer == "  42  "
+            assert postprocess_calls[0][1] is True  # prompt was passed
+            assert postprocess_calls[0][2] is True  # session was passed
             mock_response.update.assert_called_once()
-            call_kwargs = mock_response.update.call_args.kwargs
-            assert call_kwargs["output"].answer == "42"
 
     def test_execute_skips_postprocessing_when_output_is_none(
         self,