zettcode-engine 0.3.4__py3-none-any.whl

zettcode/engine/extensions/harness.py

@@ -0,0 +1,208 @@
+"""The just-in-time harness, kept behind a switch the reader can flip between two requests.
+
+The framework's :class:`~zett_agent.jit_harness.JITHarnessExtension` lets the model write small Python tools of its
+own and run them. That is opt-in by nature, and it is not behind the shell's permission level by itself: the plugin
+puts a child process (and, below `full`, a sandbox) between that code and this window. So it is registered always and
+does nothing until switched on, which is what lets the ``/plugins`` switch take effect on the next start without rebuilding
+the client.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Callable
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from zett_agent.agent import AgentRunContext
+from zett_agent.events import AgentEvent, AgentEventType
+from zett_agent.extensions.base import ModelRequestNext, ToolCallNext
+from zett_agent.jit_harness import JITHarnessExtension
+from zett_agent.messages import ToolCall, ToolMessage
+from zett_agent.model import ModelEvent, ModelRequest
+from zett_agent.tools.base import ToolResult
+
+from .harness_switch import HarnessSwitch
+from .mcp_notice import HARNESS_NOTICE_EVENT_NAME, notice_payload
+
+if TYPE_CHECKING:  # pragma: no cover - annotations only
+    from zett_agent.jit_harness import JITToolHost
+
+__all__ = ["SwitchedHarness"]
+
+
+class SwitchedHarness(JITHarnessExtension):
+    """The framework's harness, inert while its switch is off.
+
+    Every hook that changes a request first asks the switch. A hook that only releases what a request held is left
+    alone, so a request that began with the harness on still lets go of it cleanly. The switch is read once per
+    request, in ``on_tool``, because ``/harness`` waits for an idle shell: no request sees it change halfway.
+    """
+
+    def __init__(
+        self,
+        root: str | Path,
+        switch: HarnessSwitch,
+        *,
+        host_factory: Callable[[], JITToolHost] | None = None,
+        boundary: Callable[[], str] | None = None,
+        unbounded: Callable[[], bool] | None = None,
+        legacy_root: str | Path | None = None,
+    ) -> None:
+        """Keep each session's repository under ``root``, and follow ``switch``.
+
+        Args:
+            root: Directory that holds one subdirectory per session.
+            switch: Whether the harness is on for a request.
+            host_factory: Where a session's repository code runs, called once per session; ``None`` keeps the
+                framework's default, which imports it here. The plugin passes one that puts a child process between
+                the model's code and this window, inside a sandbox unless the reader gave the model everything.
+            boundary: What the host was built for — the shell level, in the plugin's words — read once per request.
+                A session's host is made once, so a level that changed since it was built (``/permission full``
+                while code was already running under a sandbox) would otherwise leave the old answer in force.
+                When it differs, the session is forgotten, which closes its child and builds the next host for the
+                level now in force.
+            unbounded: Whether the level in force promises a boundary this machine does not have. The code the model
+                wrote then runs with the reader's own access, which they are told once per session.
+            legacy_root: The root from before the workspace level, whose session directory is moved under ``root``
+                the first time that session asks for it.
+        """
+        super().__init__(root, host_factory=host_factory)
+        self._root = Path(root)
+        self._legacy_root = Path(legacy_root) if legacy_root is not None else None
+        self._switch = switch
+        self._boundary = boundary
+        self._unbounded = unbounded
+        #: Sessions already told that their repository code has no boundary on this machine.
+        self._said_unbounded: set[str] = set()
+        #: The boundary each session's host was built for, so a change of level is noticed.
+        self._built: dict[str, str] = {}
+        #: Sessions whose child must be closed once the request in flight lets go of them.
+        self._pending: set[str] = set()
+        #: Sessions whose current request has the harness, set by ``on_tool`` and cleared when it is released.
+        self._engaged: set[str] = set()
+
+    def engaged(self, context: AgentRunContext) -> bool:
+        """Return whether this request was given the harness."""
+        return context.session_id in self._engaged
+
+    async def on_tool(self, context: AgentRunContext) -> None:
+        """Mount the harness's tools, when it is on — in a host built for the level in force now."""
+        if self._legacy_root is not None:
+            from zettcode.engine.harness_paths import move_session
+
+            move_session(self._legacy_root, self._root, context.session_id)
+        if not self._switch.on:
+            return
+        self._follow_boundary(context.session_id)
+        self._engaged.add(context.session_id)
+        await self._say_unbounded(context)
+        await super().on_tool(context)
+
+    async def _say_unbounded(self, context: AgentRunContext) -> None:
+        """Tell the reader, once per session, that this machine cannot bound the code the model writes.
+
+        A machine with no sandbox passes ``None`` to the host, which runs the child with the reader's own access. The
+        harness says so because the model's own tools are not a shell command and none of the shell's fallbacks apply
+        to them. A line of news must never be what stops a request, so a failure to publish it is dropped.
+        """
+        if self._unbounded is None or not self._unbounded() or context.session_id in self._said_unbounded:
+            return
+        try:
+            await context.emit(
+                AgentEvent(
+                    type=AgentEventType.CUSTOM,
+                    session_id=context.session_id,
+                    name=HARNESS_NOTICE_EVENT_NAME,
+                    payload=notice_payload(
+                        "this machine has no sandbox, so the tools the model writes run with your own access; "
+                        "`/plugins` turns the harness plugin off",
+                        "warn",
+                    ),
+                )
+            )
+        except Exception:  # noqa: BLE001 - news is never what stops a request
+            return
+        self._said_unbounded.add(context.session_id)
+
+    def _follow_boundary(self, session_id: str) -> None:
+        """Forget a session whose host was built for a boundary that is no longer in force."""
+        if self._boundary is None:
+            return
+        now = self._boundary()
+        built = self._built.get(session_id)
+        if built is None:
+            self._built[session_id] = now
+            return
+        if built == now:
+            return
+        # Its child was started without (or with) a boundary the reader has since changed: close it, and let the next
+        # request build the host the boundary asks for. A fail-open sandbox is the case that matters. The new value is
+        # recorded only once the old child is really gone, so a forget that had to wait is tried again next time.
+        if self._forget(session_id):
+            self._built[session_id] = now
+        else:
+            self._pending.add(session_id)
+
+    def _forget(self, session_id: str) -> bool:
+        """Close a session's child now, or say that the request in flight still holds it."""
+        try:
+            self.forget_session(session_id)
+        except KeyError:
+            # The framework dropped it already (an idle session past its cap): nothing is left to close, and nothing
+            # that is keyed by it is worth keeping either.
+            self._drop(session_id)
+            return True
+        except RuntimeError:
+            return False
+        self._drop(session_id)
+        return True
+
+    def _drop(self, session_id: str) -> None:
+        """Forget what this window remembered about a session the framework no longer holds."""
+        self._built.pop(session_id, None)
+        self._said_unbounded.discard(session_id)
+
+    async def on_state(self, context: AgentRunContext) -> None:
+        """Adopt late tool registrations, for a request that has the harness."""
+        if self.engaged(context):
+            await super().on_state(context)
+
+    async def on_tool_call(self, context: AgentRunContext, call: ToolCall, call_next: ToolCallNext) -> ToolResult:
+        """Run a call through the harness when it is engaged, and straight through when not."""
+        if not self.engaged(context):
+            return await call_next()
+        return await super().on_tool_call(context, call, call_next)
+
+    async def before_tool(self, context: AgentRunContext, call: ToolCall) -> None:
+        """Announce a harness tool call, when engaged."""
+        if self.engaged(context):
+            await super().before_tool(context, call)
+
+    async def after_tool(
+        self, context: AgentRunContext, call: ToolCall, result: ToolMessage, error: Exception | None
+    ) -> None:
+        """Report a harness tool call's outcome, when engaged."""
+        if self.engaged(context):
+            await super().after_tool(context, call, result, error)
+
+    async def on_model_request(
+        self, context: AgentRunContext, request: ModelRequest, call_next: ModelRequestNext
+    ) -> AsyncIterator[ModelEvent]:
+        """Add the harness's guidance to the request, when engaged; leave it untouched when not."""
+        if not self.engaged(context):
+            async for event in call_next(request):
+                yield event
+            return
+        async for event in super().on_model_request(context, request, call_next):
+            yield event
+
+    def _release(self, context: AgentRunContext) -> None:
+        """Let go of what the request held, and forget that it had the harness."""
+        self._engaged.discard(context.session_id)
+        # Detach first: the framework refuses to forget a session an active request still holds, and the release is
+        # what makes it idle. Only then can the child built for the old boundary be closed.
+        super()._release(context)
+        if context.session_id in self._pending:
+            self._pending.discard(context.session_id)
+            if self._forget(context.session_id) and self._boundary is not None:
+                self._built[context.session_id] = self._boundary()

zettcode/engine/extensions/harness_switch.py

@@ -0,0 +1,18 @@
+"""Whether the just-in-time harness is on. Nothing here imports the framework, so a launch can hold it for free."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+__all__ = ["HarnessSwitch"]
+
+
+@dataclass
+class HarnessSwitch:
+    """Whether the harness is on; the runtime owns it, the command flips it, the extension reads it.
+
+    Attributes:
+        on: ``True`` while the model is given the harness.
+    """
+
+    on: bool = False

zettcode/engine/extensions/mcp_notice.py

@@ -0,0 +1,49 @@
+"""What an extension tells the reader as a custom event the transcript can show: MCP news, and the harness's.
+
+This is a module of its own, apart from :mod:`.mcp`, because the transcript has to recognise the
+event and :mod:`.mcp` imports the MCP SDK, which costs more than the first frame. The name and the
+payload live here, where importing them costs nothing.
+"""
+
+from __future__ import annotations
+
+from zett_agent.events import AgentEvent
+
+#: Name of the custom event that carries one line of MCP news for the reader.
+MCP_NOTICE_EVENT_NAME = "mcp_notice"
+
+#: The same, for a line about the just-in-time harness: one channel, so a reader sees both in one place.
+HARNESS_NOTICE_EVENT_NAME = "harness_notice"
+
+#: Every event that carries a line for the reader.
+NOTICE_EVENT_NAMES = frozenset({MCP_NOTICE_EVENT_NAME, HARNESS_NOTICE_EVENT_NAME})
+
+
+#: How serious a line of news is: ``info`` is progress (a server is starting, or ready), ``warn`` is something the
+#: reader may need to act on (a server is down, an entry is ignored). A payload that names neither is ``info``.
+LEVELS = ("info", "warn")
+
+
+def notice_payload(text: str, level: str = "info") -> dict[str, str]:
+    """Return the JSON-serializable payload that carries one line of MCP news and how serious it is."""
+    return {"text": text, "level": level if level in LEVELS else "info"}
+
+
+def notice_level_from(event: AgentEvent) -> str:
+    """Return how serious the notice an event carries is; ``info`` for a payload that does not say."""
+    payload = event.payload
+    level = payload.get("level") if isinstance(payload, dict) else None
+    return level if level in LEVELS else "info"
+
+
+def notice_from(event: AgentEvent) -> str | None:
+    """Return the line an MCP notice event carries, or ``None`` when it carries nothing readable.
+
+    Junk is ignored rather than raised on: an event reaches here from a hook, and a malformed one
+    must not be able to fail the request that produced it.
+    """
+    payload = event.payload
+    if event.name not in NOTICE_EVENT_NAMES or not isinstance(payload, dict):
+        return None
+    text = payload.get("text")
+    return text.strip() if isinstance(text, str) and text.strip() else None

zettcode/engine/extensions/relay.py

@@ -0,0 +1,93 @@
+"""What a child agent has to say to the person: the shell approvals it cannot decide for itself.
+
+A subagent runs in its own session, and nothing watches its event stream, so a command that needs approval would be
+put to nobody and wait for ever. That one kind of event is handed to the parent run instead, which is the one the
+person is watching. It goes out under the parent's session id, because the parent's stream refuses an event that
+belongs to another session, but the payload still names the child's session and the tool call: the shell and the SDK
+answer with those, the answer is routed to the extension that is waiting for it (the same instance serves parent and
+child), and the child's tool call carries on. Everything else a child does stays in its own log.
+
+A question (``ask_user``) is deliberately not relayed: a child does not have that tool, because a child works until it
+has finished and a question put to the person in the middle of it would be a conversation the person never started.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping
+from contextlib import aclosing
+from typing import cast
+
+from zett_agent.agent import Agent, AgentRunConfig, AgentRunContext
+from zett_agent.events import AgentEvent, AgentEventType
+from zett_agent.extensions.shell_approval import SHELL_APPROVAL_EVENT_NAME
+from zett_agent.json_types import JsonValue
+from zett_agent.messages import UserMessage
+from zett_agent.model import AgentModel, ReasoningEffort
+
+#: The one extension event a person has to see to let a child go on.
+RELAYED = frozenset({SHELL_APPROVAL_EVENT_NAME})
+
+#: Where one such event is sent.
+Relay = Callable[[AgentEvent], Awaitable[None]]
+
+
+def relay_to(parent: AgentRunContext) -> Relay:
+    """Return the function that hands a child's event to ``parent``'s run, under the parent's session id."""
+
+    async def relay(event: AgentEvent) -> None:
+        await parent.emit(
+            AgentEvent(
+                AgentEventType.CUSTOM,
+                parent.session_id,
+                tool_calls=list(event.tool_calls),
+                name=event.name,
+                payload=dict(event.payload or {}),
+            )
+        )
+
+    return relay
+
+
+class RelayingAgent(Agent):
+    """An agent that passes the events a person has to answer to :attr:`relay` as they are emitted.
+
+    Attributes:
+        relay: Where such an event goes; ``None`` drops them, as an agent nobody watches does.
+    """
+
+    relay: Relay | None = None
+
+    async def stream(
+        self,
+        message: UserMessage | str,
+        *,
+        config: AgentRunConfig | None = None,
+        model: AgentModel | None = None,
+        reasoning_effort: ReasoningEffort | None = None,
+        parallel_tool_call: bool | None = None,
+        metadata: Mapping[str, JsonValue] | None = None,
+        tags: Mapping[str, JsonValue] | None = None,
+    ) -> AsyncIterator[AgentEvent]:
+        """Stream one run, relaying approvals and questions on the way; the events themselves are unchanged."""
+        # ``Agent.stream`` is declared as an iterator and is an async generator: closing it is what stops the run
+        # underneath when the caller stops reading.
+        events = cast(
+            AsyncGenerator[AgentEvent, None],
+            super().stream(
+                message,
+                config=config,
+                model=model,
+                reasoning_effort=reasoning_effort,
+                parallel_tool_call=parallel_tool_call,
+                metadata=metadata,
+                tags=tags,
+            ),
+        )
+        async with aclosing(events):
+            async for event in events:
+                if self.relay is not None and event.type is AgentEventType.CUSTOM and event.name in RELAYED:
+                    await self.relay(event)
+                yield event
+
+
+__all__ = ["RELAYED", "Relay", "RelayingAgent", "relay_to"]