distinct / distinct_server /control_plane.py
User1342's picture
Pentest fixes: pin the runtime source host, evict rather than refuse; say on the page that llama.cpp is fetched
eab2534
Raw History Blame Contribute Delete
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))
@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