zjcode 0.0.1__py3-none-any.whl

deepagents_code/reliable_rubric.py

@@ -0,0 +1,97 @@
+"""Rubric middleware retries for transient grader transport failures."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import httpx
+from deepagents.middleware.rubric import GraderResponse, RubricMiddleware, RubricState
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+logger = logging.getLogger(__name__)
+
+
+def _exception_chain(exc: BaseException) -> Iterator[BaseException]:
+    """Yield an exception, its explicit/implicit causes, and group members once.
+
+    Descends into `BaseExceptionGroup` members as well as `__cause__` and
+    `__context__`, so a transient transport error wrapped in an async task group
+    is still discovered. Each exception is yielded at most once.
+    """
+    pending = [exc]
+    seen: set[int] = set()
+    while pending:
+        current = pending.pop()
+        if id(current) in seen:
+            continue
+        seen.add(id(current))
+        yield current
+        if isinstance(current, BaseExceptionGroup):
+            pending.extend(current.exceptions)
+        if current.__cause__ is not None:
+            pending.append(current.__cause__)
+        elif current.__context__ is not None:
+            pending.append(current.__context__)
+
+
+def _is_transient_grader_transport_error(exc: BaseException) -> bool:
+    """Return whether a grader failure is a retryable transport/read error.
+
+    Matches response-read faults (`httpx`/`httpcore` `ReadError`) and
+    response-framing faults (`RemoteProtocolError`, aiohttp
+    `TransferEncodingError`). Connect/timeout errors are intentionally excluded
+    so only mid-response transport failures trigger the retry.
+    """
+    for current in _exception_chain(exc):
+        if isinstance(current, (httpx.ReadError, httpx.RemoteProtocolError)):
+            return True
+        error_type = type(current)
+        if error_type.__module__.startswith("httpcore") and error_type.__name__ in {
+            "ReadError",
+            "RemoteProtocolError",
+        }:
+            return True
+        if (
+            error_type.__module__ == "aiohttp.http_exceptions"
+            and error_type.__name__ == "TransferEncodingError"
+            and "Not enough data to satisfy transfer length header" in str(current)
+        ):
+            return True
+    return False
+
+
+class ReliableRubricMiddleware(RubricMiddleware):
+    """Retry one transient grader transport failure without rerunning agent work.
+
+    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.
+    """
+
+    def _grade(self, state: RubricState, iteration: int) -> GraderResponse:
+        try:
+            return super()._grade(state, iteration)
+        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)
+
+    async def _agrade(self, state: RubricState, iteration: int) -> GraderResponse:
+        try:
+            return await super()._agrade(state, iteration)
+        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)

deepagents_code/resume_state.py

@@ -0,0 +1,237 @@
+"""Schema and middleware for per-checkpoint state restored when resuming.
+
+`ResumeState` declares several checkpointed, schema-private channels. They fall
+into two groups with *different* write paths:
+
+Written from inside the graph on successful model turns:
+
+- `_context_tokens` — total context tokens from the latest
+    `AIMessage.usage_metadata`, written by `ResumeStateMiddleware.after_model`.
+    Powers `/tokens` and the status bar.
+- `_model_spec` / `_model_params` — the model and invocation params effectively
+    in use for the turn, written by `ConfigurableModelMiddleware` after a
+    successful model call. Lets `dcode -r` restore the model the resumed thread
+    was actually using instead of falling back to the user's global default.
+
+Written primarily by the TUI client, via `aupdate_state` (see
+`DeepAgentsApp._persist_goal_rubric_state`) — these are user/agent-owned. Most
+have no model-node write site; the two exceptions are called out below:
+
+- `_goal_objective` / `_goal_status` / `_goal_rubric` / `_goal_status_note` —
+    the accepted goal and its lifecycle status. `_goal_objective`/`_goal_rubric`
+    are client-only, but `_goal_status`/`_goal_status_note` are *also* written
+    from inside the graph by the agent's `update_goal` tool.
+- `_pending_goal_completion_note` — an agent-requested completion awaiting the
+    post-turn rubric result and, when needed, user approval.
+- `_sticky_rubric` — the TUI-owned persistent rubric. This is separate from
+    the public `rubric` graph input so one-shot rubric turns can be checkpointed
+    without being restored as sticky state.
+- `_pending_goal_objective` / `_pending_goal_rubric` / `_pending_goal_kind` — a
+    proposed goal or amendment awaiting user acceptance of its criteria.
+
+All of these are facts the CLI reads back from `state_values` on thread resume
+so it can rehydrate the session without replaying or re-tokenizing history.
+
+The model-turn channels are persisted from inside the graph (rather than via a
+separate client-side `aupdate_state` call) so the write rides the same checkpoint
+as the model response and avoids creating a standalone `UpdateState` run in
+LangSmith. Because they are versioned channel state, resuming a specific
+checkpoint yields the values as of *that* checkpoint — not a thread-level
+aggregate. The goal/rubric channels are client-written because the user sets
+them outside any model turn (except `_goal_status`/`_goal_status_note`, which the
+`update_goal` tool also writes from inside the graph). Both paths work
+identically against local and remote (HTTP) graphs.
+"""
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Literal,
+    NotRequired,
+    cast,
+    get_args,
+)
+
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+    ContextT,
+    PrivateStateAttr,
+)
+from langchain_core.messages import AIMessage
+
+if TYPE_CHECKING:
+    from langgraph.runtime import Runtime
+
+GoalStatus = Literal["active", "paused", "blocked", "complete"]
+"""Lifecycle status of a TUI-owned goal.
+
+`active` and `blocked` are unfinished working states, `paused` preserves the goal
+without driving work, and `complete` is terminal. A blocked goal is still
+considered actionable (`active=True`) by `get_goal`, whereas a paused goal is
+unfinished but reports `active=False`.
+"""
+
+GoalProposalKind = Literal["create", "amend"]
+"""Whether a pending review creates a goal or amends the current one."""
+
+_GOAL_STATUS_VALUES: frozenset[str] = frozenset(get_args(GoalStatus))
+_GOAL_PROPOSAL_KIND_VALUES: frozenset[str] = frozenset(get_args(GoalProposalKind))
+
+
+def coerce_goal_proposal_kind(value: object) -> GoalProposalKind | None:
+    """Narrow a persisted proposal kind to a known value.
+
+    Args:
+        value: Raw value read from checkpoint state.
+
+    Returns:
+        The recognized proposal kind, otherwise `None`.
+    """
+    if isinstance(value, str) and value in _GOAL_PROPOSAL_KIND_VALUES:
+        return cast("GoalProposalKind", value)
+    return None
+
+
+def coerce_goal_status(value: object) -> GoalStatus | None:
+    """Narrow a persisted goal-status value to a known `GoalStatus`.
+
+    A corrupt or forward-version checkpoint can carry an unexpected status
+    string (or a non-string). Coercing to `None` rather than passing the raw
+    value through keeps the `GoalStatus` `Literal` load-bearing on the read
+    path, so an unknown status is treated as "no goal status" instead of a
+    silently active goal. Resume/restore callers should log the discard
+    separately so it is surfaced rather than dropped; the model-read path
+    (`_goal_snapshot`) intentionally treats an unknown status as `active`
+    without logging.
+
+    Args:
+        value: Raw value read from checkpoint state.
+
+    Returns:
+        The value when it is a recognized `GoalStatus`, otherwise `None`.
+    """
+    if isinstance(value, str) and value in _GOAL_STATUS_VALUES:
+        return cast("GoalStatus", value)
+    return None
+
+
+class GoalRubricChannels(AgentState):
+    """Goal/rubric state channels shared by every schema that touches them.
+
+    Declared once here so each schema that carries these channels —
+    `ResumeState` and `goal_tools.GoalToolState` — inherits the *same*
+    `PrivateStateAttr`-marked annotations. Middleware state schemas merge with
+    later entries winning, so an independent re-declaration that dropped the
+    `PrivateStateAttr` marker would override these and leak the field into the
+    public graph input/output schema. Inheriting from a single base makes that
+    drift unrepresentable.
+    """
+
+    _goal_objective: Annotated[NotRequired[str | None], PrivateStateAttr]
+    """Accepted goal objective restored by the TUI on resume."""
+
+    _goal_status: Annotated[NotRequired[GoalStatus | None], PrivateStateAttr]
+    """Goal lifecycle status (`active`, `paused`, `blocked`, `complete`, or `None`)."""
+
+    _goal_rubric: Annotated[NotRequired[str | None], PrivateStateAttr]
+    """Accepted rubric associated with `_goal_objective`."""
+
+    _goal_status_note: Annotated[NotRequired[str | None], PrivateStateAttr]
+    """Evidence or blocker note recorded by `update_goal`."""
+
+    _pending_goal_completion_note: Annotated[NotRequired[str | None], PrivateStateAttr]
+    """Completion evidence awaiting rubric and user approval."""
+
+    _sticky_rubric: Annotated[NotRequired[str | None], PrivateStateAttr]
+    """Persistent rubric owned by the TUI, distinct from graph input `rubric`."""
+
+
+class ResumeState(GoalRubricChannels):
+    """Extends agent state with per-checkpoint facts restored on resume.
+
+    Inherits the shared goal/rubric channels from `GoalRubricChannels` and adds
+    the channels unique to resume: the after-model token/spec facts and the
+    pending-goal proposal awaiting acceptance.
+    """
+
+    _context_tokens: Annotated[NotRequired[int], PrivateStateAttr]
+    """Total context tokens reported by the model's last `usage_metadata`."""
+
+    _model_spec: Annotated[NotRequired[str], PrivateStateAttr]
+    """`provider:model` spec effectively in use for the latest turn."""
+
+    _model_params: Annotated[NotRequired[dict[str, Any] | None], PrivateStateAttr]
+    """Invocation params effectively in use for the latest turn."""
+
+    _pending_goal_objective: Annotated[NotRequired[str | None], PrivateStateAttr]
+    """Goal objective awaiting acceptance of proposed criteria."""
+
+    _pending_goal_rubric: Annotated[NotRequired[str | None], PrivateStateAttr]
+    """Proposed criteria awaiting user acceptance."""
+
+    _pending_goal_kind: Annotated[
+        NotRequired[GoalProposalKind | None], PrivateStateAttr
+    ]
+    """Whether the pending review creates or amends a goal."""
+
+
+def _extract_context_tokens(message: AIMessage) -> int | None:
+    """Return the context-token count from an AI message, or `None` if absent.
+
+    Prefers `input_tokens + output_tokens` when both are reported; falls back
+    to `total_tokens` when the model only provides the aggregate.
+    """
+    usage = getattr(message, "usage_metadata", None)
+    if not usage:
+        return None
+    input_toks = usage.get("input_tokens", 0) or 0
+    output_toks = usage.get("output_tokens", 0) or 0
+    if input_toks or output_toks:
+        return input_toks + output_toks
+    total = usage.get("total_tokens", 0) or 0
+    return total or None
+
+
+class ResumeStateMiddleware(AgentMiddleware[ResumeState, ContextT]):
+    """Persists per-checkpoint resume facts after each model call.
+
+    See the module docstring for why this rides the model node's checkpoint
+    instead of a separate `aupdate_state` (avoids a standalone `UpdateState`
+    run in LangSmith and works identically against remote graphs).
+    """
+
+    state_schema = ResumeState
+
+    def after_model(  # noqa: PLR6301  # AgentMiddleware hook must be an instance method.
+        self,
+        state: ResumeState,
+        runtime: Runtime[ContextT],  # noqa: ARG002
+    ) -> dict[str, Any] | None:
+        """Write `_context_tokens` for the latest turn.
+
+        Model metadata is written by `ConfigurableModelMiddleware` from the
+        actual request that completed successfully; this hook only records token
+        usage from the most recent `AIMessage.usage_metadata`.
+
+        Args:
+            state: Current agent state; only `messages` is inspected.
+            runtime: LangGraph runtime required by the middleware interface.
+
+        Returns:
+            State update with `_context_tokens`, or `None` when no token count is
+            available.
+        """
+        update: dict[str, Any] = {}
+
+        for msg in reversed(state.get("messages") or []):
+            if isinstance(msg, AIMessage):
+                tokens = _extract_context_tokens(msg)
+                if tokens is not None:
+                    update["_context_tokens"] = tokens
+                break
+
+        return update or None

deepagents_code/server_graph.py

@@ -0,0 +1,310 @@
+"""Server-side graph entry point for `langgraph dev`.
+
+This module is referenced by the generated `langgraph.json` and exposes a graph
+factory that the LangGraph server can load and serve.
+
+The graph is created by `make_graph()`, which reads configuration from
+`ServerConfig.from_env()` — the same dataclass the CLI uses to *write* the
+configuration via `ServerConfig.to_env()`. This shared schema ensures the two
+sides stay in sync.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import atexit
+import logging
+import sys
+from typing import TYPE_CHECKING, Any
+
+from deepagents_code._server_config import ServerConfig
+from deepagents_code._startup_error import (
+    STARTUP_ERROR_MARKER as _STARTUP_ERROR_MARKER,
+    emit_startup_failure,
+)
+from deepagents_code.project_utils import ProjectContext, get_server_project_context
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+logger = logging.getLogger(__name__)
+
+_sandbox_cm: Any = None
+_sandbox_backend: Any = None
+_mcp_session_manager: Any = None
+
+
+def _print_startup_error(message: str) -> None:
+    """Print a startup error for both humans and the parent app process.
+
+    Args:
+        message: Concise startup failure to surface in the parent process.
+    """
+    print(message, file=sys.stderr)  # noqa: T201  # stderr fallback for logs
+    print(  # noqa: T201  # machine-readable marker consumed by server.py
+        f"{_STARTUP_ERROR_MARKER}{message}",
+        file=sys.stderr,
+    )
+
+
+def _get_mcp_session_manager() -> Any:  # noqa: ANN401
+    """Return the process-wide MCP session manager singleton.
+
+    Sessions are bound to the langgraph dev server's event loop. Cleanup
+    therefore belongs to that loop's normal shutdown path, not `atexit` —
+    an atexit handler runs after the loop is already closed and cannot
+    await `AsyncExitStack.aclose()` safely. Subprocess handles held by
+    stdio transports are released when the Python process exits.
+    """
+    global _mcp_session_manager  # noqa: PLW0603
+
+    if _mcp_session_manager is None:
+        from deepagents_code.mcp_tools import MCPSessionManager
+
+        _mcp_session_manager = MCPSessionManager()
+
+    return _mcp_session_manager
+
+
+async def _build_tools(
+    config: ServerConfig,
+    project_context: ProjectContext | None,
+) -> tuple[list[Any], list[Any] | None]:
+    """Assemble the tool list based on server config.
+
+    Loads built-in tools (conditionally including web search when Tavily is
+    available) and MCP tools when enabled.
+
+    MCP discovery is awaited on the server's event loop: LangGraph invokes this
+    async factory on its running loop, so discovery must use `await` rather than
+    `asyncio.run` (which raises inside a running loop). `stateless=True` ensures
+    discovery only uses throwaway sessions, while the shared runtime session
+    manager binds real sessions lazily inside the server loop on first tool
+    invocation. MCP adapter imports are warmed in a worker thread inside
+    `_load_tools_from_config` (only when active servers exist) because first
+    import can perform blocking package-resource scans.
+
+    Args:
+        config: Deserialized server configuration.
+        project_context: Resolved project context for MCP discovery.
+
+    Returns:
+        Tuple of `(tools, mcp_server_info)`.
+
+    Raises:
+        FileNotFoundError: If the MCP config file is not found.
+        RuntimeError: If MCP tool loading fails.
+    """
+    from deepagents_code.config import settings
+    from deepagents_code.tools import fetch_url, get_current_thread_id, web_search
+
+    tools: list[Any] = [fetch_url, get_current_thread_id]
+    if settings.has_tavily:
+        tools.append(web_search)
+
+    mcp_server_info: list[Any] | None = None
+    if not config.no_mcp:
+        from deepagents_code.mcp_tools import resolve_and_load_mcp_tools
+
+        try:
+            mcp_tools, _, mcp_server_info = await resolve_and_load_mcp_tools(
+                explicit_config_path=config.mcp_config_path,
+                no_mcp=config.no_mcp,
+                trust_project_mcp=config.trust_project_mcp,
+                project_context=project_context,
+                stateless=True,
+                session_manager=_get_mcp_session_manager(),
+            )
+        except FileNotFoundError:
+            logger.exception("MCP config file not found: %s", config.mcp_config_path)
+            raise
+        except RuntimeError:
+            logger.exception(
+                "Failed to load MCP tools (config: %s)", config.mcp_config_path
+            )
+            raise
+
+        tools.extend(mcp_tools)
+        if mcp_tools:
+            logger.info("Loaded %d MCP tool(s)", len(mcp_tools))
+
+    return tools, mcp_server_info
+
+
+async def _make_graph() -> Any:  # noqa: ANN401
+    """Create the agent graph from environment-based configuration.
+
+    Reads `DEEPAGENTS_CODE_SERVER_*` env vars via `ServerConfig.from_env()`
+    (the inverse of `ServerConfig.to_env()` used by the app process), resolves a
+    model, assembles tools, and compiles the agent graph.
+
+    Returns:
+        Compiled LangGraph agent graph.
+    """
+    config = ServerConfig.from_env()
+    project_context = get_server_project_context()
+
+    from deepagents_code.agent import create_cli_agent, load_async_subagents
+    from deepagents_code.config import (
+        configure_langsmith_secret_redaction,
+        create_model,
+        is_memory_auto_save_enabled,
+        settings,
+    )
+
+    if project_context is not None:
+        settings.reload_from_environment(start_path=project_context.user_cwd)
+    configure_langsmith_secret_redaction()
+
+    # Offload to a worker thread: `create_model` does blocking disk IO for some
+    # providers (e.g. the `openai_codex` token store currently acquires a file
+    # lock via `langchain-openai` that calls `os.mkdir`), which `blockbuster`
+    # rejects on the server event loop.
+    result = await asyncio.to_thread(
+        create_model, config.model, extra_kwargs=config.model_params
+    )
+    result.apply_to_settings()
+
+    tools, mcp_server_info = await _build_tools(config, project_context)
+
+    # Create sandbox backend if a sandbox provider is configured.
+    # The context manager is created here in the factory, but its reference is
+    # stored in a module-level global (and cleaned up via atexit) so the sandbox
+    # lives for the entire server process lifetime. `make_graph` caches the built
+    # graph, so this runs once per process despite LangGraph's per-run factory
+    # invocation.
+    global _sandbox_cm, _sandbox_backend  # noqa: PLW0603
+    sandbox_backend = None
+    if config.sandbox_type:
+        from deepagents_code.integrations.sandbox_factory import create_sandbox
+
+        try:
+            _sandbox_cm = create_sandbox(
+                config.sandbox_type,
+                sandbox_id=config.sandbox_id,
+                snapshot_name=config.sandbox_snapshot_name,
+                setup_script_path=config.sandbox_setup,
+            )
+            _sandbox_backend = _sandbox_cm.__enter__()  # noqa: PLC2801  # Context manager kept open for server process lifetime
+            sandbox_backend = _sandbox_backend
+
+            def _cleanup_sandbox() -> None:
+                if _sandbox_cm is not None:
+                    _sandbox_cm.__exit__(None, None, None)
+
+            atexit.register(_cleanup_sandbox)
+        except ImportError:
+            logger.exception(
+                "Sandbox provider '%s' is not installed", config.sandbox_type
+            )
+            _print_startup_error(
+                f"Sandbox provider '{config.sandbox_type}' is not installed"
+            )
+            sys.exit(1)
+        except NotImplementedError:
+            logger.exception("Sandbox type '%s' is not supported", config.sandbox_type)
+            _print_startup_error(
+                f"Sandbox type '{config.sandbox_type}' is not supported"
+            )
+            sys.exit(1)
+        except ValueError as exc:
+            logger.exception(
+                "Invalid sandbox configuration for '%s'", config.sandbox_type
+            )
+            _print_startup_error(f"Invalid sandbox configuration: {exc}")
+            sys.exit(1)
+        except Exception as exc:
+            logger.exception("Sandbox creation failed for '%s'", config.sandbox_type)
+            _print_startup_error(
+                f"Sandbox creation failed for '{config.sandbox_type}': {exc}"
+            )
+            sys.exit(1)
+
+    def _create_cli_agent_sync() -> Any:  # noqa: ANN401
+        async_subagents = load_async_subagents() or None
+
+        # These process-global settings writes are safe here because `make_graph`
+        # is lock-serialized and caches one graph for the server process lifetime.
+        if config.interpreter_ptc is not None:
+            settings.interpreter_ptc = config.interpreter_ptc
+        if config.interpreter_ptc_acknowledge_unsafe:
+            settings.interpreter_ptc_acknowledge_unsafe = True
+        if config.enable_interpreter:
+            settings.enable_interpreter = True
+
+        agent, _ = create_cli_agent(
+            model=result.model,
+            assistant_id=config.assistant_id,
+            tools=tools,
+            sandbox=sandbox_backend,
+            sandbox_type=config.sandbox_type,
+            system_prompt=config.system_prompt,
+            interactive=config.interactive,
+            auto_approve=config.auto_approve,
+            interrupt_shell_only=config.interrupt_shell_only,
+            shell_allow_list=config.shell_allow_list,
+            enable_ask_user=config.enable_ask_user,
+            enable_memory=config.enable_memory,
+            memory_auto_save=is_memory_auto_save_enabled(),
+            enable_skills=config.enable_skills,
+            enable_shell=config.enable_shell,
+            enable_interpreter=config.enable_interpreter,
+            rubric_model=config.rubric_model,
+            rubric_max_iterations=config.rubric_max_iterations,
+            mcp_server_info=mcp_server_info,
+            cwd=project_context.user_cwd if project_context is not None else config.cwd,
+            project_context=project_context,
+            async_subagents=async_subagents,
+        )
+        return agent
+
+    return await asyncio.to_thread(_create_cli_agent_sync)
+
+
+def _build_graph_factory() -> Callable[[], Awaitable[Any]]:
+    """Build the cached async graph factory exposed to `langgraph dev`.
+
+    The returned coroutine function is what `langgraph.json` references. It keeps
+    its cache and lock in this closure rather than in module-level globals, so
+    importing the module (e.g. for import-only checks) introduces no shared
+    mutable state.
+
+    Returns:
+        A zero-arg async factory that builds the graph once and returns the
+        cached instance on every subsequent call.
+    """
+    missing = object()
+    graph: Any = missing
+    lock = asyncio.Lock()
+
+    async def make_graph() -> Any:  # noqa: ANN401
+        """Create (or return the cached) agent graph for `langgraph dev`.
+
+        LangGraph loads this async factory from the generated `langgraph.json`
+        and invokes it lazily on its event loop — and again on every run. The
+        built graph is cached for the process lifetime so MCP discovery, sandbox
+        creation, and `atexit` registration each happen exactly once; re-running
+        them per request would re-discover MCP servers, leak sandbox sessions,
+        and stack duplicate `atexit` handlers. Any construction failure is
+        converted into a startup-error marker (scraped by the parent app
+        process) before exiting.
+
+        Returns:
+            Compiled LangGraph agent graph.
+        """
+        nonlocal graph
+        if graph is not missing:
+            return graph
+        async with lock:
+            if graph is missing:
+                try:
+                    graph = await _make_graph()
+                except Exception as exc:  # noqa: BLE001  # top-level barrier: any construction failure must surface to the parent as a marker
+                    emit_startup_failure(exc)
+                    sys.exit(1)
+            return graph
+
+    return make_graph
+
+
+make_graph = _build_graph_factory()