"""Injectable worker poll loop connecting queue, runner, and telemetry.""" from __future__ import annotations import re as _re import threading import time from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass, replace from typing import Any, Protocol from distinct_protocol import ( MAX_RUN_STEPS, PHASE_FETCHING_WEIGHTS, PHASE_GUARD_INPUT, PHASE_GUARD_OUTPUT, STEP_PHASE, AgentCapabilities, AgentSnapshot, JobResult, JobSpec, JobStatus, ) from .energy import EnergyMeter, UnavailableEnergyMeter from .guard import ( DECISION_BLOCK, GuardOverhead, GuardScreen, SelfCheckGuard, guard_is_meaningful, user_facing_refusal, ) from .harness import InferenceHarness, ProgressReporter, StructuredToolHarness, emit from .isolation import RequestSandbox, describe_request_isolation from .models import DiscoveredModel, index_discovered from .queue import DEFAULT_POLL_INTERVAL_SECONDS, InMemoryJobQueue, OfferDecision from .runners import InferenceRunner, RunnerCancelled from .sandbox import GeneratedExecutionPolicy from .servers import REASON_DENIED_BY_AGENT, AllowanceGate, Refusal from .tools import load_tool_broker def _parameters_b(model: DiscoveredModel) -> float | None: """A rough parameter count read from the model id, for the guard's floor. "qwen3-0.6b" -> 0.6, "olmoe-1b-7b-instruct" -> 7.0 (the last figure is the total). ``None`` when nothing in the id looks like a size, which the guard treats as unknown rather than as passing the floor. """ matches = _re.findall(r"(\d+(?:\.\d+)?)b", model.manifest.id.casefold()) if not matches: return None try: return float(matches[-1]) except ValueError: return None class AgentTransport(Protocol): """Network-independent interface implemented by server-specific clients. Implementations must be safe to call from worker threads for results and from the poll thread for offers/snapshots. """ def poll_offer(self, snapshot: AgentSnapshot) -> JobSpec | None: ... def poll_cancellations(self) -> Iterable[str]: ... def acknowledge_offer(self, decision: OfferDecision) -> None: ... def publish_status( self, job_id: str, status: JobStatus, reason: str | None = None ) -> None: ... def publish_result(self, result: JobResult) -> None: ... class JobExecutor(Protocol): def execute( self, job: JobSpec, *, cancel_event: threading.Event, progress: ProgressReporter, ) -> JobResult: ... class LocalInferenceExecutor: """Execute one local inference and attach measured energy when available.""" def __init__( self, *, models: Iterable[DiscoveredModel], runner: InferenceRunner, energy_meter: EnergyMeter | None = None, generated_execution: GeneratedExecutionPolicy | None = None, tool_harness: InferenceHarness | None = None, guard_enabled: bool = False, guard_model_id: str | None = None, deny_list=None, weights: Any = None, manifests: Iterable[Any] = (), request_isolation: bool = True, sandbox_runtime_paths: Iterable[Any] = (), ) -> None: self.models = index_discovered(models) # What this worker has *agreed to run*, which is no longer the same as # what is on disk. Weights arrive on demand, so a model can be offered # before its file exists; ``weights`` is the cache that closes that gap # at the moment a run first needs it. self.manifests = {manifest.id: manifest for manifest in manifests} or { model.manifest.id: model.manifest for model in self.models.values() } self.weights = weights self.runner = runner self.energy_meter = energy_meter or UnavailableEnergyMeter() self.generated_execution = generated_execution or GeneratedExecutionPolicy() self.tool_harness = tool_harness or StructuredToolHarness() # Guard screening: one extra inference pass over the input and one # over the output, Llama Guard style. The OPERATOR chooses the guard # model at setup (--guard-model); by default the model already loaded # for the run screens itself. Enforcement (blocking and deny-listing) # only happens at or above the guard's parameter floor; below it # verdicts are report-only. self.guard_enabled = bool(guard_enabled) self.guard_model_id = guard_model_id self.deny_list = deny_list # One sandbox per request, created here and destroyed when the request # ends. The alternative, one posture for the worker's lifetime, hands # each user whatever the previous user left behind. self.request_isolation = bool(request_isolation) self.sandbox_runtime_paths = tuple(sandbox_runtime_paths) def _screen( self, guard: GuardScreen, stage: str, text: str, ) -> tuple[Any, float, float | None]: """Run one guard pass, timing it and metering its energy separately.""" token = self.energy_meter.start() verdict, seconds = guard.screen(stage, text) usage = self.energy_meter.stop(token) joules = None if usage is not None: value = usage.to_dict().get("joules") if isinstance(value, (int, float)) and not isinstance(value, bool): joules = float(value) return verdict, seconds, joules def _resolve_model( self, job: JobSpec, *, cancel_event: threading.Event, progress: Callable[..., None], ) -> DiscoveredModel: """The verified local weights for this run, fetching them if needed. The failure here is deliberately loud and names the model. A worker that cannot obtain the exact pinned file fails the run; it never quietly runs a different one, and it never treats an absent digest as a passing check. """ model = self.models.get(job.model_id) if model is not None: return model manifest = self.manifests.get(job.model_id) if manifest is None: raise ValueError( f"model {job.model_id!r} is not among the models this worker approved" ) if self.weights is None: raise ValueError( f"model {job.model_id!r} is approved but its weights are not present " "and this worker has no weights cache configured" ) emit( progress, 0.01, f"Fetching {manifest.id} weights", phase=PHASE_FETCHING_WEIGHTS, step={ "kind": STEP_PHASE, "text": f"Downloading {manifest.id} from its pinned revision", }, ) def fetch_progress(fraction: float, message: str) -> None: # Weights fetching occupies the first tenth of the bar. It happens # once per model per worker, and a run that pays for it should not # look like a run that did not. emit(progress, min(0.1, fraction * 0.1), message, phase=PHASE_FETCHING_WEIGHTS) model = self.weights.ensure( manifest, progress=fetch_progress, cancel_event=cancel_event ) self.models[job.model_id] = model return model def execute( self, job: JobSpec, *, cancel_event: threading.Event, progress: Callable[..., None], ) -> JobResult: sandbox = None sandbox_report: dict[str, Any] | None = None if self.request_isolation: try: sandbox = RequestSandbox.create( request_id=_sandbox_id(job.id), read_only_paths=self.sandbox_runtime_paths, ) except OSError as exc: # A sandbox that cannot be built is reported, not pretended. # The run continues because refusing every job on a host with a # full temp directory helps nobody, and the result says plainly # that no confinement was applied. sandbox_report = { "applied": False, "reason": f"could not create a request sandbox: {exc}", } try: return self._execute_in( job, cancel_event=cancel_event, progress=progress, sandbox=sandbox, sandbox_report=sandbox_report, ) finally: if sandbox is not None: sandbox.close() # The scratchpad and the to-do list are held in this process, not # in the sandbox directory, so closing the sandbox does not reach # them. They are erased here, in the same block, because a note a # user's model wrote is prompt-derived text and this network's one # promise is that no such text outlives its run on a volunteer's # machine. Unconditional and in ``finally``: a run that raised is # exactly the run whose notes must not be left behind. _forget_workspace(job.id) def _execute_in( self, job: JobSpec, *, cancel_event: threading.Event, progress: Callable[..., None], sandbox: Any, sandbox_report: dict[str, Any] | None, ) -> JobResult: model = self._resolve_model(job, cancel_event=cancel_event, progress=progress) prompt = _render_prompt(job) guard_report: dict[str, Any] = {} guard = None guard_model = self.models.get(self.guard_model_id) or model parameters = _parameters_b(guard_model) enforce = guard_is_meaningful(parameters) if self.guard_enabled: classifier = SelfCheckGuard( lambda guard_prompt: self.runner.run( guard_model, guard_prompt, cancel_event=cancel_event, progress=None, limits={"max_output_tokens": 64}, ).text, model_ref=guard_model.manifest.id, parameters_b=parameters, ) guard = GuardScreen(classifier, model_ref=guard_model.manifest.id) emit( progress, 0.02, "Guard screening the input", phase=PHASE_GUARD_INPUT, step={"kind": STEP_PHASE, "text": "Safety screen on the request"}, ) verdict, in_seconds, in_joules = self._screen(guard, "input", job.prompt) guard_report = { "input": verdict.decision, "input_categories": list(verdict.categories), "note": classifier.note, "enforced": enforce, # Interim figures for the combined overhead computed after the # answer; stripped from the report before it leaves. "_input_seconds": in_seconds, "_input_joules": in_joules, } if verdict.decision == DECISION_BLOCK and enforce: user_key = job.limits.get("user_key") if self.deny_list is not None: self.deny_list.deny( user_key, reason_code=REASON_DENIED_BY_AGENT, evidence=verdict.evidence_id, ) overhead = GuardOverhead( input_seconds=in_seconds, input_joules=in_joules, model_ref=model.manifest.id, ) guard_report["overhead"] = overhead.to_dict() return JobResult( job_id=job.id, output=user_facing_refusal( REASON_DENIED_BY_AGENT, verdict.evidence_id ), usage={"guard": guard_report}, energy={}, ) token = self.energy_meter.start() try: if job.allowed_tools or getattr(self.tool_harness, "handles_plain_jobs", False): inference = self.tool_harness.run( job=job, model=model, runner=self.runner, broker=load_tool_broker(job), prompt=prompt, cancel_event=cancel_event, progress=progress, ) tool_events = inference.tool_events artifacts = getattr(inference, "artifacts", ()) else: inference = self.runner.run( model, prompt, cancel_event=cancel_event, progress=progress, limits=job.limits, ) tool_events = () artifacts = () finally: energy = self.energy_meter.stop(token) usage = dict(inference.usage) # What the request actually got, as opposed to what was configured. # This travels with the result so a run log can state the boundary it # ran behind rather than the boundary the worker hopes it has. if sandbox is not None: usage["request_sandbox"] = { "request_id": sandbox.request_id, "isolation": describe_request_isolation().to_dict(), "scratch": "per-request, erased with the run", } elif sandbox_report is not None: usage["request_sandbox"] = sandbox_report output_text = inference.text if guard is not None: emit( progress, 0.97, "Guard screening the output", phase=PHASE_GUARD_OUTPUT, step={"kind": STEP_PHASE, "text": "Safety screen on the answer"}, ) out_verdict, out_seconds, out_joules = self._screen( guard, "output", inference.text ) answer_joules = None value = energy.to_dict().get("joules") if isinstance(value, (int, float)) and not isinstance(value, bool): answer_joules = float(value) overhead = GuardOverhead( input_seconds=guard_report.get("_input_seconds"), output_seconds=out_seconds, input_joules=guard_report.pop("_input_joules", None), output_joules=out_joules, answer_joules=answer_joules, model_ref=model.manifest.id, ) guard_report.pop("_input_seconds", None) guard_report["output"] = out_verdict.decision guard_report["output_categories"] = list(out_verdict.categories) guard_report["overhead"] = overhead.to_dict() usage["guard"] = guard_report if out_verdict.decision == DECISION_BLOCK and enforce: user_key = job.limits.get("user_key") if self.deny_list is not None: self.deny_list.deny( user_key, reason_code=REASON_DENIED_BY_AGENT, evidence=out_verdict.evidence_id, ) output_text = user_facing_refusal( REASON_DENIED_BY_AGENT, out_verdict.evidence_id ) tool_events = () artifacts = () return JobResult( job_id=job.id, output=output_text, usage=usage, energy=energy.to_dict(), tool_events=tool_events, artifacts=artifacts, ) @dataclass(frozen=True) class WorkerLoopSettings: """One cadence, and it is the real one. There was a second, ``snapshot_interval_seconds`` (5 s), which gated a ``publish_snapshot`` call that the Gradio transport implemented as a no-op local cache. The snapshot the server actually received was the one ``poll_offer`` embeds on every cycle, so the advertised five-second heartbeat described nothing. Both the setting and the transport method are gone: the poll carries the snapshot, and ``poll_interval_seconds`` is the only rate in the system. """ poll_interval_seconds: float = DEFAULT_POLL_INTERVAL_SECONDS def __post_init__(self) -> None: if self.poll_interval_seconds <= 0: raise ValueError("poll interval must be positive") class WorkerLoop: """Conservative polling worker; one active job unless configured otherwise.""" def __init__( self, *, capabilities: AgentCapabilities, transport: AgentTransport, executor: JobExecutor, queue: InMemoryJobQueue | None = None, settings: WorkerLoopSettings | None = None, monotonic: Callable[[], float] = time.monotonic, gate: AllowanceGate | None = None, deny_list=None, ) -> None: self.capabilities = capabilities self.transport = transport self.executor = executor self.queue = queue or InMemoryJobQueue( capacity=capabilities.queue_capacity, max_active=capabilities.max_concurrency, agent_id=capabilities.agent_id, ) self.settings = settings or WorkerLoopSettings() self._monotonic = monotonic # The allowlist boundary. When present, no offer reaches the queue # without passing it, so a refused job never occupies capacity and # never becomes a JobSpec this worker has agreed to run. self.gate = gate # Agent-local refusals. Owned by this machine, sent nowhere, and asked # only one bit at a time. See distinct_agent.guard. self.deny_list = deny_list self.refusals: list[Refusal] = [] self.snapshots_sent = 0 # Lifetime energy, in the worker's own words. Reset when the process # restarts, because "since this worker started" is exactly the claim. self._started_monotonic = monotonic() self._lifetime_lock = threading.Lock() self._lifetime_joules: float | None = None self._lifetime_runs = 0 self._lifetime_measured_runs = 0 self._threads: dict[str, threading.Thread] = {} self._threads_lock = threading.Lock() def run_cycle(self) -> None: """Perform one non-blocking transport and dispatch cycle. Exactly one snapshot leaves per cycle, inside ``poll_offer``. There is no second snapshot path and no second timer. """ self._reap_threads() for job_id in tuple(self.transport.poll_cancellations()): decision = self.queue.cancel(job_id, "cancelled_by_server") if decision.found and decision.immediate: self.transport.publish_status(job_id, JobStatus.CANCELLED, decision.reason) snapshot = self._snapshot() offer = self.transport.poll_offer(snapshot) self.snapshots_sent += 1 if offer is not None: decision = self._admit(offer) self.transport.acknowledge_offer(decision) self._dispatch_ready() def _admit(self, offer: JobSpec) -> OfferDecision: """Decide whether this worker will run ``offer`` at all. The allowlist check happens **before** the queue sees the job, for two reasons. A refused job must not consume a queue slot, and admission is the last moment at which refusing is free: once a spec is queued the worker has, in every observable sense, agreed to run it. The rejection reason sent back is one of a closed set of ASCII tokens that never interpolates a model id or tool name. Worker-supplied text reaching a browser has to be escaped, and the surest way to satisfy that is for this path to produce no worker-supplied text at all. The detail an operator needs stays local, on :attr:`refusals`. """ # The deny list runs first. It is a decision about whether this machine # will serve this person at all, which is prior to what they asked for. if self.deny_list is not None: denied = self.deny_list.refusal_reason(offer.limits.get("user_key")) if denied is not None: self.refusals.append(Refusal(denied, f"job {offer.id} refused by the deny list")) return OfferDecision( job_id=offer.id, accepted=False, status=JobStatus.OFFERED, queue_position=None, reason=denied, ) if self.gate is not None: refusal = self.gate.check(offer) if refusal is not None: self.refusals.append(refusal) return OfferDecision( job_id=offer.id, accepted=False, status=JobStatus.OFFERED, queue_position=None, reason=refusal.reason, ) return self.queue.offer(offer) def apply_revocation(self, outcome) -> tuple[str, ...]: """Stop the work an operator's revocation just invalidated. Queued jobs are cancelled outright and running jobs are asked to stop through the same cooperative path a user cancel uses, so the runner unwinds cleanly rather than being killed part way through a write. Finishing "just this one" would be the wrong default: the operator revoked to stop this job, not the next one. """ stopped: list[str] = [] for job_id in getattr(outcome, "invalidated_job_ids", ()): decision = self.queue.cancel(job_id, "revoked_by_operator") if decision.found: stopped.append(job_id) if decision.immediate: self.transport.publish_status( job_id, JobStatus.CANCELLED, "revoked_by_operator" ) return tuple(stopped) def outstanding_jobs(self) -> tuple[JobSpec, ...]: """Every non-terminal job this worker holds, queued or running.""" return self.queue.queued_jobs() + self.queue.active_jobs() def run_forever(self, stop_event: threading.Event | None = None) -> None: stop = stop_event or threading.Event() while not stop.is_set(): self.run_cycle() stop.wait(self.settings.poll_interval_seconds) def shutdown(self, *, cancel_active: bool = True, join_timeout: float = 5.0) -> None: self.queue.close(cancel_waiting=True) if cancel_active: for job in self.queue.active_jobs(): self.queue.cancel(job.id, "worker_shutdown") deadline = self._monotonic() + max(0.0, join_timeout) for thread in self._thread_values(): thread.join(max(0.0, deadline - self._monotonic())) def _snapshot(self) -> AgentSnapshot: meter = getattr(self.executor, "energy_meter", None) energy_available = bool(meter is not None and meter.available) with self._lifetime_lock: joules = self._lifetime_joules runs = self._lifetime_runs measured = self._lifetime_measured_runs return replace( self.queue.snapshot( agent_id=self.capabilities.agent_id, energy_available=energy_available, ), uptime_seconds=max(0.0, self._monotonic() - self._started_monotonic), lifetime_joules=joules, lifetime_runs=runs, lifetime_measured_runs=measured, energy_scope=getattr(meter, "scope", "") or "", ) def _record_lifetime(self, result: JobResult) -> None: """Accumulate this worker's own running total since it started. Self-reported and therefore unverifiable, which is why the interface labels it as such. What it must still be is *honest about its gaps*: a run whose energy was not measured increments the run count and not the joules, so the total can declare itself a floor over N of M runs rather than passing off a partial sum as a complete one. """ energy = result.energy if isinstance(result.energy, Mapping) else {} joules = energy.get("joules") measured = ( bool(energy.get("available")) and isinstance(joules, (int, float)) and not isinstance(joules, bool) ) with self._lifetime_lock: self._lifetime_runs += 1 if measured: self._lifetime_measured_runs += 1 self._lifetime_joules = (self._lifetime_joules or 0.0) + float(joules) def _dispatch_ready(self) -> None: while True: job = self.queue.claim_next() if job is None: return thread = threading.Thread( target=self._execute_job, args=(job,), name=f"distinct-job-{job.id}", daemon=True, ) with self._threads_lock: self._threads[job.id] = thread self.transport.publish_status(job.id, JobStatus.RUNNING) thread.start() def _execute_job(self, job: JobSpec) -> None: cancel_event = self.queue.cancellation_event(job.id) # THE RUN'S OWN RECORD OF WHAT IT DID. # # Every step already went out on the live stream, which is a window on # a run in flight and is gone the moment the run finishes. A completed # turn therefore showed the model's final text and nothing about how it # was reached: which tools were considered, what they were passed, what # they returned. Kept here and returned with the result, the same steps # survive the run that produced them. taken: list[Mapping[str, Any]] = [] def report( fraction: float, message: str = "", *, phase: str | None = None, step: Mapping[str, Any] | None = None, ) -> None: if step is not None and len(taken) < MAX_RUN_STEPS: taken.append(dict(step)) self.queue.update_progress( job.id, fraction, message, phase=phase, step=step ) try: result = self.executor.execute( job, cancel_event=cancel_event, progress=report, ) if taken and not result.steps: result = replace(result, steps=tuple(taken)) if cancel_event.is_set(): self.queue.finish_cancelled(job.id, "cancelled_during_inference") self.transport.publish_status(job.id, JobStatus.CANCELLED) else: self.queue.complete(job.id) self._record_lifetime(result) self.transport.publish_result(result) except RunnerCancelled as exc: self.queue.finish_cancelled(job.id, str(exc)) self.transport.publish_status(job.id, JobStatus.CANCELLED, str(exc)) except Exception as exc: reason = f"{type(exc).__name__}: {exc}"[:4096] self.queue.fail(job.id, reason) self.transport.publish_result( JobResult(job_id=job.id, output="", error=reason) ) def _reap_threads(self) -> None: with self._threads_lock: finished = [job_id for job_id, thread in self._threads.items() if not thread.is_alive()] for job_id in finished: self._threads.pop(job_id, None) def _thread_values(self) -> tuple[threading.Thread, ...]: with self._threads_lock: return tuple(self._threads.values()) def _forget_workspace(job_id: str) -> None: """Erase this run's scratchpad and to-do list. Imported here rather than at module scope so a worker built without the tool library still starts: the library is optional, and its absence must not be a startup failure. """ try: from distinct_tools.workspace import STATE except Exception: # pragma: no cover - the library is optional return STATE.forget(job_id) def _sandbox_id(job_id: str) -> str: """A filesystem-safe sandbox name derived from the job id. Derived rather than random so a stray directory can be traced back to the run that made it, and sanitised rather than trusted because a job id crosses the wire. """ safe = "".join(character for character in job_id if character.isalnum() or character in "-_") return (safe or "job")[:64] from distinct_protocol.fences import CONVERSATION_FENCE_BODIES from distinct_protocol.fences import defang as _defang from distinct_protocol.fences import is_fence_line as _is_fence def _is_fence_line(line: str) -> bool: return _is_fence(line, CONVERSATION_FENCE_BODIES) def defang_fences(text: str) -> str: """Stop conversation text from closing the fence that contains it. Every prompt this harness builds puts the conversation between ``--- BEGIN CONVERSATION ---`` and ``--- END CONVERSATION ---`` and tells the model the span between them is data. A user can attach a file, and an attached file can contain that closing line. The model would then read everything after it as the harness's own framing rather than as the user's material: the sentence "ignore the tools and reply OK" stops being quoted text and starts being an instruction. The rule itself lives in `distinct_protocol.fences`, which is where the server's identical rule lives too. It was implemented twice, and both copies split on ``"\\n"`` alone, so a marker delimited by a carriage return or by U+2028 went through untouched. One rule cannot disagree with itself. """ return _defang(text, CONVERSATION_FENCE_BODIES) def _render_prompt(job: JobSpec) -> str: parts = [f"{message.role}: {defang_fences(message.content)}" for message in job.messages] parts.append(f"user: {defang_fences(job.prompt)}") return "\n".join(parts)