nextloop 0.1.0rc1__py3-none-any.whl

nextloop/execute/batch.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import dataclass
+
+from ..contracts import (
+    Cancelled,
+    Completed,
+    EventType,
+    NewEvent,
+    Skipped,
+    ToolCall,
+    ToolOutcome,
+    ToolRuntime,
+    ToolSpec,
+    UncertainHanging,
+    UncertainSettled,
+    ValidatedCall,
+    ValidationIssue,
+)
+
+
+@dataclass(frozen=True)
+class BatchResult:
+    outcomes: tuple[tuple[ToolCall, ToolSpec, ToolOutcome], ...] = ()
+    issues: tuple[ValidationIssue, ...] = ()
+    blocked_resources: frozenset[str] = frozenset()
+
+
+class BatchExecutor:
+    def __init__(self, registry: ToolRuntime, max_parallel: int) -> None:
+        self.registry = registry
+        self.max_parallel = max_parallel
+
+    async def execute(
+        self,
+        calls: Sequence[ToolCall],
+        *,
+        session_id: str,
+        turn: int,
+        step: int,
+        cancel: asyncio.Event,
+        active_resources: frozenset[str],
+        commit: Callable[[Sequence[NewEvent]], Awaitable[None]],
+    ) -> BatchResult:
+        validated, issues = self.registry.validate(calls)
+        if issues:
+            return BatchResult(issues=issues)
+        blocked = frozenset(
+            resource
+            for item in validated
+            for resource in item.spec.resources
+            if resource in active_resources
+        )
+        if blocked:
+            return BatchResult(
+                outcomes=tuple(
+                    (
+                        item.call,
+                        item.spec,
+                        Skipped("Blocked by unresolved resources: " + ", ".join(sorted(blocked))),
+                    )
+                    for item in validated
+                ),
+                blocked_resources=blocked,
+            )
+
+        outcomes: list[tuple[ToolCall, ToolSpec, ToolOutcome]] = []
+        index = 0
+        truncated = False
+        while index < len(validated):
+            if cancel.is_set() or truncated:
+                outcomes.extend(
+                    self._skip(validated[index:], "Cancelled or truncated before dispatch")
+                )
+                break
+            current = validated[index]
+            if current.spec.mode.value == "exclusive":
+                outcome = await self._run_one(current, session_id, turn, step, cancel, commit)
+                outcomes.append((current.call, current.spec, outcome))
+                index += 1
+                if self._closes_batch(current.spec, outcome):
+                    outcomes.extend(
+                        self._skip(validated[index:], "Truncated after exclusive Tool outcome")
+                    )
+                    break
+                continue
+
+            group = []
+            while index < len(validated) and validated[index].spec.mode.value == "parallel":
+                group.append(validated[index])
+                index += 1
+            for start in range(0, len(group), self.max_parallel):
+                chunk = group[start : start + self.max_parallel]
+                if cancel.is_set() or truncated:
+                    outcomes.extend(
+                        self._skip(chunk, "Cancelled or truncated before parallel dispatch")
+                    )
+                    continue
+                chunk_outcomes = await asyncio.gather(
+                    *(self._run_one(item, session_id, turn, step, cancel, commit) for item in chunk)
+                )
+                outcomes.extend(
+                    (item.call, item.spec, outcome)
+                    for item, outcome in zip(chunk, chunk_outcomes, strict=True)
+                )
+                if any(self._is_uncertain_or_cancelled(outcome) for outcome in chunk_outcomes):
+                    truncated = True
+                    outcomes.extend(
+                        self._skip(
+                            group[start + len(chunk) :],
+                            "Truncated after uncertain parallel outcome",
+                        )
+                    )
+        return BatchResult(outcomes=tuple(outcomes))
+
+    async def _run_one(
+        self,
+        item: ValidatedCall,
+        session_id: str,
+        turn: int,
+        step: int,
+        cancel: asyncio.Event,
+        commit: Callable[[Sequence[NewEvent]], Awaitable[None]],
+    ) -> ToolOutcome:
+        await commit(
+            [
+                NewEvent(
+                    event_type=EventType.TOOL_CALL,
+                    turn=turn,
+                    step=step,
+                    payload={
+                        "call_id": item.call.call_id,
+                        "name": item.call.name,
+                        "arguments": item.call.arguments,
+                        "resources": list(item.spec.resources),
+                    },
+                )
+            ]
+        )
+        return await self.registry.execute(
+            item.call,
+            item.arguments,
+            session_id=session_id,
+            turn=turn,
+            step=step,
+            cancel=cancel,
+        )
+
+    @staticmethod
+    def _skip(
+        items: Sequence[ValidatedCall], reason: str
+    ) -> list[tuple[ToolCall, ToolSpec, ToolOutcome]]:
+        return [(item.call, item.spec, Skipped(reason)) for item in items]
+
+    @staticmethod
+    def _is_uncertain_or_cancelled(outcome: ToolOutcome) -> bool:
+        return isinstance(outcome, (Cancelled, UncertainSettled, UncertainHanging))
+
+    @classmethod
+    def _closes_batch(cls, spec: ToolSpec, outcome: ToolOutcome) -> bool:
+        if cls._is_uncertain_or_cancelled(outcome):
+            return True
+        return isinstance(outcome, Completed) and (outcome.is_error or outcome.concludes_turn)

nextloop/execute/model.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+
+from ..contracts import (
+    CandidateRequest,
+    EventType,
+    LoopConfig,
+    ModelProvider,
+    ModelResponse,
+    NewEvent,
+    ProviderError,
+)
+
+
+@dataclass(frozen=True)
+class ModelCallResult:
+    events: tuple[NewEvent, ...]
+    response: ModelResponse | None = None
+    error: ProviderError | None = None
+    cancelled: bool = False
+    rebuild_for_steer: bool = False
+
+
+class ModelCaller:
+    def __init__(
+        self,
+        provider: ModelProvider,
+        config: LoopConfig,
+        has_steer: Callable[[], Awaitable[bool]],
+    ) -> None:
+        self.provider = provider
+        self.config = config
+        self.has_steer = has_steer
+
+    async def call(self, request: CandidateRequest, cancel: asyncio.Event) -> ModelCallResult:
+        events: list[NewEvent] = []
+        for attempt in range(1, self.config.max_provider_attempts + 1):
+            if cancel.is_set():
+                return ModelCallResult(events=tuple(events), cancelled=True)
+            events.append(self._event(request, attempt, "started"))
+            try:
+                response = await _await_with_cancel(
+                    self.provider.complete(request, attempt=attempt, cancel=cancel),
+                    cancel,
+                )
+            except _CancelledByUser:
+                events.append(self._event(request, attempt, "cancelled"))
+                return ModelCallResult(events=tuple(events), cancelled=True)
+            except ProviderError as exc:
+                events.append(
+                    self._event(
+                        request,
+                        attempt,
+                        "error",
+                        error=str(exc),
+                        retryable=exc.retryable,
+                        partial_output=exc.partial_output,
+                        code=exc.code,
+                    )
+                )
+                if cancel.is_set():
+                    return ModelCallResult(events=tuple(events), cancelled=True)
+                if (
+                    exc.partial_output
+                    or not exc.retryable
+                    or attempt >= self.config.max_provider_attempts
+                ):
+                    return ModelCallResult(events=tuple(events), error=exc)
+                if await self.has_steer():
+                    return ModelCallResult(events=tuple(events), rebuild_for_steer=True)
+                delay = self.config.retry_backoff_seconds * (2 ** (attempt - 1))
+                if delay and await _sleep_or_cancel(delay, cancel):
+                    return ModelCallResult(events=tuple(events), cancelled=True)
+                continue
+            except asyncio.CancelledError:
+                events.append(self._event(request, attempt, "cancelled"))
+                return ModelCallResult(events=tuple(events), cancelled=True)
+            except Exception as exc:
+                error = ProviderError(f"{type(exc).__name__}: {exc}")
+                events.append(self._event(request, attempt, "error", error=str(error)))
+                return ModelCallResult(events=tuple(events), error=error)
+            events.append(self._event(request, attempt, "success"))
+            return ModelCallResult(events=tuple(events), response=response)
+        return ModelCallResult(
+            events=tuple(events),
+            error=ProviderError("Provider retry loop exited without a response"),
+        )
+
+    @staticmethod
+    def _event(request: CandidateRequest, attempt: int, status: str, **payload: object) -> NewEvent:
+        return NewEvent(
+            event_type=EventType.ASSISTANT_ATTEMPT,
+            turn=request.turn,
+            step=request.step,
+            attempt=attempt,
+            payload={"status": status, **payload},
+        )
+
+
+class _CancelledByUser(Exception):
+    pass
+
+
+async def _await_with_cancel(
+    awaitable: Awaitable[ModelResponse], cancel: asyncio.Event
+) -> ModelResponse:
+    if cancel.is_set():
+        raise _CancelledByUser()
+    operation = asyncio.ensure_future(awaitable)
+    cancellation = asyncio.create_task(cancel.wait())
+    try:
+        done, _ = await asyncio.wait({operation, cancellation}, return_when=asyncio.FIRST_COMPLETED)
+        if cancellation in done and not operation.done():
+            operation.cancel()
+            await asyncio.gather(operation, return_exceptions=True)
+            raise _CancelledByUser()
+        cancellation.cancel()
+        return await operation
+    finally:
+        cancellation.cancel()
+        if not operation.done():
+            operation.cancel()
+        await asyncio.gather(operation, cancellation, return_exceptions=True)
+
+
+async def _sleep_or_cancel(delay: float, cancel: asyncio.Event) -> bool:
+    try:
+        await asyncio.wait_for(cancel.wait(), timeout=delay)
+        return True
+    except TimeoutError:
+        return cancel.is_set()

nextloop/execute/tools.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import jsonschema
+
+from ..contracts import (
+    Cancelled,
+    Completed,
+    ToolCall,
+    ToolOutcome,
+    ToolSpec,
+    UncertainHanging,
+    ValidatedCall,
+    ValidationIssue,
+)
+
+
+@dataclass(frozen=True)
+class ToolContext:
+    session_id: str
+    turn: int
+    step: int
+    cancel: asyncio.Event
+
+
+ToolHandler = Callable[[Mapping[str, Any], ToolContext], Any]
+
+
+class ToolRegistry:
+    def __init__(self) -> None:
+        self._specs: dict[str, ToolSpec] = {}
+        self._handlers: dict[str, ToolHandler] = {}
+
+    def register(self, spec: ToolSpec, handler: ToolHandler) -> None:
+        if spec.name in self._specs:
+            raise ValueError(f"Duplicate Tool name: {spec.name}")
+        self._specs[spec.name] = spec
+        self._handlers[spec.name] = handler
+
+    def catalog(self) -> Sequence[ToolSpec]:
+        return tuple(self._specs.values())
+
+    def validate(
+        self, calls: Sequence[ToolCall]
+    ) -> tuple[tuple[ValidatedCall, ...], tuple[ValidationIssue, ...]]:
+        validated = []
+        issues = []
+        seen_ids: set[str] = set()
+        for index, call in enumerate(calls):
+            call_id = call.call_id or f"unidentified_call_{index}"
+            if not call.call_id:
+                issues.append(ValidationIssue(call_id, call.name, "Tool Call identity is empty"))
+            elif call.call_id in seen_ids:
+                issues.append(
+                    ValidationIssue(call_id, call.name, "Tool Call identity is duplicated")
+                )
+            seen_ids.add(call.call_id)
+            spec = self._specs.get(call.name)
+            if spec is None:
+                issues.append(ValidationIssue(call_id, call.name, "Unknown Tool name"))
+                continue
+            try:
+                arguments = json.loads(call.arguments)
+            except (json.JSONDecodeError, TypeError) as exc:
+                issues.append(ValidationIssue(call_id, call.name, f"Invalid argument JSON: {exc}"))
+                continue
+            if not isinstance(arguments, dict):
+                issues.append(
+                    ValidationIssue(call_id, call.name, "Tool arguments must be a JSON object")
+                )
+                continue
+            validator = jsonschema.Draft202012Validator(spec.parameters)
+            for error in sorted(validator.iter_errors(arguments), key=lambda item: list(item.path)):
+                path = "$" + "".join(
+                    f"[{part}]" if isinstance(part, int) else f".{part}" for part in error.path
+                )
+                issues.append(ValidationIssue(call_id, call.name, f"{path}: {error.message}"))
+            validated.append(ValidatedCall(call=call, spec=spec, arguments=arguments))
+        return tuple(validated), tuple(issues)
+
+    async def execute(
+        self,
+        call: ToolCall,
+        arguments: Mapping[str, Any],
+        *,
+        session_id: str,
+        turn: int,
+        step: int,
+        cancel: asyncio.Event,
+    ) -> ToolOutcome:
+        spec = self._specs[call.name]
+        if cancel.is_set():
+            return Cancelled(
+                settled=True, partial_content="Cancelled before Tool handler execution"
+            )
+        handler = self._handlers[call.name]
+        try:
+            operation = handler(arguments, ToolContext(session_id, turn, step, cancel))
+            if inspect.isawaitable(operation):
+                cancellable = _await_with_cancel(operation, cancel)
+                if spec.timeout_seconds is not None:
+                    cancellable = asyncio.wait_for(cancellable, timeout=spec.timeout_seconds)
+                value = await cancellable
+            else:
+                value = operation
+            if isinstance(value, (Completed, Cancelled, UncertainHanging)):
+                return value
+            from ..contracts import Skipped, UncertainSettled
+
+            if isinstance(value, (Skipped, UncertainSettled)):
+                return value
+            if isinstance(value, str):
+                return Completed(value)
+            return Completed(str(value), structured=value)
+        except TimeoutError:
+            return UncertainHanging(
+                f"Tool {call.name} timed out; external effects are unknown",
+                resource_locks=spec.resources,
+            )
+        except asyncio.CancelledError:
+            return Cancelled(
+                settled=False,
+                partial_content="Tool handler was cancelled; external effects are unknown",
+                resource_locks=spec.resources,
+            )
+        except Exception as exc:
+            return Completed(f"{type(exc).__name__}: {exc}", is_error=True)
+
+
+async def _await_with_cancel(awaitable: Awaitable[Any], cancel: asyncio.Event) -> Any:
+    operation = asyncio.ensure_future(awaitable)
+    cancellation = asyncio.create_task(cancel.wait())
+    try:
+        done, _ = await asyncio.wait({operation, cancellation}, return_when=asyncio.FIRST_COMPLETED)
+        if cancellation in done and not operation.done():
+            operation.cancel()
+            await asyncio.gather(operation, return_exceptions=True)
+            raise asyncio.CancelledError()
+        cancellation.cancel()
+        return await operation
+    finally:
+        cancellation.cancel()
+        if not operation.done():
+            operation.cancel()
+        await asyncio.gather(operation, cancellation, return_exceptions=True)

nextloop/loop.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from .contracts import (
+    DeliveryOwner,
+    EventStore,
+    EventType,
+    NewEvent,
+    RecoveryOutcome,
+    TurnContext,
+    TurnOutcome,
+)
+from .driver import SessionActor, SessionExecution, TurnDriver
+from .events.repair import pending_final_messages, repair_events
+
+
+class AgentLoop:
+    def __init__(self) -> None:
+        self._actors: dict[str, SessionActor] = {}
+        self._active: dict[str, SessionExecution] = {}
+        self._resource_locks: dict[str, set[str]] = {}
+
+    async def run(self, context: TurnContext) -> TurnOutcome:
+        actor = self._actors.setdefault(context.session_id, SessionActor())
+        async with actor.run_lock:
+            writer_epoch, last_seq, last_turn = await context.store.attach_writer(
+                context.session_id
+            )
+            execution = SessionExecution(
+                context=context,
+                actor=actor,
+                writer_epoch=writer_epoch,
+                seq=last_seq,
+                turn=last_turn + 1,
+                resource_locks=set(self._resource_locks.get(context.session_id, set())),
+            )
+            self._active[context.session_id] = execution
+            try:
+                await execution.append(
+                    [
+                        NewEvent(
+                            event_type=EventType.INPUT_ACCEPTED,
+                            turn=execution.turn,
+                            step=0,
+                            payload={
+                                "input_id": context.input_id,
+                                "content": context.initial_input,
+                            },
+                        ),
+                        NewEvent(event_type=EventType.TURN_START, turn=execution.turn, step=0),
+                    ]
+                )
+                outcome = await TurnDriver(execution).run()
+                self._resource_locks[context.session_id] = set(execution.resource_locks)
+                return outcome
+            finally:
+                self._active.pop(context.session_id, None)
+
+    async def steer(self, session_id: str, content: str, *, input_id: str | None = None) -> bool:
+        execution = self._active.get(session_id)
+        if execution is None or not content.strip():
+            return False
+        async with execution.actor.commit_lock:
+            if execution.final_committed:
+                return False
+            accepted_id = (
+                input_id or f"in_{len(execution.actor.inbox) + 1}_{execution.turn}_{execution.step}"
+            )
+            await execution._append_locked(
+                [
+                    NewEvent(
+                        event_type=EventType.INPUT_ACCEPTED,
+                        turn=execution.turn,
+                        step=execution.step,
+                        payload={"input_id": accepted_id, "content": content},
+                    )
+                ]
+            )
+            execution.actor.inbox.append((accepted_id, content))
+            return True
+
+    async def cancel(self, session_id: str, cause: str = "user_stop") -> bool:
+        execution = self._active.get(session_id)
+        if execution is None:
+            return False
+        async with execution.actor.commit_lock:
+            if execution.final_committed:
+                return False
+            await execution._append_locked(
+                [
+                    NewEvent(
+                        event_type=EventType.CANCELLATION_REQUESTED,
+                        turn=execution.turn,
+                        step=execution.step,
+                        payload={"cause": cause},
+                    )
+                ]
+            )
+            execution.cancel.set()
+            return True
+
+    async def recover(
+        self,
+        session_id: str,
+        store: EventStore,
+        delivery: DeliveryOwner,
+    ) -> RecoveryOutcome:
+        actor = self._actors.setdefault(session_id, SessionActor())
+        async with actor.run_lock:
+            if session_id in self._active:
+                raise RuntimeError("Cannot recover a Session with an active Turn")
+            writer_epoch, last_seq, _ = await store.attach_writer(session_id)
+            events = list(await store.load_recovery(session_id))
+            repaired = repair_events(events)
+            seq = last_seq
+
+            async def append(new_events: Sequence[NewEvent]) -> None:
+                nonlocal seq
+                if not new_events:
+                    return
+                committed = await store.append(session_id, writer_epoch, seq + 1, new_events)
+                seq = committed[-1].seq
+
+            await append(repaired.events)
+            self._resource_locks[session_id] = set(repaired.active_resource_locks)
+            delivered = []
+            delivery_unknown = []
+            for message in pending_final_messages(events):
+                try:
+                    await delivery.deliver(message)
+                    await append(
+                        [
+                            NewEvent(
+                                event_type=EventType.DELIVERY_COMMITTED,
+                                turn=message.turn,
+                                step=message.step,
+                                payload={
+                                    "final_message_id": message.final_message_id,
+                                    "recovered": True,
+                                },
+                            )
+                        ]
+                    )
+                    delivered.append(message.final_message_id)
+                except Exception as exc:
+                    await append(
+                        [
+                            NewEvent(
+                                event_type=EventType.DELIVERY_UNKNOWN,
+                                turn=message.turn,
+                                step=message.step,
+                                payload={
+                                    "final_message_id": message.final_message_id,
+                                    "error": str(exc),
+                                    "recovered": True,
+                                },
+                            )
+                        ]
+                    )
+                    delivery_unknown.append(message.final_message_id)
+            return RecoveryOutcome(
+                session_id=session_id,
+                writer_epoch=writer_epoch,
+                interrupted_turns=repaired.interrupted_turns,
+                delivered_final_ids=tuple(delivered),
+                delivery_unknown_final_ids=tuple(delivery_unknown),
+                active_resource_locks=repaired.active_resource_locks,
+            )