Download distinct_server/control_plane.py from User1342/distinct: direct link, hf CLI and curl.
- Browser
- Download file 93.4 kB
-
https://huggingface.co/spaces/User1342/distinct/resolve/main/distinct_server/control_plane.py
- Command line
-
hf download hf://spaces/User1342/distinct/distinct_server/control_plane.py
-
curl -L -o control_plane.py https://huggingface.co/spaces/User1342/distinct/resolve/main/distinct_server/control_plane.py
93.4 kB
| """Thread-safe, bounded, in-memory control plane for the Distinct MVP. | |
| All public methods take the same re-entrant lock. Multi-record operations such | |
| as pairing, offer acceptance, dependency release, and snapshot reconciliation | |
| therefore commit atomically from the perspective of HTTP request threads. | |
| """ | |
| from __future__ import annotations | |
| import copy | |
| import hmac | |
| import json | |
| import math | |
| import secrets | |
| import threading | |
| import time | |
| import uuid | |
| from dataclasses import replace | |
| from typing import Any, Callable, Mapping, Optional, Sequence, Tuple, Union | |
| from distinct_protocol import ( | |
| AgentCapabilities, | |
| AgentSnapshot, | |
| AgentStatus, | |
| ChatMessage, | |
| JobSpec, | |
| JobStatus, | |
| ToolSelection, | |
| ) | |
| from distinct_protocol.handshake import ( | |
| ALG_ED25519, | |
| NOTHING, | |
| Allowance, | |
| ServerCatalogue, | |
| SignedSnapshot, | |
| ) | |
| from distinct_protocol.handshake import sign_response as _sign_response_payload | |
| from .models import ( | |
| AgentCredential, | |
| AgentRecord, | |
| AgentView, | |
| AuthenticationError, | |
| CapacityError, | |
| CleanupReport, | |
| ConflictError, | |
| ControlPlaneLimits, | |
| ConversationRecord, | |
| ConversationView, | |
| DependencyError, | |
| JobRecord, | |
| JobView, | |
| NoCompatibleAgentError, | |
| NotFoundError, | |
| PairingCode, | |
| PairingRecord, | |
| RunSelection, | |
| SessionRecord, | |
| SessionView, | |
| ValidationError, | |
| ) | |
| from .security import ( | |
| access_code_digest, | |
| canonical_request, | |
| derive_user_key, | |
| encode_secret, | |
| generate_secret, | |
| new_pairing_code, | |
| pairing_code_digest, | |
| validate_nonce, | |
| ) | |
| _TERMINAL_STATUSES = frozenset( | |
| { | |
| JobStatus.COMPLETED, | |
| JobStatus.FAILED, | |
| JobStatus.CANCELLED, | |
| JobStatus.EXPIRED, | |
| } | |
| ) | |
| _AGENT_EXECUTION_STATUSES = frozenset( | |
| {JobStatus.ACCEPTED, JobStatus.QUEUED, JobStatus.RUNNING} | |
| ) | |
| #: Statuses from which a job can strand when its agent goes silent, and to | |
| #: which a stranded job can return when the agent reappears. | |
| _STRANDABLE_STATUSES = frozenset( | |
| {JobStatus.OFFERED, JobStatus.ACCEPTED, JobStatus.QUEUED, JobStatus.RUNNING} | |
| ) | |
| _ELIGIBLE_AGENT_STATUSES = frozenset({AgentStatus.ONLINE, AgentStatus.BUSY}) | |
| # ``JobSpec.prompt`` may not be empty, so an erased prompt is replaced with this | |
| # tombstone rather than removed. It is not user text and must never be shown as | |
| # though it were: ``JobView.prompt_erased`` is the flag callers should read. | |
| ERASED_PROMPT = "[erased]" | |
| # One string for both "no such agent" and "a private agent you may not see". | |
| # Two different messages would let anyone enumerate other people's private | |
| # workers by reading the error text, so there is exactly one message and it is | |
| # a constant rather than two literals that could drift apart. | |
| _AGENT_NOT_FOUND = "agent was not found" | |
| #: Placeholder target for a job submitted before any compatible agent exists. | |
| #: Such a job is born STRANDED and picked up by the recovery pass the moment a | |
| #: compatible agent is online, which is what lets a user queue a request and | |
| #: then go start their worker. | |
| UNASSIGNED_AGENT = "unassigned" | |
| # The bounded conversation window carried forward to a released follow-up. | |
| # ``JobSpec`` caps messages at 256; this is the smaller product-level window. | |
| _MAX_CONTEXT_MESSAGES = 40 | |
| _MAX_SYSTEM_MESSAGES = 8 | |
| def _new_id(prefix: str) -> str: | |
| return f"{prefix}_{uuid.uuid4().hex}" | |
| def _bounded_context(messages: Sequence[ChatMessage]) -> tuple[ChatMessage, ...]: | |
| """Trim a context window from the front, keeping leading system messages.""" | |
| system = [message for message in messages if message.role == "system"][:_MAX_SYSTEM_MESSAGES] | |
| rest = [message for message in messages if message.role != "system"] | |
| budget = max(0, _MAX_CONTEXT_MESSAGES - len(system)) | |
| return tuple(system) + tuple(rest[-budget:]) | |
| def _single_line_identifier(value: str, label: str) -> str: | |
| if not isinstance(value, str) or not value or len(value) > 128: | |
| raise ValidationError(f"{label} must be 1-128 characters") | |
| if any(character in value for character in "\r\n"): | |
| raise ValidationError(f"{label} must be single-line text") | |
| return value | |
| def _tool_capability(tool: ToolSelection) -> str: | |
| return f"{tool.id}@{tool.version}" | |
| #: The capabilities a phantom record carries. Built once because it is | |
| #: immutable and never read, and named so nothing mistakes it for a machine. | |
| _PHANTOM_CAPABILITIES = AgentCapabilities( | |
| agent_id="phantom", | |
| name="phantom", | |
| os="phantom", | |
| arch="phantom", | |
| cpu="phantom", | |
| ram_gb=0.0, | |
| models=(), | |
| tools=(), | |
| energy_provider="none", | |
| modes=(), | |
| queue_capacity=1, | |
| ) | |
| #: The one thing an unauthenticated agent request is ever told. Named once so | |
| #: the two paths that use it cannot drift into being distinguishable again. | |
| UNAUTHENTICATED = ( | |
| "this request could not be authenticated. If this worker was paired " | |
| "before the server restarted, pair it again" | |
| ) | |
| class ControlPlane: | |
| """Own all volatile server scheduling and worker-authentication state.""" | |
| def __init__( | |
| self, | |
| limits: Optional[ControlPlaneLimits] = None, | |
| *, | |
| clock: Callable[[], float] = time.time, | |
| pairing_pepper: Optional[bytes] = None, | |
| server_id: str = "distinct-server", | |
| catalogue: Optional[Allowance] = None, | |
| ) -> None: | |
| self.limits = limits or ControlPlaneLimits() | |
| self._clock = clock | |
| self.server_id = _single_line_identifier(server_id, "server_id") | |
| # What this server may ask an agent to do. Injected rather than | |
| # imported from the model catalogue so a test can state its own, and so | |
| # this module does not take ownership of a catalogue another part of the | |
| # project maintains. | |
| self.catalogue = catalogue if isinstance(catalogue, Allowance) else NOTHING | |
| self._pairing_pepper = pairing_pepper or generate_secret() | |
| if len(self._pairing_pepper) < 32: | |
| raise ValueError("pairing pepper must contain at least 32 bytes") | |
| self._lock = threading.RLock() | |
| self._pairing_codes: dict[str, PairingRecord] = {} | |
| self._agents: dict[str, AgentRecord] = {} | |
| self._sessions: dict[str, SessionRecord] = {} | |
| self._conversations: dict[tuple[str, str], ConversationRecord] = {} | |
| self._jobs: dict[str, JobRecord] = {} | |
| self._children: dict[str, set[str]] = {} | |
| # Prompt-free cancellation tombstones let a worker safely learn that a | |
| # job vanished when a browser session is deleted. | |
| self._orphan_cancellations: dict[str, dict[str, float]] = {} | |
| # Pairing and request authentication --------------------------------- | |
| def create_pairing_code( | |
| self, | |
| owner_id: str, | |
| *, | |
| private: bool = True, | |
| now: Optional[float] = None, | |
| ) -> PairingCode: | |
| """Mint a one-use code for an agent only its owner and invitees see. | |
| ``private`` is decided here, by the authenticated user generating the | |
| code, and not by the agent that redeems it. A worker cannot make itself | |
| private to hide, and cannot make itself public to be seen. | |
| **The default was ``False`` and that was a hole, not a preference.** | |
| The one caller in the product, ``DistinctUI.pairing_code``, did not | |
| pass the argument, so every worker anybody paired was visible to every | |
| signed-in stranger and to every signed-out visitor, and the access code | |
| that is supposed to decide who may use a machine granted access nobody | |
| needed. The feature was implemented, tested and off. It is now on | |
| unless somebody types the word ``private=False``, which is the right | |
| way round: a machine becoming public to the internet should be a | |
| sentence somebody wrote, not a default they inherited. | |
| """ | |
| owner_id = _single_line_identifier(owner_id, "owner_id") | |
| with self._lock: | |
| current = self._now(now) | |
| self._cleanup_pairing_codes_locked(current) | |
| if len(self._pairing_codes) >= self.limits.max_pairing_codes: | |
| raise CapacityError("too many active pairing codes") | |
| # THE GLOBAL CAP PROTECTS NOBODY ON ITS OWN. | |
| # | |
| # Filling it is cheap — one signed-in account asking in a loop — | |
| # and the codes live ten minutes, so the pool can be held full and | |
| # every other person on this server is then unable to pair a | |
| # worker at all. A per-owner share makes that account spend only | |
| # its own allowance. Counted after the expiry sweep above, so the | |
| # number is codes still alive rather than codes ever asked for. | |
| mine = sum( | |
| 1 for record in self._pairing_codes.values() if record.owner_id == owner_id | |
| ) | |
| if mine >= self.limits.max_pairing_codes_per_owner: | |
| raise CapacityError( | |
| "you already have the maximum number of unused pairing codes; " | |
| "use one, or wait for them to expire" | |
| ) | |
| while True: | |
| code = new_pairing_code() | |
| digest = pairing_code_digest(code, self._pairing_pepper) | |
| if digest not in self._pairing_codes: | |
| break | |
| expires_at = current + self.limits.pairing_ttl_seconds | |
| self._pairing_codes[digest] = PairingRecord(owner_id, expires_at, bool(private)) | |
| return PairingCode(code=code, expires_at=expires_at) | |
| def pair_agent( | |
| self, | |
| pairing_code: str, | |
| capabilities: AgentCapabilities, | |
| *, | |
| signing_public_key: bytes = b"", | |
| access_code: str = "", | |
| now: Optional[float] = None, | |
| ) -> AgentCredential: | |
| """Register an agent, and remember how to recognise its access code. | |
| ``access_code`` is generated by the agent when it starts and printed | |
| for the people who are allowed to use it. Only its digest is kept: the | |
| server has no use for the code itself and every reason not to hold it. | |
| """ | |
| if not isinstance(capabilities, AgentCapabilities): | |
| raise ValidationError("capabilities must be AgentCapabilities") | |
| with self._lock: | |
| current = self._now(now) | |
| self._cleanup_pairing_codes_locked(current) | |
| try: | |
| digest = pairing_code_digest(pairing_code, self._pairing_pepper) | |
| except ValueError as exc: | |
| raise AuthenticationError("pairing code is invalid or expired") from exc | |
| pairing = self._pairing_codes.get(digest) | |
| if pairing is None or pairing.expires_at <= current: | |
| self._pairing_codes.pop(digest, None) | |
| raise AuthenticationError("pairing code is invalid or expired") | |
| if len(self._agents) >= self.limits.max_agents: | |
| raise CapacityError("agent registry is full") | |
| agent_id = capabilities.agent_id or _new_id("agent") | |
| if agent_id in self._agents: | |
| raise ConflictError("agent id is already registered") | |
| if not capabilities.models: | |
| raise ValidationError("agent must advertise at least one model") | |
| if len(capabilities.models) > self.limits.max_models_per_agent: | |
| raise CapacityError("agent advertises too many models") | |
| if len(capabilities.tools) > self.limits.max_tools_per_agent: | |
| raise CapacityError("agent advertises too many tools") | |
| if not capabilities.energy_provider.strip(): | |
| raise ValidationError("agent must describe its energy provider") | |
| if len(capabilities.energy_provider) > 256: | |
| raise ValidationError("energy provider description is too long") | |
| if not capabilities.agent_id: | |
| capabilities = replace(capabilities, agent_id=agent_id) | |
| if signing_public_key and ( | |
| not isinstance(signing_public_key, (bytes, bytearray)) | |
| or len(signing_public_key) != 32 | |
| ): | |
| raise ValidationError("an Ed25519 public key is exactly 32 bytes") | |
| secret = generate_secret() | |
| self._agents[agent_id] = AgentRecord( | |
| signing_public_key=bytes(signing_public_key), | |
| owner_id=pairing.owner_id, | |
| access_digest=( | |
| access_code_digest(access_code, self._pairing_pepper) | |
| if access_code | |
| else "" | |
| ), | |
| capabilities=copy.deepcopy(capabilities), | |
| secret=secret, | |
| issued_at=current, | |
| last_snapshot_at=current, | |
| private=pairing.private, | |
| registered_capabilities=copy.deepcopy(capabilities), | |
| ) | |
| # Delete only after all validation and insertion succeed: one use. | |
| del self._pairing_codes[digest] | |
| return AgentCredential(agent_id, encode_secret(secret), current) | |
| def register_unclaimed( | |
| self, | |
| capabilities: AgentCapabilities, | |
| claim_code: str, | |
| *, | |
| signing_public_key: bytes = b"", | |
| access_code: str = "", | |
| private: bool = False, | |
| now: Optional[float] = None, | |
| ) -> AgentCredential: | |
| """Register a worker that has not been claimed by anybody yet. | |
| PAIRING RUNS THE OTHER WAY ROUND NOW, AND THIS IS WHY. | |
| The old flow asked a volunteer to sign in, find a pairing code on the | |
| page, and paste it into a terminal. The reliable failure was somebody | |
| typing a value of their own where the code went and being told, one | |
| program later, that their code was invalid -- which reads as the server | |
| rejecting *them* rather than as a copy that never happened. | |
| So the code travels the other way. The worker invents one, prints it, | |
| and registers here with nothing but its digest. It is issued a | |
| credential immediately, because it needs one to ask whether it has been | |
| claimed yet, but the record it gets is inert: ``owner_id`` is empty and | |
| ``claim_digest`` is set, which ``_visible_locked`` treats as "this | |
| worker does not exist". Nothing can list it, select it or send it work. | |
| That inertness is what makes an unauthenticated registration safe. What | |
| remains is resource use, which is bounded two ways: a hard ceiling on | |
| how many unclaimed records may exist at once, and a short expiry after | |
| which they are swept. | |
| """ | |
| if not isinstance(capabilities, AgentCapabilities): | |
| raise ValidationError("capabilities must be AgentCapabilities") | |
| with self._lock: | |
| current = self._now(now) | |
| self._cleanup_unclaimed_locked(current) | |
| try: | |
| digest = pairing_code_digest(claim_code, self._pairing_pepper) | |
| except ValueError as exc: | |
| raise ValidationError("claim code is not a usable code") from exc | |
| # A FULL POOL MUST NOT BE A CLOSED DOOR. | |
| # | |
| # A penetration test filled this pool from one unauthenticated | |
| # client at about six registrations a second, and the ceiling then | |
| # did exactly what a ceiling does: it refused everybody, including | |
| # the volunteer who had just started a worker for the first time. | |
| # Protecting the registry by turning the feature off for real users | |
| # is the wrong trade, and it is the trade an attacker was choosing. | |
| # | |
| # So a full pool evicts its oldest waiter instead. A legitimate | |
| # registration now always succeeds; what a flood costs is the slot | |
| # of whoever registered longest ago, and a flood mostly evicts its | |
| # own earlier entries because it owns nearly all of them. | |
| # | |
| # RESIDUAL RISK, STATED RATHER THAN PAPERED OVER. Sustained flooding | |
| # can still evict a real worker in the window between it printing a | |
| # code and its owner typing that code in. The worker says so, in | |
| # those words, and starting it again gives a fresh code. Closing | |
| # that gap properly needs per-source limiting, which this layer | |
| # cannot do: it has no address, by design. The place for it is the | |
| # route in front, and it is not pretended to be here. | |
| waiting = sorted( | |
| ( | |
| (record.claim_expires_at, agent_id) | |
| for agent_id, record in self._agents.items() | |
| if record.claim_digest | |
| ), | |
| ) | |
| while len(waiting) >= self.limits.max_unclaimed_agents: | |
| _, oldest = waiting.pop(0) | |
| del self._agents[oldest] | |
| if len(self._agents) >= self.limits.max_agents: | |
| raise CapacityError("agent registry is full") | |
| if any( | |
| hmac.compare_digest(record.claim_digest, digest) | |
| for record in self._agents.values() | |
| if record.claim_digest | |
| ): | |
| # Two workers offering one code would make the claim ambiguous, | |
| # and the second one to arrive is the one that must lose. | |
| raise ConflictError("that claim code is already waiting") | |
| agent_id = capabilities.agent_id or _new_id("agent") | |
| if agent_id in self._agents: | |
| raise ConflictError("agent id is already registered") | |
| if not capabilities.models: | |
| raise ValidationError("agent must advertise at least one model") | |
| if len(capabilities.models) > self.limits.max_models_per_agent: | |
| raise CapacityError("agent advertises too many models") | |
| if len(capabilities.tools) > self.limits.max_tools_per_agent: | |
| raise CapacityError("agent advertises too many tools") | |
| if not capabilities.energy_provider.strip(): | |
| raise ValidationError("agent must describe its energy provider") | |
| if len(capabilities.energy_provider) > 256: | |
| raise ValidationError("energy provider description is too long") | |
| if not capabilities.agent_id: | |
| capabilities = replace(capabilities, agent_id=agent_id) | |
| if signing_public_key and ( | |
| not isinstance(signing_public_key, (bytes, bytearray)) | |
| or len(signing_public_key) != 32 | |
| ): | |
| raise ValidationError("an Ed25519 public key is exactly 32 bytes") | |
| secret = generate_secret() | |
| self._agents[agent_id] = AgentRecord( | |
| signing_public_key=bytes(signing_public_key), | |
| owner_id="", | |
| access_digest=( | |
| access_code_digest(access_code, self._pairing_pepper) | |
| if access_code | |
| else "" | |
| ), | |
| capabilities=copy.deepcopy(capabilities), | |
| secret=secret, | |
| issued_at=current, | |
| last_snapshot_at=current, | |
| private=bool(private), | |
| registered_capabilities=copy.deepcopy(capabilities), | |
| claim_digest=digest, | |
| claim_expires_at=current + self.limits.claim_ttl_seconds, | |
| ) | |
| return AgentCredential(agent_id, encode_secret(secret), current) | |
| def claim_agent( | |
| self, | |
| claim_code: str, | |
| owner_id: str, | |
| *, | |
| now: Optional[float] = None, | |
| ) -> str: | |
| """Bind a waiting worker to the signed-in person who typed its code. | |
| Returns the agent id. The failure message says nothing about why, for | |
| the same reason ``redeem_access_code`` says nothing: an error that told | |
| a caller apart "no such code" from "expired" from "already claimed" | |
| would be an oracle, and this one can be asked without an account on the | |
| machine that printed the code. | |
| Signing in is the whole check. Ownership is what a claim confers -- the | |
| right to be sent other people's prompts is downstream of it -- so it | |
| cannot be granted to an identity the server has not authenticated. | |
| """ | |
| if not isinstance(owner_id, str) or not owner_id.strip(): | |
| raise AuthenticationError("that code is not valid") | |
| owner_id = _single_line_identifier(owner_id, "owner_id") | |
| with self._lock: | |
| current = self._now(now) | |
| self._cleanup_unclaimed_locked(current) | |
| try: | |
| digest = pairing_code_digest(claim_code, self._pairing_pepper) | |
| except ValueError as exc: | |
| raise AuthenticationError("that code is not valid") from exc | |
| for agent_id, record in self._agents.items(): | |
| if not record.claim_digest: | |
| continue | |
| if not hmac.compare_digest(record.claim_digest, digest): | |
| continue | |
| if record.revoked or record.claim_expires_at <= current: | |
| raise AuthenticationError("that code is not valid") | |
| record.owner_id = owner_id | |
| # One use. Clearing the digest is also what makes the record | |
| # visible, so the two can never disagree. | |
| record.claim_digest = "" | |
| record.claim_expires_at = 0.0 | |
| return agent_id | |
| raise AuthenticationError("that code is not valid") | |
| def claim_state(self, agent_id: str, *, now: Optional[float] = None) -> str: | |
| """``"waiting"``, ``"claimed"``, or ``"gone"`` for a worker's own poll. | |
| The worker asks about itself, having authenticated as itself, so this | |
| tells it only what it already put there. It exists so a worker can stop | |
| printing "waiting to be claimed" at the moment somebody claims it. | |
| """ | |
| with self._lock: | |
| current = self._now(now) | |
| self._cleanup_unclaimed_locked(current) | |
| record = self._agents.get(agent_id) | |
| if record is None or record.revoked: | |
| return "gone" | |
| return "waiting" if record.claim_digest else "claimed" | |
| def _phantom_record(self) -> AgentRecord: | |
| """A record for an agent that does not exist, so the checks can run. | |
| A fresh random secret each time, and never stored. It exists only so | |
| that a request naming an unknown agent goes through exactly the checks | |
| a request naming a known one does, and therefore produces exactly the | |
| same answers. See `verify_agent_request` for why that matters. | |
| The capabilities are the smallest valid ones rather than anything | |
| meaningful: nothing reads them on this path, and a phantom that looked | |
| like a plausible machine would be a thing that could be mistaken for | |
| one somewhere else later. | |
| """ | |
| return AgentRecord( | |
| owner_id="phantom", | |
| capabilities=_PHANTOM_CAPABILITIES, | |
| secret=secrets.token_bytes(32), | |
| issued_at=0.0, | |
| last_snapshot_at=0.0, | |
| ) | |
| def verify_agent_request( | |
| self, | |
| agent_id: str, | |
| timestamp: float, | |
| nonce: str, | |
| method: str, | |
| path: str, | |
| body: Union[bytes, str], | |
| signature: str, | |
| *, | |
| now: Optional[float] = None, | |
| ) -> None: | |
| """Authenticate once and reserve the nonce atomically. | |
| Success returns ``None``. A nonce is recorded only after a valid HMAC, | |
| preventing unauthenticated requests from exhausting an agent's cache. | |
| """ | |
| with self._lock: | |
| current = self._now(now) | |
| # AN UNKNOWN AGENT IS CHECKED, NOT REFUSED. | |
| # | |
| # This used to raise "agent was not found" straight away, and a | |
| # known agent with a bad signature raised something else. So an | |
| # unauthenticated caller could ask about any id and learn whether a | |
| # worker with that id exists, which is exactly what | |
| # `_visible_locked` calls "whether `viewer_id` may know this agent | |
| # exists at all" and refuses everywhere else. A default agent id is | |
| # the machine's hostname, so the ids worth guessing are guessable. | |
| # | |
| # Refusing with one uniform message was not enough on its own: the | |
| # checks below reject a bad timestamp or a malformed nonce before | |
| # they ever reach the signature, so a caller sending a deliberately | |
| # bad timestamp still saw two different answers. The only way for | |
| # the answers to match is for the same checks to run either way, so | |
| # an unknown id is given a phantom record with a random secret and | |
| # put through the identical sequence. It fails at the signature, | |
| # which is where a known agent with the wrong secret fails, with | |
| # the same sentence. | |
| # | |
| # The phantom is built fresh per call and never stored: it must not | |
| # accumulate nonces, and two calls about the same unknown id must | |
| # not be able to tell each other apart either. | |
| known = True | |
| try: | |
| record = self._agent_locked(agent_id) | |
| except NotFoundError: | |
| known = False | |
| record = self._phantom_record() | |
| if record.revoked: | |
| raise AuthenticationError("agent credential is revoked") | |
| try: | |
| validate_nonce(nonce) | |
| message = canonical_request(agent_id, timestamp, nonce, method, path, body) | |
| except (TypeError, ValueError) as exc: | |
| raise AuthenticationError(str(exc)) from exc | |
| if not math.isfinite(float(timestamp)): | |
| raise AuthenticationError("timestamp must be finite") | |
| if abs(current - int(timestamp)) > self.limits.request_clock_skew_seconds: | |
| raise AuthenticationError("request timestamp is outside the accepted window") | |
| self._cleanup_agent_nonces_locked(record, current) | |
| if nonce in record.nonces: | |
| raise AuthenticationError("request nonce has already been used") | |
| if len(record.nonces) >= self.limits.max_nonces_per_agent: | |
| raise CapacityError("agent nonce cache is full") | |
| if not isinstance(signature, str): | |
| raise AuthenticationError("signature must be hexadecimal text") | |
| expected = hmac.digest(record.secret, message, "sha256").hex() | |
| matched = hmac.compare_digest(expected, signature.lower()) | |
| if not matched or not known: | |
| # `not known` as well as `not matched`: a phantom's random | |
| # secret will not match, but saying so explicitly means a | |
| # freak collision cannot authenticate a caller against an agent | |
| # that does not exist. | |
| raise AuthenticationError(UNAUTHENTICATED) | |
| record.nonces[nonce] = current + self.limits.nonce_ttl_seconds | |
| def issue_catalogue( | |
| self, | |
| challenge: str, | |
| *, | |
| secret: Optional[bytes] = None, | |
| now: Optional[float] = None, | |
| ) -> ServerCatalogue: | |
| """State what this server may ask for, bound to the agent's challenge. | |
| Echoing the agent's own nonce is what makes the catalogue fresh: a | |
| catalogue captured from an earlier pairing carries an earlier challenge | |
| and is refused before an operator is ever shown it. When ``secret`` is | |
| supplied the catalogue is signed with it, which is the pairing case. | |
| """ | |
| with self._lock: | |
| catalogue = ServerCatalogue( | |
| server_id=self.server_id, | |
| challenge=challenge, | |
| allowance=self.catalogue, | |
| issued_at=self._now(now), | |
| ) | |
| return catalogue if secret is None else catalogue.signed(secret) | |
| def catalogue_for_agent( | |
| self, | |
| agent_id: str, | |
| challenge: str, | |
| *, | |
| now: Optional[float] = None, | |
| ) -> ServerCatalogue: | |
| """The signed catalogue one paired agent asked for. | |
| Signed with that agent's own channel secret, bound to the agent's own | |
| challenge nonce, so the operator can be shown a catalogue that is | |
| provably fresh and provably from the server the agent paired with. | |
| The secret never leaves the control plane. | |
| """ | |
| with self._lock: | |
| record = self._agent_locked(agent_id) | |
| if record.revoked or not record.secret: | |
| raise AuthenticationError("agent credential is revoked") | |
| return self.issue_catalogue(challenge, secret=record.secret, now=now) | |
| def sign_response( | |
| self, | |
| agent_id: str, | |
| path: str, | |
| request_nonce: str, | |
| body: Mapping[str, Any], | |
| ) -> Optional[dict]: | |
| """Prove a response came from this server, to this agent, for this call. | |
| The secret never leaves the control plane: callers hand in a body and | |
| receive a proof. Returns ``None`` for an unknown or revoked agent, | |
| because there is no key to sign with and inventing one would make an | |
| unauthenticated answer look authenticated. | |
| """ | |
| with self._lock: | |
| record = self._agents.get(agent_id) | |
| if record is None or record.revoked or not record.secret: | |
| return None | |
| return _sign_response_payload(record.secret, path, request_nonce, body) | |
| def verify_signed_snapshot(self, agent_id: str, envelope: Any) -> bool: | |
| """Check the detached signature and refuse a replayed or stale counter. | |
| Both halves matter. The signature says the agent made this statement; | |
| the counter says it is the newest such statement. A captured snapshot | |
| replayed later carries a counter that is no longer ahead, so it is | |
| refused even though its signature is perfectly valid. | |
| """ | |
| if not isinstance(envelope, SignedSnapshot): | |
| raise ValidationError("signed snapshot must be a SignedSnapshot") | |
| with self._lock: | |
| record = self._agent_locked(agent_id) | |
| if record.revoked or not record.secret: | |
| return False | |
| # Verified against the pinned public key, under the algorithm the | |
| # server expects rather than the one the message names. An envelope | |
| # arriving as HMAC is refused even though the server holds a key | |
| # that could check it: accepting it would let the server itself | |
| # have produced the advertisement. | |
| if not record.signing_public_key: | |
| return False | |
| if not envelope.verify(record.signing_public_key, algorithm=ALG_ED25519): | |
| return False | |
| if envelope.counter <= record.snapshot_counter: | |
| return False | |
| record.snapshot_counter = envelope.counter | |
| return True | |
| def record_advertised(self, agent_id: str, advertised: Allowance) -> Allowance: | |
| """Narrow an agent's advertised capability to what its operator approved. | |
| Narrowing only. The set registered at pairing is a ceiling this can | |
| never rise above, so a compromised worker cannot advertise its way into | |
| work its operator never offered, and an approval can only ever take | |
| capability away. | |
| """ | |
| if not isinstance(advertised, Allowance): | |
| raise ValidationError("advertised capability must be an Allowance") | |
| with self._lock: | |
| record = self._agent_locked(agent_id) | |
| ceiling = record.registered_capabilities or record.capabilities | |
| registered = Allowance( | |
| models=ceiling.models, | |
| tools=tuple( | |
| item if "@" in item else f"{item}@1.0.0" | |
| for item in ceiling.tools | |
| ), | |
| ) | |
| effective = advertised.intersect(registered) | |
| record.capabilities = replace( | |
| record.capabilities, | |
| models=effective.models, | |
| tools=effective.tools, | |
| ) | |
| return effective | |
| def rotate_agent_credential( | |
| self, agent_id: str, *, now: Optional[float] = None | |
| ) -> AgentCredential: | |
| with self._lock: | |
| current = self._now(now) | |
| record = self._agent_locked(agent_id) | |
| secret = generate_secret() | |
| record.secret = secret | |
| record.issued_at = current | |
| record.nonces.clear() | |
| record.revoked = False | |
| return AgentCredential(agent_id, encode_secret(secret), current) | |
| def revoke_agent(self, agent_id: str) -> None: | |
| with self._lock: | |
| record = self._agent_locked(agent_id) | |
| record.revoked = True | |
| record.secret = b"" | |
| record.nonces.clear() | |
| # Sessions and conversations ---------------------------------------- | |
| def create_session( | |
| self, session_id: Optional[str] = None, *, now: Optional[float] = None | |
| ) -> SessionView: | |
| with self._lock: | |
| current = self._now(now) | |
| session_id = session_id or _new_id("session") | |
| _single_line_identifier(session_id, "session_id") | |
| if session_id in self._sessions: | |
| return self._session_view_locked(self._sessions[session_id]) | |
| if len(self._sessions) >= self.limits.max_sessions: | |
| raise CapacityError("session registry is full") | |
| record = SessionRecord(session_id=session_id, created_at=current) | |
| self._sessions[session_id] = record | |
| return self._session_view_locked(record) | |
| def delete_session( | |
| self, | |
| session_id: str, | |
| *, | |
| now: Optional[float] = None, | |
| ) -> int: | |
| """Irreversibly remove a browser session and all of its prompt data.""" | |
| with self._lock: | |
| current = self._now(now) | |
| session = self._sessions.pop(session_id, None) | |
| if session is None: | |
| return 0 | |
| job_ids: set[str] = set() | |
| for conversation_id in session.conversation_ids: | |
| conversation = self._conversations.pop((session_id, conversation_id), None) | |
| if conversation is not None: | |
| job_ids.update(conversation.job_ids) | |
| for children in self._children.values(): | |
| children.difference_update(job_ids) | |
| for job_id in job_ids: | |
| job = self._jobs.pop(job_id, None) | |
| if job is not None and job.status not in _TERMINAL_STATUSES: | |
| self._orphan_cancellations.setdefault( | |
| job.spec.target_agent_id, | |
| {}, | |
| )[job_id] = current + self.limits.job_ttl_seconds | |
| self._children.pop(job_id, None) | |
| return len(job_ids) | |
| def delete_conversation( | |
| self, | |
| session_id: str, | |
| conversation_id: str, | |
| *, | |
| now: Optional[float] = None, | |
| ) -> int: | |
| """Remove one conversation and the jobs it held. Returns how many. | |
| The same erasure :meth:`delete_session` performs, narrowed to a single | |
| conversation, because a person wanting one exchange gone should not | |
| have to take the whole session with it. | |
| A job still in flight is cancelled rather than merely forgotten. Its | |
| worker is holding it and will come back for it, so the cancellation is | |
| left where that worker will find it on its next poll; dropping the | |
| record alone would leave the machine working on an answer for a | |
| conversation that no longer exists. | |
| Unknown ids are not an error. A conversation whose runs have all | |
| finished is cleaned up on its own, so a browser asking to delete one | |
| that has already gone is asking for a state that already holds. | |
| """ | |
| with self._lock: | |
| current = self._now(now) | |
| session = self._sessions.get(session_id) | |
| if session is None: | |
| return 0 | |
| conversation = self._conversations.pop((session_id, conversation_id), None) | |
| if conversation is None: | |
| return 0 | |
| if conversation_id in session.conversation_ids: | |
| session.conversation_ids.remove(conversation_id) | |
| job_ids = set(conversation.job_ids) | |
| for children in self._children.values(): | |
| children.difference_update(job_ids) | |
| for job_id in job_ids: | |
| job = self._jobs.pop(job_id, None) | |
| if job is not None and job.status not in _TERMINAL_STATUSES: | |
| self._orphan_cancellations.setdefault( | |
| job.spec.target_agent_id, | |
| {}, | |
| )[job_id] = current + self.limits.job_ttl_seconds | |
| self._children.pop(job_id, None) | |
| return len(job_ids) | |
| def create_conversation( | |
| self, | |
| session_id: str, | |
| conversation_id: Optional[str] = None, | |
| *, | |
| now: Optional[float] = None, | |
| ) -> ConversationView: | |
| with self._lock: | |
| current = self._now(now) | |
| session = self._session_locked(session_id) | |
| conversation_id = conversation_id or _new_id("conversation") | |
| _single_line_identifier(conversation_id, "conversation_id") | |
| key = (session_id, conversation_id) | |
| if key in self._conversations: | |
| return self._conversation_view_locked(self._conversations[key]) | |
| if len(session.conversation_ids) >= self.limits.max_conversations_per_session: | |
| raise CapacityError("session has too many conversations") | |
| record = ConversationRecord(session_id, conversation_id, current) | |
| self._conversations[key] = record | |
| session.conversation_ids.append(conversation_id) | |
| return self._conversation_view_locked(record) | |
| def get_session(self, session_id: str) -> SessionView: | |
| with self._lock: | |
| return self._session_view_locked(self._session_locked(session_id)) | |
| def get_conversation(self, session_id: str, conversation_id: str) -> ConversationView: | |
| with self._lock: | |
| return self._conversation_view_locked( | |
| self._conversation_locked(session_id, conversation_id) | |
| ) | |
| # Agent registry and snapshots -------------------------------------- | |
| def record_snapshot( | |
| self, | |
| agent_id: str, | |
| snapshot: AgentSnapshot, | |
| *, | |
| now: Optional[float] = None, | |
| ) -> AgentView: | |
| if not isinstance(snapshot, AgentSnapshot): | |
| raise ValidationError("snapshot must be AgentSnapshot") | |
| with self._lock: | |
| current = self._now(now) | |
| record = self._agent_locked(agent_id) | |
| if record.revoked: | |
| raise AuthenticationError("agent credential is revoked") | |
| if snapshot.agent_id != agent_id: | |
| raise ValidationError("snapshot agent id does not match credential") | |
| if snapshot.queue_capacity > record.capabilities.queue_capacity: | |
| raise ValidationError("snapshot exceeds the registered queue capacity") | |
| active = tuple(snapshot.active_job_ids) | |
| queued = tuple(snapshot.queued_job_ids) | |
| if len(active) > record.capabilities.max_concurrency: | |
| raise ValidationError("snapshot exceeds registered concurrency") | |
| if len(active) + len(queued) > snapshot.queue_capacity: | |
| raise ValidationError("snapshot exceeds its declared queue capacity") | |
| all_ids = active + queued | |
| if len(set(all_ids)) != len(all_ids): | |
| raise ValidationError("snapshot contains duplicate job ids") | |
| if not set(snapshot.progress).issubset(set(all_ids)): | |
| raise ValidationError("snapshot progress references an unlisted job") | |
| if not set(snapshot.progress_notes).issubset(set(all_ids)): | |
| raise ValidationError("snapshot progress notes reference an unlisted job") | |
| orphan_ids = set(self._orphan_cancellations.get(agent_id, ())) | |
| stale_cancelled: set[str] = set() | |
| # Validate the full snapshot before changing any job state. A | |
| # just-cancelled job is allowed for one or more stale snapshots so | |
| # this same sync can return its cancellation instead of deadlocking | |
| # the worker behind a validation error. | |
| for job_id in all_ids: | |
| job = self._jobs.get(job_id) | |
| if job_id in orphan_ids: | |
| stale_cancelled.add(job_id) | |
| continue | |
| if job is None or job.selection.target_agent_id != agent_id: | |
| raise ValidationError("snapshot references a job not owned by this agent") | |
| if job.accepted_by == agent_id and job.status is JobStatus.CANCELLED: | |
| stale_cancelled.add(job_id) | |
| continue | |
| if job.accepted_by != agent_id or job.status not in ( | |
| _AGENT_EXECUTION_STATUSES | {JobStatus.STRANDED} | |
| ): | |
| # STRANDED is accepted here deliberately: an agent that | |
| # went silent and came back is still holding the job, and | |
| # this same pass returns the job to its live status below. | |
| raise ValidationError("snapshot references a job that was not accepted") | |
| for job_id in queued: | |
| if job_id not in stale_cancelled and self._jobs[job_id].status is JobStatus.RUNNING: | |
| raise ConflictError("a running job cannot return to the queue") | |
| active = tuple(job_id for job_id in active if job_id not in stale_cancelled) | |
| queued = tuple(job_id for job_id in queued if job_id not in stale_cancelled) | |
| for job in self._jobs.values(): | |
| if job.selection.target_agent_id == agent_id: | |
| job.queue_position = None | |
| for job_id in active: | |
| job = self._jobs[job_id] | |
| job.status = JobStatus.RUNNING | |
| job.updated_at = current | |
| for position, job_id in enumerate(queued, start=1): | |
| job = self._jobs[job_id] | |
| job.status = JobStatus.QUEUED | |
| job.queue_position = position | |
| job.updated_at = current | |
| clean_snapshot = replace( | |
| snapshot, | |
| active_job_ids=active, | |
| queued_job_ids=queued, | |
| progress={ | |
| job_id: value | |
| for job_id, value in snapshot.progress.items() | |
| if job_id not in stale_cancelled | |
| }, | |
| progress_notes={ | |
| job_id: value | |
| for job_id, value in snapshot.progress_notes.items() | |
| if job_id not in stale_cancelled | |
| }, | |
| ) | |
| record.snapshot = AgentSnapshot.from_dict(clean_snapshot.to_dict()) | |
| # Client last_seen is informational; liveness uses receipt time. | |
| record.last_snapshot_at = current | |
| tombstones = self._orphan_cancellations.get(agent_id) | |
| if tombstones is not None: | |
| for job_id in tuple(tombstones): | |
| if job_id not in all_ids: | |
| del tombstones[job_id] | |
| if not tombstones: | |
| del self._orphan_cancellations[agent_id] | |
| return self._agent_view_locked(agent_id, record, current) | |
| def get_agent( | |
| self, | |
| agent_id: str, | |
| *, | |
| viewer_id: Optional[str] = None, | |
| now: Optional[float] = None, | |
| ) -> AgentView: | |
| with self._lock: | |
| record = self._agent_locked(agent_id) | |
| if not self._visible_locked(record, viewer_id): | |
| raise NotFoundError(_AGENT_NOT_FOUND) | |
| return self._agent_view_locked(agent_id, record, self._now(now)) | |
| def _visible_locked(record: AgentRecord, viewer_id: Optional[str]) -> bool: | |
| """Whether ``viewer_id`` may know this agent exists at all. | |
| A public agent is visible to everyone, including signed-out users. A | |
| private one is visible only to the authenticated identity that generated | |
| its pairing code. ``viewer_id=None`` means an unauthenticated or | |
| unscoped caller and never sees a private agent. | |
| The check is a plain comparison with no early exit that depends on | |
| whether the agent exists, so the cost of asking about a private agent | |
| does not differ from asking about a missing one. | |
| """ | |
| if record.revoked: | |
| # A revoked credential is a worker this server will not talk to, so | |
| # nobody should still see it listed. It was visible to its owner | |
| # and its grantees, which is a worker that cannot be used sitting | |
| # in somebody's list looking like one that can. | |
| return False | |
| if record.claim_digest: | |
| # AN UNCLAIMED WORKER DOES NOT EXIST YET. | |
| # | |
| # Registering without a pairing code is the one operation on this | |
| # server that needs no account, which is what lets a volunteer | |
| # start a worker and claim it afterwards. The price is that anybody | |
| # can create one of these records, so until a signed-in person | |
| # enters its claim code it must be visible to nobody and reachable | |
| # by nothing. | |
| # | |
| # This is the single place that has to be right. Every listing, | |
| # every selection, both dispatch paths and the per-agent view all | |
| # ask this function first, so one `False` here keeps an unclaimed | |
| # worker out of all of them at once rather than relying on six | |
| # call sites to remember. | |
| return False | |
| if not record.private: | |
| return True | |
| if not viewer_id: | |
| return False | |
| if record.owner_id == viewer_id: | |
| return True | |
| # Redeeming the agent's own access code is the other way in, and the | |
| # only one that does not require having created the pairing code. | |
| return viewer_id in record.granted | |
| def redeem_access_code( | |
| self, | |
| access_code: str, | |
| viewer_id: str, | |
| *, | |
| now: Optional[float] = None, | |
| ) -> str: | |
| """Give a signed-in identity access to the agent that printed this code. | |
| Returns the agent id. Raises ``AuthenticationError`` when no agent | |
| matches, and deliberately says nothing about why: an error that | |
| distinguished "no such code" from "expired" from "already yours" would | |
| be an oracle worth querying. | |
| The comparison is over digests of a 192-bit random string, so guessing | |
| is not a strategy and enumeration is not possible: there is no order to | |
| walk and nothing shorter to try. | |
| """ | |
| if not isinstance(viewer_id, str) or not viewer_id.strip(): | |
| # No identity, no grant. This is the whole point of signing in. | |
| raise AuthenticationError("that code is not valid") | |
| with self._lock: | |
| self._expire_jobs_locked(self._now(now)) | |
| try: | |
| digest = access_code_digest(access_code, self._pairing_pepper) | |
| except ValueError as exc: | |
| raise AuthenticationError("that code is not valid") from exc | |
| for agent_id, record in self._agents.items(): | |
| if not record.access_digest: | |
| continue | |
| if hmac.compare_digest(record.access_digest, digest): | |
| if record.revoked or record.claim_digest: | |
| # An unclaimed worker prints its access code before | |
| # anybody owns it. Honouring that code here would let a | |
| # stranger take a share of a machine whose operator has | |
| # not finished setting it up, so it grants nothing | |
| # until the worker is claimed. | |
| raise AuthenticationError("that code is not valid") | |
| record.granted.add(viewer_id) | |
| return agent_id | |
| raise AuthenticationError("that code is not valid") | |
| def granted_agents(self, viewer_id: Optional[str]) -> Tuple[str, ...]: | |
| """Which agents this identity has redeemed a code for.""" | |
| if not viewer_id: | |
| return () | |
| with self._lock: | |
| return tuple( | |
| agent_id | |
| for agent_id, record in self._agents.items() | |
| if viewer_id in record.granted and not record.claim_digest | |
| ) | |
| def list_agents( | |
| self, | |
| *, | |
| viewer_id: Optional[str] = None, | |
| now: Optional[float] = None, | |
| ) -> Tuple[AgentView, ...]: | |
| """Every agent this viewer may see, and no trace of the rest. | |
| The agent table, its queue depths and its wait estimates are all | |
| rendered from this, so filtering here is what stops a private agent | |
| leaking through a capacity figure as well as through a name. | |
| """ | |
| with self._lock: | |
| current = self._now(now) | |
| return tuple( | |
| self._agent_view_locked(agent_id, self._agents[agent_id], current) | |
| for agent_id in sorted(self._agents) | |
| if self._visible_locked(self._agents[agent_id], viewer_id) | |
| ) | |
| def compatible_agents( | |
| self, | |
| model_id: str, | |
| allowed_tools: Sequence[ToolSelection] = (), | |
| *, | |
| viewer_id: Optional[str] = None, | |
| include_offline: bool = False, | |
| now: Optional[float] = None, | |
| ) -> Tuple[AgentView, ...]: | |
| tools = self._normalize_tools(allowed_tools) | |
| with self._lock: | |
| current = self._now(now) | |
| candidates = [] | |
| for agent_id, record in self._agents.items(): | |
| if not self._visible_locked(record, viewer_id): | |
| continue | |
| if not self._is_compatible_locked(record, model_id, tools): | |
| continue | |
| if not include_offline and not self._eligible_for_offer_locked(record, current): | |
| continue | |
| candidates.append( | |
| (self._agent_load_key_locked(agent_id, record), agent_id, record) | |
| ) | |
| candidates.sort(key=lambda value: (value[0], value[1])) | |
| return tuple( | |
| self._agent_view_locked(agent_id, record, current) | |
| for _, agent_id, record in candidates | |
| ) | |
| # Jobs and dispatch -------------------------------------------------- | |
| def submit_job( | |
| self, | |
| session_id: str, | |
| conversation_id: str, | |
| prompt: str, | |
| model_id: str, | |
| allowed_tools: Sequence[ToolSelection] = (), | |
| *, | |
| parent_job_id: Optional[str] = None, | |
| target_agent_id: Optional[str] = None, | |
| messages: Sequence[ChatMessage] = (), | |
| inference_limits: Optional[Mapping[str, Any]] = None, | |
| viewer_id: Optional[str] = None, | |
| avoid_agent_ids: Sequence[str] = (), | |
| now: Optional[float] = None, | |
| ) -> JobView: | |
| tools = self._normalize_tools(allowed_tools) | |
| with self._lock: | |
| current = self._now(now) | |
| conversation = self._conversation_locked(session_id, conversation_id) | |
| if len(self._jobs) >= self.limits.max_jobs: | |
| raise CapacityError("job registry is full") | |
| if len(conversation.job_ids) >= self.limits.max_jobs_per_conversation: | |
| raise CapacityError("conversation has too many jobs") | |
| parent: Optional[JobRecord] = None | |
| if parent_job_id is not None: | |
| parent = self._job_in_scope_locked(session_id, parent_job_id) | |
| if parent.spec.conversation_id != conversation_id: | |
| raise DependencyError("parent job belongs to another conversation") | |
| if parent.status in { | |
| JobStatus.FAILED, | |
| JobStatus.CANCELLED, | |
| JobStatus.EXPIRED, | |
| }: | |
| raise DependencyError("parent job cannot complete successfully") | |
| born_stranded = False | |
| if target_agent_id is None: | |
| try: | |
| target_agent_id = self._choose_agent_locked( | |
| model_id, tools, current, viewer_id=viewer_id, avoid=avoid_agent_ids | |
| ) | |
| except (CapacityError, NoCompatibleAgentError): | |
| # No compatible agent right now is not a refusal: the run | |
| # is queued as STRANDED and assigned automatically when a | |
| # compatible agent comes online. This is what lets the | |
| # user submit first and start their worker second. | |
| target_agent_id = UNASSIGNED_AGENT | |
| born_stranded = True | |
| else: | |
| target = self._agent_locked(target_agent_id) | |
| if not self._visible_locked(target, viewer_id): | |
| # Deliberately the same error type and the same text that | |
| # _agent_locked raises for an id that does not exist. | |
| # Distinguishing "not yours" from "no such agent" would let | |
| # anyone enumerate other people's private workers. | |
| raise NotFoundError(_AGENT_NOT_FOUND) | |
| if not self._is_compatible_locked(target, model_id, tools): | |
| raise NoCompatibleAgentError( | |
| "target agent cannot run the exact model and tool versions" | |
| ) | |
| if not self._eligible_for_offer_locked(target, current): | |
| raise NoCompatibleAgentError("target agent is not currently eligible") | |
| if self._outstanding_count_locked(target_agent_id) >= self._capacity_locked(target): | |
| raise CapacityError("target agent has no queue capacity") | |
| job_id = _new_id("job") | |
| if born_stranded: | |
| status = ( | |
| JobStatus.STRANDED | |
| if parent is None or parent.status is JobStatus.COMPLETED | |
| else JobStatus.BLOCKED | |
| ) | |
| else: | |
| status = ( | |
| JobStatus.OFFERED | |
| if parent is None or parent.status is JobStatus.COMPLETED | |
| else JobStatus.BLOCKED | |
| ) | |
| limits = dict(inference_limits or {}) | |
| if viewer_id and not born_stranded: | |
| # The guard's pseudonymous handle: computed here, where the | |
| # real identity is, keyed per agent so two agents cannot | |
| # correlate a user by comparing keys. The agent's deny list | |
| # keys on it; the agent never sees the identity behind it. | |
| target_record = self._agents.get(target_agent_id) | |
| if target_record is not None and target_record.secret: | |
| limits["user_key"] = derive_user_key(target_record.secret, viewer_id) | |
| try: | |
| spec = JobSpec( | |
| id=job_id, | |
| session_id=session_id, | |
| conversation_id=conversation_id, | |
| parent_job_id=parent_job_id, | |
| prompt=prompt, | |
| model_id=model_id, | |
| allowed_tools=tools, | |
| target_agent_id=target_agent_id, | |
| created_at=current, | |
| messages=tuple(messages), | |
| limits=limits, | |
| ) | |
| except (TypeError, ValueError) as exc: | |
| raise ValidationError(str(exc)) from exc | |
| selection = RunSelection( | |
| model_id=spec.model_id, | |
| tool_versions=tuple(_tool_capability(tool) for tool in spec.allowed_tools), | |
| target_agent_id=spec.target_agent_id, | |
| ) | |
| record = JobRecord( | |
| spec=spec, | |
| selection=selection, | |
| status=status, | |
| updated_at=current, | |
| expires_at=current + self.limits.job_ttl_seconds, | |
| viewer_subject=viewer_id, | |
| ) | |
| self._jobs[job_id] = record | |
| self._children.setdefault(job_id, set()) | |
| if parent_job_id is not None: | |
| self._children.setdefault(parent_job_id, set()).add(job_id) | |
| conversation.job_ids.append(job_id) | |
| return self._job_view_locked(record) | |
| def offers_for_agent( | |
| self, agent_id: str, *, now: Optional[float] = None | |
| ) -> Tuple[JobView, ...]: | |
| with self._lock: | |
| current = self._now(now) | |
| self._expire_jobs_locked(current) | |
| self._agent_locked(agent_id) | |
| jobs = [ | |
| job | |
| for job in self._jobs.values() | |
| if job.selection.target_agent_id == agent_id | |
| and job.status is JobStatus.OFFERED | |
| ] | |
| jobs.sort(key=lambda job: (job.spec.created_at, job.spec.id)) | |
| return tuple(self._job_view_locked(job) for job in jobs) | |
| def accept_offer( | |
| self, | |
| agent_id: str, | |
| job_id: str, | |
| *, | |
| queue_position: Optional[int] = None, | |
| now: Optional[float] = None, | |
| ) -> JobView: | |
| """Record an admission the agent has already decided. | |
| ``queue_position`` is the position the agent's own queue assigned at | |
| admission time. Accepting it here is what lets the user see "Queued #2" | |
| immediately instead of a blank until the next snapshot arrives. The | |
| next ``record_snapshot`` remains authoritative and may overwrite it -- | |
| the server still never *reorders*, it just stops showing nothing. | |
| """ | |
| with self._lock: | |
| current = self._now(now) | |
| self._expire_jobs_locked(current) | |
| agent = self._agent_locked(agent_id) | |
| if agent.revoked: | |
| raise AuthenticationError("agent credential is revoked") | |
| job = self._job_locked(job_id) | |
| if job.selection.target_agent_id != agent_id: | |
| raise ConflictError("offer is targeted to another agent") | |
| position = self._validated_position_locked(queue_position, agent) | |
| if job.status in _AGENT_EXECUTION_STATUSES and job.accepted_by == agent_id: | |
| return self._job_view_locked(job) # idempotent retry | |
| if job.status is not JobStatus.OFFERED: | |
| raise ConflictError(f"job cannot be accepted while {job.status.value}") | |
| if not self._is_compatible_locked(agent, job.spec.model_id, job.spec.allowed_tools): | |
| # Capability changes never silently remove tools or redirect work. | |
| raise NoCompatibleAgentError("agent no longer matches the immutable job selection") | |
| job.accepted_by = agent_id | |
| job.status = JobStatus.ACCEPTED | |
| job.queue_position = position | |
| job.updated_at = current | |
| return self._job_view_locked(job) | |
| def _validated_position_locked( | |
| self, queue_position: Optional[int], agent: AgentRecord | |
| ) -> Optional[int]: | |
| if queue_position is None: | |
| return None | |
| if isinstance(queue_position, bool) or not isinstance(queue_position, int): | |
| raise ValidationError("queue position must be an integer") | |
| if not 1 <= queue_position <= self._capacity_locked(agent): | |
| raise ValidationError("queue position is outside the agent's capacity") | |
| return queue_position | |
| def reject_offer( | |
| self, | |
| agent_id: str, | |
| job_id: str, | |
| reason: str, | |
| *, | |
| now: Optional[float] = None, | |
| ) -> JobView: | |
| """Fail an offer that the targeted worker cannot safely admit. | |
| Submission already reserves queue capacity from the worker's last | |
| snapshot, so rejection indicates a changed or inconsistent local | |
| state. Failing explicitly is preferable to leaving the user with a | |
| permanently offered job and an invisible worker-side rejection. | |
| """ | |
| if not isinstance(reason, str) or not reason or len(reason) > 512: | |
| raise ValidationError("offer rejection reason must be 1-512 characters") | |
| with self._lock: | |
| current = self._now(now) | |
| self._expire_jobs_locked(current) | |
| agent = self._agent_locked(agent_id) | |
| if agent.revoked: | |
| raise AuthenticationError("agent credential is revoked") | |
| job = self._job_locked(job_id) | |
| if job.selection.target_agent_id != agent_id: | |
| raise ConflictError("offer is targeted to another agent") | |
| if job.status is JobStatus.FAILED and job.error == f"agent rejected offer: {reason}": | |
| return self._job_view_locked(job) | |
| if job.status is not JobStatus.OFFERED: | |
| raise ConflictError(f"job cannot be rejected while {job.status.value}") | |
| job.status = JobStatus.FAILED | |
| job.error = f"agent rejected offer: {reason}" | |
| job.queue_position = None | |
| job.updated_at = current | |
| self._fail_children_locked(job_id, current, "dependency failed") | |
| self._scrub_terminal_locked(job, drop_result=True) | |
| return self._job_view_locked(job) | |
| def complete_job( | |
| self, | |
| agent_id: str, | |
| job_id: str, | |
| result: Any, | |
| *, | |
| now: Optional[float] = None, | |
| ) -> JobView: | |
| return self._finish_job(agent_id, job_id, result=result, error=None, now=now) | |
| def fail_job( | |
| self, | |
| agent_id: str, | |
| job_id: str, | |
| error: str, | |
| *, | |
| result: Any = None, | |
| now: Optional[float] = None, | |
| ) -> JobView: | |
| if not isinstance(error, str) or not error or len(error) > 4_096: | |
| raise ValidationError("error must be 1-4096 characters") | |
| return self._finish_job(agent_id, job_id, result=result, error=error, now=now) | |
| def cancel_job( | |
| self, session_id: str, job_id: str, *, now: Optional[float] = None | |
| ) -> JobView: | |
| with self._lock: | |
| current = self._now(now) | |
| job = self._job_in_scope_locked(session_id, job_id) | |
| if job.status is JobStatus.CANCELLED: | |
| return self._job_view_locked(job) | |
| if job.status in _TERMINAL_STATUSES: | |
| raise ConflictError(f"job cannot be cancelled while {job.status.value}") | |
| self._cancel_tree_locked(job_id, current, "cancelled by user") | |
| return self._job_view_locked(job) | |
| def cancellations_for_agent(self, agent_id: str) -> Tuple[JobView, ...]: | |
| with self._lock: | |
| self._agent_locked(agent_id) | |
| jobs = [ | |
| job | |
| for job in self._jobs.values() | |
| if job.accepted_by == agent_id and job.status is JobStatus.CANCELLED | |
| ] | |
| jobs.sort(key=lambda job: (job.updated_at, job.spec.id)) | |
| return tuple(self._job_view_locked(job) for job in jobs) | |
| def cancellation_ids_for_agent(self, agent_id: str) -> tuple[str, ...]: | |
| """Return cancellations, including prompt-free session-deletion tombstones.""" | |
| with self._lock: | |
| self._agent_locked(agent_id) | |
| values = { | |
| job.spec.id | |
| for job in self._jobs.values() | |
| if job.accepted_by == agent_id and job.status is JobStatus.CANCELLED | |
| } | |
| values.update(self._orphan_cancellations.get(agent_id, ())) | |
| return tuple(sorted(values)) | |
| def claim_result( | |
| self, session_id: str, job_id: str, *, now: Optional[float] = None | |
| ) -> Any: | |
| """Hand the answer to its owning session, then schedule its erasure. | |
| The browser keeps the transcript; the server does not. Calling this | |
| marks the answer delivered, and the next ``cleanup_expired`` drops it. | |
| Repeat calls before that pass still return the answer, so a dropped | |
| round trip does not lose the user's reply -- but the window is one | |
| cleanup pass, not the job TTL. | |
| """ | |
| with self._lock: | |
| current = self._now(now) | |
| job = self._job_in_scope_locked(session_id, job_id) | |
| if job.status is not JobStatus.COMPLETED: | |
| raise ConflictError(f"job has no deliverable answer while {job.status.value}") | |
| if job.result is None: | |
| return None | |
| if job.delivered_at is None: | |
| job.delivered_at = current | |
| return self._copy_result(job.result) | |
| def get_job(self, session_id: str, job_id: str) -> JobView: | |
| with self._lock: | |
| return self._job_view_locked(self._job_in_scope_locked(session_id, job_id)) | |
| def list_jobs( | |
| self, session_id: str, conversation_id: Optional[str] = None | |
| ) -> Tuple[JobView, ...]: | |
| with self._lock: | |
| session = self._session_locked(session_id) | |
| if conversation_id is not None: | |
| records = [self._conversation_locked(session_id, conversation_id)] | |
| else: | |
| records = [ | |
| self._conversation_locked(session_id, item) | |
| for item in session.conversation_ids | |
| ] | |
| return tuple( | |
| self._job_view_locked(self._jobs[job_id]) | |
| for conversation in records | |
| for job_id in conversation.job_ids | |
| if job_id in self._jobs | |
| ) | |
| def reassign_job( | |
| self, | |
| session_id: str, | |
| job_id: str, | |
| target_agent_id: Optional[str] = None, | |
| *, | |
| viewer_id: Optional[str] = None, | |
| now: Optional[float] = None, | |
| ) -> JobView: | |
| """Move a stranded run to another agent, or fail it visibly. | |
| Rebuilt from the immutable :class:`RunSelection`, so the exact model | |
| and the exact tool versions the user consented to cannot be dropped or | |
| widened on the way to the new worker. The previous agent, if it ever | |
| returns, finds a cancellation tombstone rather than a job two agents | |
| both believe they own. | |
| """ | |
| with self._lock: | |
| current = self._now(now) | |
| job = self._job_in_scope_locked(session_id, job_id) | |
| if job.status is not JobStatus.STRANDED: | |
| raise ConflictError( | |
| f"only a stranded run can be reassigned; this one is {job.status.value}" | |
| ) | |
| previous_agent = job.selection.target_agent_id | |
| tools = job.spec.allowed_tools | |
| if target_agent_id is None: | |
| target_agent_id = self._choose_agent_locked( | |
| job.selection.model_id, | |
| tools, | |
| current, | |
| viewer_id=viewer_id, | |
| avoid=(previous_agent,), | |
| ) | |
| else: | |
| target = self._agent_locked(target_agent_id) | |
| if not self._visible_locked(target, viewer_id): | |
| raise NotFoundError(_AGENT_NOT_FOUND) | |
| if not self._is_compatible_locked(target, job.selection.model_id, tools): | |
| raise NoCompatibleAgentError( | |
| "target agent cannot run the exact model and tool versions" | |
| ) | |
| if not self._eligible_for_offer_locked(target, current): | |
| raise NoCompatibleAgentError("target agent is not currently eligible") | |
| # The old agent must never run this job if it reappears. | |
| if previous_agent != UNASSIGNED_AGENT: | |
| self._orphan_cancellations.setdefault(previous_agent, {})[job_id] = ( | |
| current + self.limits.job_ttl_seconds | |
| ) | |
| self._assign_stranded_locked(job, target_agent_id, current) | |
| return self._job_view_locked(job) | |
| def _assign_stranded_locked(self, job: JobRecord, target_agent_id: str, now: float) -> None: | |
| """Point a stranded job at ``target_agent_id`` and re-offer it. | |
| The pseudonymous ``user_key`` is re-derived for the new agent: it is | |
| keyed per agent by design, so the one minted for the previous agent | |
| would be meaningless (and would quietly bypass the new agent's deny | |
| list) if carried across. | |
| """ | |
| limits = {key: value for key, value in job.spec.limits.items() if key != "user_key"} | |
| target_record = self._agents.get(target_agent_id) | |
| if job.viewer_subject and target_record is not None and target_record.secret: | |
| limits["user_key"] = derive_user_key(target_record.secret, job.viewer_subject) | |
| job.spec = replace(job.spec, target_agent_id=target_agent_id, limits=limits) | |
| job.selection = RunSelection( | |
| model_id=job.selection.model_id, | |
| tool_versions=job.selection.tool_versions, | |
| target_agent_id=target_agent_id, | |
| ) | |
| job.accepted_by = None | |
| job.queue_position = None | |
| job.status = JobStatus.OFFERED | |
| job.updated_at = now | |
| def _strand_and_recover_locked(self, now: float) -> int: | |
| """Flip live jobs on silent agents to STRANDED, and back on return. | |
| Detection alone was already present (missed heartbeats); this is what | |
| follows detection. STRANDED is non-terminal: nothing is erased, and a | |
| job returns to its previous live state when its agent's snapshot next | |
| reports it (record_snapshot) or, for a never-accepted offer, here. | |
| """ | |
| stranded = 0 | |
| for job in self._jobs.values(): | |
| agent = self._agents.get(job.selection.target_agent_id) | |
| agent_online = agent is not None and self._is_online_locked(agent, now) | |
| if job.status in _STRANDABLE_STATUSES and not agent_online: | |
| job.status = JobStatus.STRANDED | |
| job.queue_position = None | |
| job.updated_at = now | |
| stranded += 1 | |
| elif ( | |
| job.status is JobStatus.STRANDED | |
| and agent_online | |
| and job.accepted_by is None | |
| ): | |
| # An offer the agent never saw: re-offer it. Accepted work is | |
| # recovered by record_snapshot instead, which knows whether | |
| # the agent still holds it. | |
| job.status = JobStatus.OFFERED | |
| job.updated_at = now | |
| elif job.status is JobStatus.STRANDED and not agent_online: | |
| # A run waiting for an agent that does not exist or has not | |
| # returned: assign it automatically the moment a compatible | |
| # agent is available. This covers both a run submitted before | |
| # any worker joined and a run whose worker never came back. | |
| tools = job.spec.allowed_tools | |
| try: | |
| chosen = self._choose_agent_locked( | |
| job.selection.model_id, | |
| tools, | |
| now, | |
| viewer_id=job.viewer_subject, | |
| avoid=( | |
| (job.selection.target_agent_id,) | |
| if job.selection.target_agent_id != UNASSIGNED_AGENT | |
| and job.accepted_by is None | |
| else () | |
| ), | |
| ) | |
| except (CapacityError, NoCompatibleAgentError, NotFoundError): | |
| continue | |
| if job.accepted_by is not None: | |
| # Accepted work is only moved by an explicit reassign; the | |
| # user chooses between waiting and repeating side effects. | |
| continue | |
| previous = job.selection.target_agent_id | |
| if previous != UNASSIGNED_AGENT: | |
| self._orphan_cancellations.setdefault(previous, {})[job.spec.id] = ( | |
| now + self.limits.job_ttl_seconds | |
| ) | |
| self._assign_stranded_locked(job, chosen, now) | |
| return stranded | |
| # Expiry ------------------------------------------------------------- | |
| def cleanup_expired(self, *, now: Optional[float] = None) -> CleanupReport: | |
| with self._lock: | |
| current = self._now(now) | |
| pairing_count = self._cleanup_pairing_codes_locked(current) | |
| nonce_count = sum( | |
| self._cleanup_agent_nonces_locked(record, current) | |
| for record in self._agents.values() | |
| ) | |
| self._strand_and_recover_locked(current) | |
| expired_count = self._expire_jobs_locked(current) | |
| self._cleanup_orphan_cancellations_locked(current) | |
| erased_count = self._scrub_delivered_locked(current) | |
| offline_count = sum( | |
| not self._is_online_locked(record, current) | |
| for record in self._agents.values() | |
| ) | |
| return CleanupReport( | |
| pairing_codes_removed=pairing_count, | |
| nonces_removed=nonce_count, | |
| jobs_expired=expired_count, | |
| agents_offline=offline_count, | |
| results_erased=erased_count, | |
| ) | |
| # Locked helpers ----------------------------------------------------- | |
| def _now(self, supplied: Optional[float]) -> float: | |
| value = self._clock() if supplied is None else supplied | |
| if isinstance(value, bool) or not isinstance(value, (int, float)): | |
| raise ValidationError("time must be numeric") | |
| value = float(value) | |
| if not math.isfinite(value) or value < 0: | |
| raise ValidationError("time must be finite and non-negative") | |
| return value | |
| def _agent_locked(self, agent_id: str) -> AgentRecord: | |
| record = self._agents.get(agent_id) | |
| if record is None: | |
| raise NotFoundError(_AGENT_NOT_FOUND) | |
| return record | |
| def _session_locked(self, session_id: str) -> SessionRecord: | |
| record = self._sessions.get(session_id) | |
| if record is None: | |
| raise NotFoundError("session was not found") | |
| return record | |
| def _conversation_locked(self, session_id: str, conversation_id: str) -> ConversationRecord: | |
| record = self._conversations.get((session_id, conversation_id)) | |
| if record is None: | |
| raise NotFoundError("conversation was not found in this session") | |
| return record | |
| def _job_locked(self, job_id: str) -> JobRecord: | |
| record = self._jobs.get(job_id) | |
| if record is None: | |
| raise NotFoundError("job was not found") | |
| return record | |
| def _job_in_scope_locked(self, session_id: str, job_id: str) -> JobRecord: | |
| job = self._job_locked(job_id) | |
| if job.spec.session_id != session_id: | |
| # Do not reveal cross-session record existence at the API boundary. | |
| raise NotFoundError("job was not found") | |
| return job | |
| def _cleanup_pairing_codes_locked(self, now: float) -> int: | |
| expired = [ | |
| digest for digest, record in self._pairing_codes.items() if record.expires_at <= now | |
| ] | |
| for digest in expired: | |
| del self._pairing_codes[digest] | |
| return len(expired) | |
| def _cleanup_unclaimed_locked(self, now: float) -> int: | |
| """Forget workers nobody claimed in time. | |
| An unclaimed record is the only thing on this server an unauthenticated | |
| caller can create, so it must not be something an unauthenticated caller | |
| can accumulate. It costs a slot for at most ``claim_ttl_seconds`` and is | |
| then gone; the worker holding it discovers this on its next poll and | |
| says so rather than waiting for ever. | |
| """ | |
| expired = [ | |
| agent_id | |
| for agent_id, record in self._agents.items() | |
| if record.claim_digest and record.claim_expires_at <= now | |
| ] | |
| for agent_id in expired: | |
| del self._agents[agent_id] | |
| return len(expired) | |
| def _cleanup_agent_nonces_locked(record: AgentRecord, now: float) -> int: | |
| expired = [nonce for nonce, expires_at in record.nonces.items() if expires_at <= now] | |
| for nonce in expired: | |
| del record.nonces[nonce] | |
| return len(expired) | |
| def _cleanup_orphan_cancellations_locked(self, now: float) -> None: | |
| for agent_id, tombstones in tuple(self._orphan_cancellations.items()): | |
| for job_id, expires_at in tuple(tombstones.items()): | |
| if expires_at <= now: | |
| del tombstones[job_id] | |
| if not tombstones: | |
| del self._orphan_cancellations[agent_id] | |
| def _is_online_locked(self, record: AgentRecord, now: float) -> bool: | |
| if record.revoked: | |
| return False | |
| elapsed = max(0.0, now - record.last_snapshot_at) | |
| if elapsed >= self.limits.offline_after_seconds: | |
| return False | |
| return record.snapshot is None or record.snapshot.status is not AgentStatus.OFFLINE | |
| def _missed_polls_locked(self, record: AgentRecord, now: float) -> int: | |
| elapsed = max(0.0, now - record.last_snapshot_at) | |
| return int(elapsed // self.limits.poll_interval_seconds) | |
| def _eligible_for_offer_locked(self, record: AgentRecord, now: float) -> bool: | |
| if not self._is_online_locked(record, now): | |
| return False | |
| if record.snapshot is not None and record.snapshot.status not in _ELIGIBLE_AGENT_STATUSES: | |
| return False | |
| return True | |
| def _normalize_tools(tools: Sequence[ToolSelection]) -> Tuple[ToolSelection, ...]: | |
| try: | |
| normalized = tuple( | |
| tool if isinstance(tool, ToolSelection) else ToolSelection.from_dict(tool) | |
| for tool in tools | |
| ) | |
| except (TypeError, ValueError) as exc: | |
| raise ValidationError(str(exc)) from exc | |
| if len(normalized) > 16 or len({tool.id for tool in normalized}) != len(normalized): | |
| raise ValidationError("tools must be unique and limited to 16") | |
| return normalized | |
| def _is_compatible_locked( | |
| record: AgentRecord, | |
| model_id: str, | |
| tools: Sequence[ToolSelection], | |
| ) -> bool: | |
| if record.revoked or model_id not in record.capabilities.models: | |
| return False | |
| advertised = { | |
| item if "@" in item else f"{item}@1.0.0" | |
| for item in record.capabilities.tools | |
| } | |
| required = {_tool_capability(tool) for tool in tools} | |
| return required.issubset(advertised) | |
| def _capacity_locked(self, record: AgentRecord) -> int: | |
| if record.snapshot is None: | |
| return record.capabilities.queue_capacity | |
| return min(record.capabilities.queue_capacity, record.snapshot.queue_capacity) | |
| def _outstanding_count_locked(self, agent_id: str) -> int: | |
| return sum( | |
| job.selection.target_agent_id == agent_id and job.status not in _TERMINAL_STATUSES | |
| for job in self._jobs.values() | |
| ) | |
| def _agent_load_key_locked(self, agent_id: str, record: AgentRecord) -> tuple[float, int, float]: | |
| outstanding = self._outstanding_count_locked(agent_id) | |
| capacity = self._capacity_locked(record) | |
| wait = record.snapshot.estimated_wait_s if record.snapshot is not None else 0.0 | |
| return (outstanding / capacity, outstanding, wait) | |
| def _choose_agent_locked( | |
| self, | |
| model_id: str, | |
| tools: Sequence[ToolSelection], | |
| now: float, | |
| *, | |
| viewer_id: Optional[str] = None, | |
| avoid: Sequence[str] = (), | |
| ) -> str: | |
| """Pick the least loaded agent this viewer is allowed to be given. | |
| ``avoid`` carries agents that have already refused this user. Automatic | |
| selection skips them rather than offering work that will be refused | |
| again, which is what stops a denial turning into a loop. It is a routing | |
| hint only: the server learns nothing it was not already told by the | |
| refusal, it never learns the agent's deny list, and losing the hint on a | |
| restart costs one wasted offer rather than correctness. | |
| """ | |
| skip = frozenset(avoid) | |
| candidates = [] | |
| compatible_seen = False | |
| for agent_id, record in self._agents.items(): | |
| if not self._visible_locked(record, viewer_id): | |
| continue | |
| if not self._is_compatible_locked(record, model_id, tools): | |
| continue | |
| compatible_seen = True | |
| if agent_id in skip: | |
| continue | |
| if not self._eligible_for_offer_locked(record, now): | |
| continue | |
| if self._outstanding_count_locked(agent_id) >= self._capacity_locked(record): | |
| continue | |
| candidates.append((self._agent_load_key_locked(agent_id, record), agent_id)) | |
| if not candidates: | |
| if compatible_seen: | |
| raise CapacityError("compatible agents are offline or at capacity") | |
| raise NoCompatibleAgentError("no agent supports the exact model and tool versions") | |
| candidates.sort(key=lambda value: (value[0], value[1])) | |
| return candidates[0][1] | |
| def _agent_view_locked(self, agent_id: str, record: AgentRecord, now: float) -> AgentView: | |
| return AgentView( | |
| capabilities=copy.deepcopy(record.capabilities), | |
| owner_id=record.owner_id, | |
| snapshot=( | |
| AgentSnapshot.from_dict(record.snapshot.to_dict()) | |
| if record.snapshot is not None | |
| else None | |
| ), | |
| online=self._is_online_locked(record, now), | |
| missed_polls=self._missed_polls_locked(record, now), | |
| last_snapshot_at=record.last_snapshot_at, | |
| outstanding_jobs=self._outstanding_count_locked(agent_id), | |
| revoked=record.revoked, | |
| ) | |
| def _session_view_locked(record: SessionRecord) -> SessionView: | |
| return SessionView(record.session_id, record.created_at, tuple(record.conversation_ids)) | |
| def _conversation_view_locked(record: ConversationRecord) -> ConversationView: | |
| return ConversationView( | |
| record.session_id, | |
| record.conversation_id, | |
| record.created_at, | |
| tuple(record.job_ids), | |
| ) | |
| def _job_view_locked(record: JobRecord) -> JobView: | |
| # Wire objects are frozen, but nested values may still be mappings. | |
| # Rebuild through the wire representation so callers never share state. | |
| spec = JobSpec.from_dict(record.spec.to_dict()) | |
| return JobView( | |
| spec=spec, | |
| status=record.status, | |
| accepted_by=record.accepted_by, | |
| queue_position=record.queue_position, | |
| updated_at=record.updated_at, | |
| expires_at=record.expires_at, | |
| result=ControlPlane._copy_result(record.result), | |
| error=record.error, | |
| prompt_erased=record.prompt_erased, | |
| result_erased=record.result_erased, | |
| ) | |
| def _bounded_result_copy(self, result: Any) -> Any: | |
| try: | |
| if hasattr(result, "to_dict"): | |
| encoded = json.dumps(result.to_dict(), ensure_ascii=False, default=str) | |
| else: | |
| encoded = json.dumps(result, ensure_ascii=False, default=str) | |
| except (TypeError, ValueError) as exc: | |
| raise ValidationError("result must be serializable") from exc | |
| if len(encoded) > self.limits.max_result_characters: | |
| raise CapacityError("job result exceeds the configured size limit") | |
| return self._copy_result(result) | |
| def _copy_result(result: Any) -> Any: | |
| if result is None: | |
| return None | |
| result_type = type(result) | |
| if hasattr(result, "to_dict") and hasattr(result_type, "from_dict"): | |
| return result_type.from_dict(result.to_dict()) | |
| return copy.deepcopy(result) | |
| def _finish_job( | |
| self, | |
| agent_id: str, | |
| job_id: str, | |
| *, | |
| result: Any, | |
| error: Optional[str], | |
| now: Optional[float], | |
| ) -> JobView: | |
| result_copy = self._bounded_result_copy(result) | |
| with self._lock: | |
| current = self._now(now) | |
| self._expire_jobs_locked(current) | |
| self._agent_locked(agent_id) | |
| job = self._job_locked(job_id) | |
| if job.selection.target_agent_id != agent_id or job.accepted_by != agent_id: | |
| raise ConflictError("agent does not own this accepted job") | |
| desired = JobStatus.FAILED if error is not None else JobStatus.COMPLETED | |
| if job.status is desired: | |
| return self._job_view_locked(job) # idempotent retry | |
| if job.status not in _AGENT_EXECUTION_STATUSES: | |
| raise ConflictError(f"job cannot finish while {job.status.value}") | |
| job.status = desired | |
| job.result = result_copy | |
| job.error = error | |
| job.queue_position = None | |
| job.updated_at = current | |
| if desired is JobStatus.COMPLETED: | |
| # Children must be released before the parent is scrubbed: the | |
| # released context is built from the parent's still-intact spec. | |
| self._release_children_locked(job_id, current) | |
| self._scrub_terminal_locked(job, drop_result=False) | |
| else: | |
| self._fail_children_locked(job_id, current, "dependency failed") | |
| self._scrub_terminal_locked(job, drop_result=True) | |
| return self._job_view_locked(job) | |
| def _release_children_locked(self, parent_job_id: str, now: float) -> None: | |
| parent = self._jobs[parent_job_id] | |
| context: Optional[tuple[ChatMessage, ...]] = None | |
| for child_id in sorted(self._children.get(parent_job_id, ())): | |
| child = self._jobs[child_id] | |
| if child.status is JobStatus.BLOCKED: | |
| # A user may queue this turn before its parent has answered. | |
| # Build the context now, before the spec can be offered, and | |
| # build it from the parent chain rather than by re-reading every | |
| # earlier job: the parent's own messages already carry the | |
| # history, so nothing here depends on prompts that erasure has | |
| # since removed. This is also a single provenance -- the server | |
| # derives the context, the client never supplies it for a | |
| # released turn. ``release_with_context`` is the one sanctioned | |
| # spec mutation and raises on any second attempt or any | |
| # non-blocked status, so the immutability of an offered spec is | |
| # held by the record itself rather than by ordering here. | |
| if context is None: | |
| context = self._context_from_parent_locked(parent) | |
| child.release_with_context(context, now) | |
| def _context_from_parent_locked(self, parent: JobRecord) -> tuple[ChatMessage, ...]: | |
| """Context for a follow-up: the parent's own window, plus its turn. | |
| Called from ``_release_children_locked`` while the parent record is | |
| still intact, i.e. before ``_scrub_terminal_locked`` runs. The result | |
| is transitively complete because each parent's ``messages`` were built | |
| the same way when *it* was released. | |
| """ | |
| context: list[ChatMessage] = list(parent.spec.messages) | |
| if not parent.prompt_erased: | |
| context.append(ChatMessage("user", parent.spec.prompt)) | |
| output = self._completed_output_locked(parent) | |
| if output is not None: | |
| context.append(ChatMessage("assistant", output)) | |
| return _bounded_context(context) | |
| def _completed_output_locked(job: JobRecord) -> str | None: | |
| if job.status is not JobStatus.COMPLETED: | |
| return None | |
| result = job.result | |
| if isinstance(result, Mapping): | |
| result = result.get("output", result.get("text")) | |
| if not isinstance(result, str) or not result: | |
| return None | |
| return result[:65_536] | |
| def _scrub_terminal_locked(job: JobRecord, *, drop_result: bool) -> None: | |
| """Erase prompt text -- and optionally the answer -- from a finished job. | |
| Called on every terminal transition. The prompt goes immediately: once | |
| a run is over the server has no further use for it, and holding it is | |
| the thing the privacy notice says does not happen. | |
| A *completed* run keeps its answer until the owning browser session has | |
| read it once (``claim_result``), because otherwise the user would never | |
| see the reply they waited for. Every other terminal state -- failed, | |
| cancelled, expired -- drops both at once; there is nothing to deliver. | |
| """ | |
| if not job.prompt_erased: | |
| job.spec = replace(job.spec, prompt=ERASED_PROMPT, messages=()) | |
| job.prompt_erased = True | |
| if drop_result and not job.result_erased: | |
| job.result = None | |
| job.result_erased = True | |
| def _fail_children_locked(self, parent_job_id: str, now: float, reason: str) -> None: | |
| for child_id in sorted(self._children.get(parent_job_id, ())): | |
| child = self._jobs[child_id] | |
| if child.status not in _TERMINAL_STATUSES: | |
| child.status = JobStatus.FAILED | |
| child.error = f"{reason}: {parent_job_id}" | |
| child.queue_position = None | |
| child.updated_at = now | |
| self._scrub_terminal_locked(child, drop_result=True) | |
| self._fail_children_locked(child_id, now, reason) | |
| def _cancel_tree_locked(self, job_id: str, now: float, reason: str) -> None: | |
| job = self._jobs[job_id] | |
| if job.status not in _TERMINAL_STATUSES: | |
| job.status = JobStatus.CANCELLED | |
| job.error = reason | |
| job.queue_position = None | |
| job.updated_at = now | |
| # Cancellation is an explicit request to stop; there is no answer | |
| # owed to anyone, so prompt and any partial result go together. | |
| self._scrub_terminal_locked(job, drop_result=True) | |
| for child_id in sorted(self._children.get(job_id, ())): | |
| self._cancel_tree_locked(child_id, now, "dependency was cancelled") | |
| def _expire_jobs_locked(self, now: float) -> int: | |
| expired = 0 | |
| for job_id, job in tuple(self._jobs.items()): | |
| if job.status not in _TERMINAL_STATUSES and job.expires_at <= now: | |
| job.status = JobStatus.EXPIRED | |
| job.error = "job expired before completion" | |
| job.queue_position = None | |
| job.updated_at = now | |
| self._scrub_terminal_locked(job, drop_result=True) | |
| expired += 1 | |
| expired += self._expire_children_locked(job_id, now) | |
| return expired | |
| def _expire_children_locked(self, parent_job_id: str, now: float) -> int: | |
| expired = 0 | |
| for child_id in sorted(self._children.get(parent_job_id, ())): | |
| child = self._jobs[child_id] | |
| if child.status not in _TERMINAL_STATUSES: | |
| child.status = JobStatus.EXPIRED | |
| child.error = f"dependency expired: {parent_job_id}" | |
| child.queue_position = None | |
| child.updated_at = now | |
| self._scrub_terminal_locked(child, drop_result=True) | |
| expired += 1 | |
| expired += self._expire_children_locked(child_id, now) | |
| return expired | |
| def _scrub_delivered_locked(self, now: float) -> int: | |
| """Drop answers the owning session has already read. | |
| Retention for a completed run's answer is therefore: until the browser | |
| fetches it, plus one cleanup pass. It is not indefinite, and it is not | |
| the previous 24-hour ``job_ttl_seconds`` window. | |
| """ | |
| erased = 0 | |
| for job in self._jobs.values(): | |
| if ( | |
| job.delivered_at is not None | |
| and not job.result_erased | |
| and job.delivered_at < now | |
| ): | |
| job.result = None | |
| job.result_erased = True | |
| erased += 1 | |
| return erased | |