"""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)) @staticmethod 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) @staticmethod 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 @staticmethod 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 @staticmethod 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, ) @staticmethod def _session_view_locked(record: SessionRecord) -> SessionView: return SessionView(record.session_id, record.created_at, tuple(record.conversation_ids)) @staticmethod def _conversation_view_locked(record: ConversationRecord) -> ConversationView: return ConversationView( record.session_id, record.conversation_id, record.created_at, tuple(record.job_ids), ) @staticmethod 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) @staticmethod 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) @staticmethod 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] @staticmethod 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