feat(code): let the rubric grader inspect working-directory files

libs/code/deepagents_code/reliable_rubric.py

@@ -3,13 +3,27 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, NotRequired
 
 import httpx
-from deepagents.middleware.rubric import GraderResponse, RubricMiddleware, RubricState
+from deepagents.middleware.rubric import (
+    RUBRIC_GRADER_MESSAGE_SOURCE,
+    GraderResponse,
+    RubricMiddleware,
+    RubricState,
+    _strategy_from_result,  # noqa: PLC2701
+)
+from langchain.agents.middleware.types import AgentMiddleware, AgentState, hook_config
+from langchain_core.messages import HumanMessage
+from langgraph.errors import GraphBubbleUp
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
+    from collections.abc import Callable, Iterator, Sequence
+
+    from deepagents.middleware.rubric import RubricEvaluation
+    from langchain_core.language_models import BaseChatModel
+    from langchain_core.tools import BaseTool
+    from langgraph.runtime import Runtime
 
 logger = logging.getLogger(__name__)
 
@@ -63,35 +77,245 @@ def _is_transient_grader_transport_error(exc: BaseException) -> bool:
     return False
 
 
+class RubricGraderState(AgentState[GraderResponse]):
+    """Nested-grader state used to scope verification-tool budgets."""
+
+    rubric_grading_operation_id: NotRequired[str]
+
+
 class ReliableRubricMiddleware(RubricMiddleware):
-    """Retry one transient grader transport failure without rerunning agent work.
+    """Run a context-aware nested grader and retry transient transport failures.
 
-    The retry re-invokes only the grader sub-agent, never the task agent, so it
-    relies on the grader's own tools being read-only/idempotent. A second
-    failure — transient or not — propagates to the base middleware, which
-    surfaces it as a `grader_error` result rather than a silent success.
+    The nested grader receives Deep Agents Code's verification middleware and
+    runtime context without requiring those application-specific capabilities in
+    the SDK's `RubricMiddleware`. A transport retry re-invokes only the grader,
+    never the task agent, so grader tools must be read-only or idempotent.
     """
 
-    def _grade(self, state: RubricState, iteration: int) -> GraderResponse:
+    def __init__(  # noqa: D107
+        self,
+        *,
+        model: str | BaseChatModel,
+        system_prompt: str | None = None,
+        tools: Sequence[BaseTool] | None = None,
+        grader_middleware: Sequence[AgentMiddleware[Any, Any]] | None = None,
+        grader_context_schema: type[Any] | None = None,
+        max_iterations: int = 3,
+        on_evaluation: Callable[[RubricEvaluation], None] | None = None,
+    ) -> None:
+        super().__init__(
+            model=model,
+            system_prompt=system_prompt,
+            tools=tools,
+            max_iterations=max_iterations,
+            on_evaluation=on_evaluation,
+        )
+        self._grader_middleware = list(grader_middleware or ())
+        self._grader_context_schema = grader_context_schema
+
+    @hook_config(can_jump_to=["model"])
+    def after_agent(
+        self,
+        state: RubricState,
+        runtime: Runtime[Any],
+    ) -> dict[str, Any] | None:
+        """Grade synchronously while preserving nested graph interrupts.
+
+        Returns:
+            The rubric state update, or `None` when no rubric is active.
+
+        Raises:
+            GraphBubbleUp: If the nested grader pauses or otherwise bubbles control.
+        """
+        prep = self._prepare_evaluation(state, runtime)
+        if prep is None:
+            return None
+        grading_run_id, iteration = prep
+
+        try:
+            graded = self._grade(
+                state,
+                iteration,
+                context=getattr(runtime, "context", None),
+            )
+        except GraphBubbleUp:
+            raise
+        except Exception as exc:  # noqa: BLE001
+            return self._handle_grader_exception(
+                runtime,
+                state,
+                grading_run_id,
+                iteration,
+                exc,
+            )
+
+        return self._finalize_evaluation(
+            graded,
+            state,
+            runtime,
+            grading_run_id,
+            iteration,
+        )
+
+    async def aafter_agent(
+        self,
+        state: RubricState,
+        runtime: Runtime[Any],
+    ) -> dict[str, Any] | None:
+        """Grade asynchronously while preserving nested graph interrupts.
+
+        Returns:
+            The rubric state update, or `None` when no rubric is active.
+
+        Raises:
+            GraphBubbleUp: If the nested grader pauses or otherwise bubbles control.
+        """
+        prep = self._prepare_evaluation(state, runtime)
+        if prep is None:
+            return None
+        grading_run_id, iteration = prep
+
+        try:
+            graded = await self._agrade(
+                state,
+                iteration,
+                context=getattr(runtime, "context", None),
+            )
+        except GraphBubbleUp:
+            raise
+        except Exception as exc:  # noqa: BLE001
+            return self._handle_grader_exception(
+                runtime,
+                state,
+                grading_run_id,
+                iteration,
+                exc,
+            )
+
+        return self._finalize_evaluation(
+            graded,
+            state,
+            runtime,
+            grading_run_id,
+            iteration,
+        )
+
+    def _ensure_grader(self) -> Any:  # noqa: ANN401
+        if self._grader is not None:
+            return self._grader
+
+        from deepagents._models import (  # noqa: PLC2701
+            resolve_model,
+        )
+        from langchain.agents import create_agent
+
+        resolved_model = resolve_model(self._model)
+        self._resolved_model = resolved_model
+        self._grader = create_agent(
+            model=resolved_model,
+            system_prompt=self._system_prompt,
+            tools=self._tools,
+            middleware=self._grader_middleware,
+            name=RUBRIC_GRADER_MESSAGE_SOURCE,
+            response_format=GraderResponse,
+            state_schema=RubricGraderState,
+            context_schema=self._grader_context_schema,
+        )
+        return self._grader
+
+    def _grader_input(
+        self,
+        state: RubricState,
+        iteration: int,
+    ) -> dict[str, Any]:
+        """Build nested-grader input with a stable verification-operation ID.
+
+        Returns:
+            The nested grader's input state.
+        """
+        grading_run_id = state.get("_current_grading_run_id") or "untracked"
+        payload = self._build_grader_payload(state, iteration)
+        return {
+            "messages": [HumanMessage(content=payload)],
+            "rubric_grading_operation_id": f"{grading_run_id}:{iteration}",
+        }
+
+    def _grade_once(
+        self,
+        state: RubricState,
+        iteration: int,
+        *,
+        context: object | None,
+    ) -> GraderResponse:
+        grader = self._ensure_grader()
+        metadata = self._grader_trace_metadata()
+        self._record_grader_trace_metadata(metadata)
+        result = grader.invoke(
+            self._grader_input(state, iteration),
+            config=self._grader_invocation_config(metadata),
+            context=context,
+        )
+        self._record_grader_trace_metadata(
+            self._grader_trace_metadata(
+                effective_strategy=_strategy_from_result(result),
+            )
+        )
+        return self._extract_graded(result)
+
+    async def _agrade_once(
+        self,
+        state: RubricState,
+        iteration: int,
+        *,
+        context: object | None,
+    ) -> GraderResponse:
+        grader = self._ensure_grader()
+        metadata = self._grader_trace_metadata()
+        self._record_grader_trace_metadata(metadata)
+        result = await grader.ainvoke(
+            self._grader_input(state, iteration),
+            config=self._grader_invocation_config(metadata),
+            context=context,
+        )
+        self._record_grader_trace_metadata(
+            self._grader_trace_metadata(
+                effective_strategy=_strategy_from_result(result),
+            )
+        )
+        return self._extract_graded(result)
+
+    def _grade(
+        self,
+        state: RubricState,
+        iteration: int,
+        *,
+        context: object | None = None,
+    ) -> GraderResponse:
         try:
-            return super()._grade(state, iteration)
+            return self._grade_once(state, iteration, context=context)
         except Exception as exc:
             if not _is_transient_grader_transport_error(exc):
                 raise
             logger.warning(
                 "Rubric grader transport failed; retrying grading once",
                 exc_info=True,
             )
-        return super()._grade(state, iteration)
+        return self._grade_once(state, iteration, context=context)
 
-    async def _agrade(self, state: RubricState, iteration: int) -> GraderResponse:
+    async def _agrade(
+        self,
+        state: RubricState,
+        iteration: int,
+        *,
+        context: object | None = None,
+    ) -> GraderResponse:
         try:
-            return await super()._agrade(state, iteration)
+            return await self._agrade_once(state, iteration, context=context)
         except Exception as exc:
             if not _is_transient_grader_transport_error(exc):
                 raise
             logger.warning(
                 "Rubric grader transport failed; retrying grading once",
                 exc_info=True,
             )
-        return await super()._agrade(state, iteration)
+        return await self._agrade_once(state, iteration, context=context)

libs/code/deepagents_code/server_graph.py

@@ -149,16 +149,16 @@ def _criteria_context_tools(
     tools: list[Any],
     mcp_tools: list[Any],
 ) -> list[Any]:
-    """Select external context tools from the normal agent tool list.
+    """Select read-only external tools for criteria drafting and rubric grading.
 
     Args:
         tools: Main agent tools in execution order.
         mcp_tools: Exact tool objects returned by MCP discovery.
 
     Returns:
-        External context tools available to criteria generation. MCP tools are
-        included only when their protocol annotations explicitly declare them
-        read-only.
+        External context tools available to criteria generation and grading.
+        MCP tools are included only when their protocol annotations explicitly
+        declare them read-only.
     """
     from deepagents_code.tools import fetch_url, web_search
 
@@ -223,6 +223,7 @@ async def _make_graph() -> Any:  # noqa: ANN401
     result.apply_to_settings()
 
     tools, mcp_server_info, mcp_tools = await _build_tools(config, project_context)
+    read_only_context_tools = _criteria_context_tools(tools, mcp_tools)
 
     # Create sandbox backend if a sandbox provider is configured.
     # The context manager is created here in the factory, but its reference is
@@ -322,7 +323,8 @@ def _create_cli_agent_sync() -> Any:  # noqa: ANN401
             cwd=project_context.user_cwd if project_context is not None else config.cwd,
             project_context=project_context,
             async_subagents=async_subagents,
-            goal_criteria_tools=_criteria_context_tools(tools, mcp_tools),
+            goal_criteria_tools=read_only_context_tools,
+            rubric_grader_tools=read_only_context_tools,
         )
         return agent