Add request preprocessing and response postprocessing to MainLoop

Add two overridable methods for transforming requests before agent
processing and responses before returning to callers:

- preprocess_request(request) - override to transform incoming requests
- postprocess_response(response) - override to transform outgoing responses

Both methods return their input unchanged by default. Subclass TriviaAgentLoop
and override these methods to implement custom preprocessing/postprocessing logic.

Author: claude

src/trivia_agent/worker.py

@@ -14,20 +14,21 @@
 
 import os
 import sys
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from dataclasses import field
 from datetime import UTC, datetime, timedelta
 from pathlib import Path
 from typing import TYPE_CHECKING, TextIO
 
-from weakincentives import FrozenDataclass, Prompt
+from weakincentives import Budget, FrozenDataclass, Prompt, PromptResponse
 from weakincentives.adapters import ProviderAdapter
 from weakincentives.adapters.claude_agent_sdk import ClaudeAgentWorkspaceSection, HostMount
 from weakincentives.deadlines import Deadline
 from weakincentives.debug.bundle import BundleConfig
 from weakincentives.prompt import PromptTemplate
 from weakincentives.prompt.overrides import LocalPromptOverridesStore, PromptOverridesStore
 from weakincentives.runtime import (
+    Heartbeat,
     LoopGroup,
     MainLoop,
     MainLoopConfig,
@@ -204,6 +205,10 @@ class TriviaAgentLoop(MainLoop[TriviaRequest, TriviaResponse]):
       customization of prompt content without code changes
     - **Evaluation integration**: Compatible with EvalLoop for running
       evaluations with session-aware scoring
+    - **Request preprocessing**: Override preprocess_request() to transform
+      requests before agent processing
+    - **Response postprocessing**: Override postprocess_response() to transform
+      responses before returning to callers
 
     To use this loop, instantiate it with an adapter and mailbox, then either:
     1. Call run() directly for single-threaded processing
@@ -265,6 +270,34 @@ def __init__(
         self._base_template = build_prompt_template()
         self._overrides_store = overrides_store
 
+    def preprocess_request(self, request: TriviaRequest) -> TriviaRequest:
+        """Transform request before agent processing.
+
+        Override this method in subclasses to implement custom preprocessing
+        logic such as validation, normalization, or enrichment.
+
+        Args:
+            request: The incoming TriviaRequest.
+
+        Returns:
+            TriviaRequest: The preprocessed request.
+        """
+        return request
+
+    def postprocess_response(self, response: TriviaResponse) -> TriviaResponse:
+        """Transform response before returning to caller.
+
+        Override this method in subclasses to implement custom postprocessing
+        logic such as formatting, cleanup, or validation.
+
+        Args:
+            response: The TriviaResponse from the agent.
+
+        Returns:
+            TriviaResponse: The postprocessed response.
+        """
+        return response
+
     def prepare(
         self,
         request: TriviaRequest,
@@ -277,6 +310,10 @@ 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
@@ -348,6 +385,53 @@ def prepare(
 
         return prompt, session
 
+    def execute(
+        self,
+        request: TriviaRequest,
+        *,
+        budget: Budget | None = None,
+        deadline: Deadline | None = None,
+        resources: Mapping[type[object], object] | None = None,
+        heartbeat: Heartbeat | None = None,
+        experiment: Experiment | None = None,
+    ) -> 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.
+
+        Args:
+            request: TriviaRequest containing the question to process.
+            budget: Optional Budget for token/cost limits.
+            deadline: Optional Deadline for time limits.
+            resources: Optional mapping of resource types to instances.
+            heartbeat: Optional Heartbeat for progress reporting.
+            experiment: Optional Experiment for evaluation runs.
+
+        Returns:
+            tuple[PromptResponse[TriviaResponse], Session]: Response and session.
+        """
+        # Apply preprocessing to the request
+        preprocessed_request = self.preprocess_request(request)
+
+        # Execute with the preprocessed request
+        prompt_response, session = super().execute(
+            preprocessed_request,
+            budget=budget,
+            deadline=deadline,
+            resources=resources,
+            heartbeat=heartbeat,
+            experiment=experiment,
+        )
+
+        # Apply postprocessing to the response output (if present)
+        output = prompt_response.output
+        if output is not None:
+            postprocessed_output = self.postprocess_response(output)
+            return prompt_response.update(output=postprocessed_output), session  # type: ignore[return-value]
+
+        return prompt_response, session
+
 
 @FrozenDataclass()
 class TriviaRuntime:

tests/trivia_agent/test_worker.py

@@ -238,6 +238,189 @@ def test_prepare_seeds_overrides_store(
         assert call_args.kwargs.get("tag") == "latest"
 
 
+class TestTriviaAgentLoopPreprocessing:
+    """Tests for TriviaAgentLoop.preprocess_request() method."""
+
+    def test_preprocess_request_returns_unchanged_by_default(
+        self,
+        fake_mailboxes: TriviaMailboxes,
+    ) -> None:
+        """Test that preprocess_request returns request unchanged by default."""
+        mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
+
+        loop = TriviaAgentLoop(
+            adapter=mock_adapter,
+            requests=fake_mailboxes.requests,
+        )
+
+        request = TriviaRequest(question="test question")
+        result = loop.preprocess_request(request)
+
+        assert result is request
+
+    def test_preprocess_request_can_be_overridden(
+        self,
+        fake_mailboxes: TriviaMailboxes,
+    ) -> None:
+        """Test that preprocess_request can be overridden in subclass."""
+
+        class CustomLoop(TriviaAgentLoop):
+            def preprocess_request(self, request: TriviaRequest) -> TriviaRequest:
+                return TriviaRequest(question=request.question.upper())
+
+        mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
+
+        loop = CustomLoop(
+            adapter=mock_adapter,
+            requests=fake_mailboxes.requests,
+        )
+
+        request = TriviaRequest(question="hello world")
+        result = loop.preprocess_request(request)
+
+        assert result.question == "HELLO WORLD"
+
+
+class TestTriviaAgentLoopPostprocessing:
+    """Tests for TriviaAgentLoop.postprocess_response() method."""
+
+    def test_postprocess_response_returns_unchanged_by_default(
+        self,
+        fake_mailboxes: TriviaMailboxes,
+    ) -> None:
+        """Test that postprocess_response returns response unchanged by default."""
+        mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
+
+        loop = TriviaAgentLoop(
+            adapter=mock_adapter,
+            requests=fake_mailboxes.requests,
+        )
+
+        response = TriviaResponse(answer="42")
+        result = loop.postprocess_response(response)
+
+        assert result is response
+
+    def test_postprocess_response_can_be_overridden(
+        self,
+        fake_mailboxes: TriviaMailboxes,
+    ) -> None:
+        """Test that postprocess_response can be overridden in subclass."""
+
+        class CustomLoop(TriviaAgentLoop):
+            def postprocess_response(self, response: TriviaResponse) -> TriviaResponse:
+                return TriviaResponse(answer=f"Answer: {response.answer}")
+
+        mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
+
+        loop = CustomLoop(
+            adapter=mock_adapter,
+            requests=fake_mailboxes.requests,
+        )
+
+        response = TriviaResponse(answer="42")
+        result = loop.postprocess_response(response)
+
+        assert result.answer == "Answer: 42"
+
+
+class TestTriviaAgentLoopExecute:
+    """Tests for TriviaAgentLoop.execute() with preprocessing/postprocessing."""
+
+    def test_execute_calls_preprocess_request(
+        self,
+        fake_mailboxes: TriviaMailboxes,
+    ) -> None:
+        """Test that execute() calls preprocess_request."""
+        from weakincentives.runtime import MainLoop
+
+        class CustomLoop(TriviaAgentLoop):
+            def preprocess_request(self, request: TriviaRequest) -> TriviaRequest:
+                return TriviaRequest(question=request.question.strip())
+
+        mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
+
+        loop = CustomLoop(
+            adapter=mock_adapter,
+            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)
+
+            assert len(captured_requests) == 1
+            assert captured_requests[0].question == "What is the answer?"
+
+    def test_execute_calls_postprocess_response(
+        self,
+        fake_mailboxes: TriviaMailboxes,
+    ) -> None:
+        """Test that execute() calls postprocess_response."""
+        from weakincentives.runtime import MainLoop
+
+        class CustomLoop(TriviaAgentLoop):
+            def postprocess_response(self, response: TriviaResponse) -> TriviaResponse:
+                return TriviaResponse(answer=response.answer.strip())
+
+        mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
+
+        loop = CustomLoop(
+            adapter=mock_adapter,
+            requests=fake_mailboxes.requests,
+        )
+
+        mock_response = MagicMock()
+        mock_response.output = TriviaResponse(answer="  42  ")
+        mock_response.update = MagicMock(return_value=mock_response)
+        mock_session = MagicMock()
+
+        with patch.object(MainLoop, "execute", return_value=(mock_response, mock_session)):
+            request = TriviaRequest(question="What is the answer?")
+            loop.execute(request)
+
+            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,
+        fake_mailboxes: TriviaMailboxes,
+    ) -> None:
+        """Test that execute() skips postprocessing when output is None."""
+        from weakincentives.runtime import MainLoop
+
+        mock_adapter: ProviderAdapter[TriviaResponse] = MagicMock()
+
+        loop = TriviaAgentLoop(
+            adapter=mock_adapter,
+            requests=fake_mailboxes.requests,
+        )
+
+        mock_response = MagicMock()
+        mock_response.output = None
+        mock_response.update = MagicMock()
+        mock_session = MagicMock()
+
+        with patch.object(MainLoop, "execute", return_value=(mock_response, mock_session)):
+            request = TriviaRequest(question="What is the answer?")
+            result_response, result_session = loop.execute(request)
+
+            mock_response.update.assert_not_called()
+            assert result_response is mock_response
+            assert result_session is mock_session
+
+
 class TestTriviaRuntime:
     """Tests for TriviaRuntime dataclass."""