fix(runtime): recognize client identity environment and refresh worker registrations #976
@@ -1193,6 +1193,16 @@ RECOGNIZED_GITEA_ENV_KEYS = frozenset({
|
||||
"GITEA_IRRECOVERABLE_HMAC_SECRET",
|
||||
"GITEA_FORCE_MCP_RUNTIME_CHECK",
|
||||
"GITEA_FORCE_CLIENT_MANAGED",
|
||||
# #975: the client-identity inputs the server actually consumes at startup
|
||||
# (CLIENT_NAME_ENV / CLIENT_INSTANCE_ENV / CLIENT_SESSION_ENV in
|
||||
# gitea_mcp_server). Production read them while this allowlist omitted them,
|
||||
# so the peer-env scan classified them as unsupported overrides and the
|
||||
# capability resolver refused every mutation fleet-wide. Named individually
|
||||
# on purpose: no prefix is added, so an unrecognised GITEA_* override is
|
||||
# still refused exactly as it was before.
|
||||
"GITEA_MCP_CLIENT",
|
||||
"GITEA_MCP_CLIENT_INSTANCE",
|
||||
"GITEA_MCP_CLIENT_SESSION",
|
||||
})
|
||||
|
||||
RECOGNIZED_GITEA_ENV_PREFIXES = (
|
||||
|
||||
+113
-2
@@ -1916,8 +1916,16 @@ def _verify_role_mutation_workspace(
|
||||
if runtime_reasons:
|
||||
raise RuntimeError("; ".join(runtime_reasons))
|
||||
except Exception as exc:
|
||||
if "stale-runtime:" in str(exc):
|
||||
raise RuntimeError(str(exc))
|
||||
# #975: ``_check_mcp_runtimes_diagnostics`` raises every reason it
|
||||
# produces through this one RuntimeError, but only ``stale-runtime:``
|
||||
# was re-raised here — an ``unsupported-env:`` reason was swallowed
|
||||
# while still failing the capability resolver, so this preflight and
|
||||
# the resolver disagreed about the identical diagnostic. Both
|
||||
# authoritative prefixes now propagate the same way. This can only ever
|
||||
# widen what is refused, never widen what is permitted.
|
||||
message = str(exc)
|
||||
if any(prefix in message for prefix in RUNTIME_DIAGNOSTIC_HARD_PREFIXES):
|
||||
raise RuntimeError(message)
|
||||
pass
|
||||
|
||||
role = _effective_workspace_role()
|
||||
@@ -15714,6 +15722,10 @@ _WORKER_REGISTRY = None
|
||||
_WORKER_IDENTITY: str | None = None
|
||||
_WORKER_GENERATION: str | None = None
|
||||
_WORKER_REGISTRATION_ATTEMPTED = False
|
||||
#: #975: the one heartbeat supervisor for this process's registration. One
|
||||
#: process registers exactly one worker identity, so there is exactly one
|
||||
#: supervisor and it is never replaced.
|
||||
_WORKER_HEARTBEAT_SUPERVISOR = None
|
||||
|
||||
#: Env a client launcher may set to name itself and its session. Absent values
|
||||
#: are reported as unknown; they are never guessed at, because guessing is what
|
||||
@@ -15800,6 +15812,17 @@ def _active_worker_identity() -> str | None:
|
||||
if outcome.get("registered"):
|
||||
_WORKER_IDENTITY = identity
|
||||
_WORKER_GENERATION = generation
|
||||
# #975: registration is the only moment identity and fencing
|
||||
# epoch are both known, so the heartbeat supervisor is started
|
||||
# here. Without it ``last_heartbeat_at`` never left
|
||||
# ``started_at`` and every healthy client lost ownership at the
|
||||
# TTL.
|
||||
_start_worker_heartbeat(
|
||||
identity=identity,
|
||||
fencing_epoch=outcome.get("fencing_epoch"),
|
||||
hints=hints,
|
||||
generation=generation,
|
||||
)
|
||||
return identity
|
||||
if not outcome.get("collision"):
|
||||
return None
|
||||
@@ -15810,6 +15833,78 @@ def _active_worker_identity() -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def _start_worker_heartbeat(
|
||||
*,
|
||||
identity: str,
|
||||
fencing_epoch,
|
||||
hints: dict,
|
||||
generation: str | None,
|
||||
) -> None:
|
||||
"""Attach a heartbeat supervisor to the registration just created (#975).
|
||||
|
||||
Never raises: a supervisor that cannot start leaves the row exactly as
|
||||
``register()`` wrote it, which is the pre-#975 behaviour, rather than
|
||||
failing the tool call that happened to trigger lazy registration.
|
||||
|
||||
Under pytest the thread is deliberately not started. Tests drive
|
||||
``WorkerHeartbeatSupervisor`` directly with an injected clock, so the suite
|
||||
proves the lifecycle without leaving background sqlite writers behind.
|
||||
"""
|
||||
global _WORKER_HEARTBEAT_SUPERVISOR
|
||||
if _WORKER_HEARTBEAT_SUPERVISOR is not None:
|
||||
return
|
||||
registry = _worker_registry()
|
||||
if registry is None or fencing_epoch is None:
|
||||
return
|
||||
try:
|
||||
supervisor = mcp_worker_identity.WorkerHeartbeatSupervisor(
|
||||
registry,
|
||||
worker_identity=identity,
|
||||
fencing_epoch=int(fencing_epoch),
|
||||
session_id=hints.get("session_id"),
|
||||
generation_id=generation,
|
||||
client_name=hints.get("client_name"),
|
||||
pid=os.getpid(),
|
||||
ttl_seconds=mcp_worker_identity.DEFAULT_HEARTBEAT_TTL_SECONDS,
|
||||
interval_seconds=os.environ.get(
|
||||
mcp_worker_identity.HEARTBEAT_INTERVAL_ENV
|
||||
),
|
||||
)
|
||||
_WORKER_HEARTBEAT_SUPERVISOR = supervisor
|
||||
if not mcp_daemon_guard.is_pytest_runtime():
|
||||
supervisor.start()
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
def _worker_heartbeat_status() -> dict:
|
||||
"""Read-only heartbeat observability for ``gitea_get_runtime_context`` (#975).
|
||||
|
||||
Reports the unsupervised case explicitly rather than omitting the key, so an
|
||||
operator can tell "no supervisor" apart from "supervisor with no beats yet".
|
||||
"""
|
||||
supervisor = _WORKER_HEARTBEAT_SUPERVISOR
|
||||
if supervisor is None:
|
||||
return {
|
||||
"supervised": False,
|
||||
"running": False,
|
||||
"reasons": [
|
||||
"no worker heartbeat supervisor is attached to this process; the "
|
||||
"registration is not being renewed and will go stale at its TTL"
|
||||
],
|
||||
}
|
||||
try:
|
||||
status = dict(supervisor.status())
|
||||
except Exception as exc:
|
||||
return {
|
||||
"supervised": True,
|
||||
"running": False,
|
||||
"reasons": [f"heartbeat status unavailable: {type(exc).__name__}: {exc}"],
|
||||
}
|
||||
status.setdefault("reasons", [])
|
||||
return status
|
||||
|
||||
|
||||
def _active_role_kind_safe() -> str | None:
|
||||
"""Best-effort role for the registry record; never raises into a tool call.
|
||||
|
||||
@@ -19451,6 +19546,11 @@ def gitea_get_runtime_context(
|
||||
"fencing_epoch": provenance_assessment["fencing_epoch"],
|
||||
"conflicting_live_sessions": provenance_assessment["conflicting_live_sessions"],
|
||||
"provenance_assessment": provenance_assessment,
|
||||
# #975: whether this registration is actually being renewed. Read-only,
|
||||
# and it grants nothing — ownership still comes from the attachment
|
||||
# record above. It exists so "my heartbeat stopped" is diagnosable
|
||||
# before the TTL turns it into session_attachment_missing.
|
||||
"worker_heartbeat": _worker_heartbeat_status(),
|
||||
"unconsumed_gitea_env": unconsumed_env,
|
||||
"preflight_ready": preflight["preflight_ready"],
|
||||
"preflight_block_reasons": preflight["preflight_block_reasons"],
|
||||
@@ -21959,6 +22059,17 @@ def gitea_route_task_session(
|
||||
# self-recovery was removed from the read-only path (was _trigger_mcp_auto_restart).
|
||||
# Recovery is owned exclusively by the IDE/client reconnect path.
|
||||
|
||||
#: Every reason prefix ``_check_mcp_runtimes_diagnostics`` can emit. Callers
|
||||
#: raise its reasons as one RuntimeError, and #975 found that the preflight
|
||||
#: re-raise recognised only ``stale-runtime:``, silently dropping
|
||||
#: ``unsupported-env:`` while the capability resolver still failed on it. This
|
||||
#: lives beside the producer so a newly added reason family cannot be forgotten
|
||||
#: by a distant re-raise predicate again.
|
||||
RUNTIME_DIAGNOSTIC_HARD_PREFIXES: tuple[str, ...] = (
|
||||
"stale-runtime:",
|
||||
"unsupported-env:",
|
||||
)
|
||||
|
||||
|
||||
def _check_mcp_runtimes_diagnostics(task: str, matching_profiles: list[str]) -> list[str]:
|
||||
"""Read-only: report missing or stale MCP runtimes (no config or process mutation).
|
||||
|
||||
@@ -29,6 +29,7 @@ implements:
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import hashlib
|
||||
import os
|
||||
import re
|
||||
@@ -123,6 +124,103 @@ DEFAULT_REGISTRY_PATH = os.path.expanduser(
|
||||
#: worker within a single operator coffee break.
|
||||
DEFAULT_HEARTBEAT_TTL_SECONDS = 900.0
|
||||
|
||||
#: Optional operator override for the beat interval, in seconds. Clamped by
|
||||
#: :func:`heartbeat_interval_for` so it can never be set at or above the TTL.
|
||||
HEARTBEAT_INTERVAL_ENV = "GITEA_WORKER_HEARTBEAT_INTERVAL_SECONDS"
|
||||
|
||||
#: Never beat faster than this, so a misconfigured interval cannot turn the
|
||||
#: supervisor into a busy sqlite writer.
|
||||
MIN_HEARTBEAT_INTERVAL_SECONDS = 1.0
|
||||
|
||||
|
||||
def heartbeat_interval_for(
|
||||
ttl_seconds: float = DEFAULT_HEARTBEAT_TTL_SECONDS,
|
||||
override: str | float | None = None,
|
||||
) -> float:
|
||||
"""Beat interval for *ttl_seconds*: one third of the TTL, capped at a half.
|
||||
|
||||
#975. A third means two consecutive beats can be lost before the row is
|
||||
allowed to look stale, and the half-TTL cap is a hard ceiling so no
|
||||
override can produce an interval that expires the registration it is
|
||||
supposed to renew. That is the whole liveness contract: a healthy worker
|
||||
stays owned, and a worker that has genuinely stopped beating still expires
|
||||
on the configured TTL — this function never touches the TTL itself.
|
||||
"""
|
||||
try:
|
||||
ttl = float(ttl_seconds)
|
||||
except (TypeError, ValueError):
|
||||
ttl = DEFAULT_HEARTBEAT_TTL_SECONDS
|
||||
if not ttl > 0:
|
||||
ttl = DEFAULT_HEARTBEAT_TTL_SECONDS
|
||||
|
||||
interval = ttl / 3.0
|
||||
if override is not None:
|
||||
candidate = str(override).strip()
|
||||
if candidate:
|
||||
try:
|
||||
parsed = float(candidate)
|
||||
except (TypeError, ValueError):
|
||||
parsed = None
|
||||
if parsed is not None and parsed > 0:
|
||||
interval = parsed
|
||||
|
||||
ceiling = ttl / 2.0
|
||||
if interval > ceiling:
|
||||
interval = ceiling
|
||||
if interval < MIN_HEARTBEAT_INTERVAL_SECONDS:
|
||||
interval = min(MIN_HEARTBEAT_INTERVAL_SECONDS, ceiling)
|
||||
return interval
|
||||
|
||||
|
||||
#: Recorded-vs-presented pairs compared by :func:`_heartbeat_expectation_drift`.
|
||||
_HEARTBEAT_EXPECTATION_FIELDS = (
|
||||
"session_id",
|
||||
"generation_id",
|
||||
"client_name",
|
||||
"pid",
|
||||
)
|
||||
|
||||
|
||||
def _heartbeat_expectation_drift(
|
||||
record: dict[str, Any],
|
||||
*,
|
||||
expected_session_id: str | None = None,
|
||||
expected_generation_id: str | None = None,
|
||||
expected_client_name: str | None = None,
|
||||
expected_pid: int | None = None,
|
||||
) -> list[tuple[str, Any, Any]]:
|
||||
"""Return ``(field, recorded, presented)`` for every mismatched expectation.
|
||||
|
||||
Only supplied expectations are compared, so a caller that presents nothing
|
||||
gets the pre-#975 behaviour. ``client_name`` is compared through
|
||||
:func:`normalize_client_name` so one application's namespaces stay one
|
||||
client while genuinely different clients stay distinct.
|
||||
"""
|
||||
presented = {
|
||||
"session_id": expected_session_id,
|
||||
"generation_id": expected_generation_id,
|
||||
"client_name": (
|
||||
normalize_client_name(expected_client_name)
|
||||
if expected_client_name is not None
|
||||
else None
|
||||
),
|
||||
"pid": expected_pid,
|
||||
}
|
||||
drift: list[tuple[str, Any, Any]] = []
|
||||
for field in _HEARTBEAT_EXPECTATION_FIELDS:
|
||||
want = presented[field]
|
||||
if want is None:
|
||||
continue
|
||||
got = record.get(field)
|
||||
if field == "pid":
|
||||
match = got is not None and int(got) == int(want)
|
||||
else:
|
||||
match = got == want
|
||||
if not match:
|
||||
drift.append((field, got, want))
|
||||
return drift
|
||||
|
||||
|
||||
STATUS_ACTIVE = "active"
|
||||
STATUS_SUPERSEDED = "superseded"
|
||||
STATUS_RELEASED = "released"
|
||||
@@ -786,11 +884,25 @@ class WorkerRegistry:
|
||||
worker_identity: str,
|
||||
fencing_epoch: int,
|
||||
now: datetime | None = None,
|
||||
expected_session_id: str | None = None,
|
||||
expected_generation_id: str | None = None,
|
||||
expected_client_name: str | None = None,
|
||||
expected_pid: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Renew only the owning registration (#948 AC11).
|
||||
|
||||
A stale epoch is refused rather than silently renewed, so a superseded
|
||||
session that resumes cannot heartbeat its way back into ownership.
|
||||
|
||||
#975 adds optional keyword-only *expectations*. Every one supplied must
|
||||
match the recorded row or the renewal is refused. They exist because
|
||||
identity plus epoch cannot express "renew the row I registered, and only
|
||||
that row": a recycled PID, or a second session of the same client, would
|
||||
otherwise be renewable by the wrong beater. They are fencing tokens,
|
||||
never assertions — supplying one can only cause a refusal, never grant
|
||||
anything, and omitting them preserves the pre-#975 behaviour exactly.
|
||||
Drift reports the existing ``BLOCKER_FENCED`` rather than a new
|
||||
``blocker_kind``, because consumers switch on that value.
|
||||
"""
|
||||
stamp = _ts(now or _utc_now())
|
||||
with self._tx() as conn:
|
||||
@@ -807,6 +919,30 @@ class WorkerRegistry:
|
||||
"reasons": [f"no registration for {worker_identity!r}"],
|
||||
}
|
||||
record = self._row_to_record(row)
|
||||
drift = _heartbeat_expectation_drift(
|
||||
record,
|
||||
expected_session_id=expected_session_id,
|
||||
expected_generation_id=expected_generation_id,
|
||||
expected_client_name=expected_client_name,
|
||||
expected_pid=expected_pid,
|
||||
)
|
||||
if drift:
|
||||
return {
|
||||
"success": False,
|
||||
"renewed": False,
|
||||
"mutation_performed": False,
|
||||
"blocker_kind": BLOCKER_FENCED,
|
||||
"expectation_drift": drift,
|
||||
"reasons": [
|
||||
"presented worker expectations do not match the recorded "
|
||||
"registration, so this beater does not own the row: "
|
||||
+ "; ".join(
|
||||
f"{field} recorded {recorded!r}, presented {presented!r}"
|
||||
for field, recorded, presented in drift
|
||||
)
|
||||
+ " (#975)"
|
||||
],
|
||||
}
|
||||
if record["status"] != STATUS_ACTIVE:
|
||||
return {
|
||||
"success": False,
|
||||
@@ -1430,3 +1566,253 @@ def resolve_bound_remote(
|
||||
"explicitly to avoid host drift (#948)."
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
# --- Production heartbeat lifecycle (#975) --------------------------------
|
||||
#
|
||||
# #948 delivered ``WorkerRegistry.heartbeat()`` and nothing ever called it, so
|
||||
# ``last_heartbeat_at`` stayed pinned to ``started_at`` for every production
|
||||
# worker and ``heartbeat_ttl_seconds`` stopped being a liveness window at all —
|
||||
# it became a hard cap on how long any client could stay attached. This
|
||||
# supervisor is the missing caller.
|
||||
#
|
||||
# It is a daemon thread rather than an asyncio task or a per-request refresh
|
||||
# because the renewal has to survive an *idle* session: a client waiting on a
|
||||
# lease with no tool call in flight must not lose ownership, which rules out
|
||||
# request-driven refresh. Registration itself happens lazily inside a tool call
|
||||
# on the server's event loop, and the registry performs blocking
|
||||
# ``BEGIN IMMEDIATE`` sqlite writes, which must not run on that loop.
|
||||
# ``daemon=True`` is deliberate: a hard kill takes the thread down with the
|
||||
# process, so a dead worker still goes stale on the normal TTL.
|
||||
|
||||
#: Refusals that mean this worker no longer owns its row. Beating again could
|
||||
#: only ever be an attempt to renew ownership it has already lost, so the
|
||||
#: supervisor stops permanently instead of retrying.
|
||||
TERMINAL_HEARTBEAT_BLOCKERS = frozenset({BLOCKER_FENCED, BLOCKER_NO_ATTACHMENT})
|
||||
|
||||
|
||||
class WorkerHeartbeatSupervisor:
|
||||
"""Periodically renew exactly one worker registration.
|
||||
|
||||
Contract:
|
||||
|
||||
* It never registers. A supervisor exists only for an already-registered
|
||||
identity, so it cannot create a second registration or a second identity
|
||||
system.
|
||||
* Every beat presents the full expectation set, so it can renew only the row
|
||||
matching this exact client name, worker identity, session, generation and
|
||||
pid.
|
||||
* A terminal refusal (fenced or missing) stops it permanently and records
|
||||
why. A fenced session must never beat its way back into ownership.
|
||||
* A transient failure (a locked database, say) is counted and the loop
|
||||
continues, so one contended write does not silently end the heartbeat.
|
||||
* Nothing here raises into a caller. ``beat_once`` returns its outcome and
|
||||
the thread body swallows everything, because a heartbeat failure must
|
||||
degrade to "not renewed" and never crash a tool call.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
registry: WorkerRegistry,
|
||||
*,
|
||||
worker_identity: str,
|
||||
fencing_epoch: int,
|
||||
session_id: str | None = None,
|
||||
generation_id: str | None = None,
|
||||
client_name: str | None = None,
|
||||
pid: int | None = None,
|
||||
ttl_seconds: float = DEFAULT_HEARTBEAT_TTL_SECONDS,
|
||||
interval_seconds: float | None = None,
|
||||
clock=None,
|
||||
) -> None:
|
||||
self._registry = registry
|
||||
self.worker_identity = worker_identity
|
||||
self.fencing_epoch = int(fencing_epoch)
|
||||
self.session_id = session_id
|
||||
self.generation_id = generation_id
|
||||
self.client_name = (
|
||||
normalize_client_name(client_name) if client_name is not None else None
|
||||
)
|
||||
self.pid = int(pid) if pid is not None else None
|
||||
self.ttl_seconds = float(ttl_seconds)
|
||||
self.interval_seconds = (
|
||||
heartbeat_interval_for(self.ttl_seconds)
|
||||
if interval_seconds is None
|
||||
else heartbeat_interval_for(self.ttl_seconds, interval_seconds)
|
||||
)
|
||||
#: Injectable so every TTL test uses controlled time and no test waits
|
||||
#: for a real interval or a real TTL to elapse.
|
||||
self._clock = clock or _utc_now
|
||||
|
||||
self._stop_event = threading.Event()
|
||||
self._thread: threading.Thread | None = None
|
||||
self._lock = threading.Lock()
|
||||
self._atexit_registered = False
|
||||
|
||||
self.started = False
|
||||
self.stopped_reason: str | None = None
|
||||
self.beats_attempted = 0
|
||||
self.beats_renewed = 0
|
||||
self.transient_failures = 0
|
||||
self.last_beat_at: str | None = None
|
||||
self.last_result: dict[str, Any] | None = None
|
||||
|
||||
# -- one beat --
|
||||
|
||||
def beat_once(self, now: datetime | None = None) -> dict[str, Any]:
|
||||
"""Renew once. Never raises; returns the registry outcome or a failure."""
|
||||
if self.stopped_reason is not None:
|
||||
return {
|
||||
"success": False,
|
||||
"renewed": False,
|
||||
"beat_attempted": False,
|
||||
"reasons": [
|
||||
f"supervisor already stopped: {self.stopped_reason}"
|
||||
],
|
||||
}
|
||||
|
||||
stamp = now or self._clock()
|
||||
self.beats_attempted += 1
|
||||
try:
|
||||
result = self._registry.heartbeat(
|
||||
worker_identity=self.worker_identity,
|
||||
fencing_epoch=self.fencing_epoch,
|
||||
now=stamp,
|
||||
expected_session_id=self.session_id,
|
||||
expected_generation_id=self.generation_id,
|
||||
expected_client_name=self.client_name,
|
||||
expected_pid=self.pid,
|
||||
)
|
||||
except Exception as exc:
|
||||
# Transient by assumption: an unexpected error is not proof this
|
||||
# worker lost ownership, so it must not silently end the heartbeat.
|
||||
self.transient_failures += 1
|
||||
result = {
|
||||
"success": False,
|
||||
"renewed": False,
|
||||
"transient": True,
|
||||
"blocker_kind": None,
|
||||
"reasons": [f"heartbeat raised {type(exc).__name__}: {exc}"],
|
||||
}
|
||||
self.last_result = result
|
||||
return result
|
||||
|
||||
self.last_result = result
|
||||
if result.get("renewed"):
|
||||
self.beats_renewed += 1
|
||||
self.last_beat_at = result.get("last_heartbeat_at") or _ts(stamp)
|
||||
return result
|
||||
|
||||
blocker = result.get("blocker_kind")
|
||||
if blocker in TERMINAL_HEARTBEAT_BLOCKERS:
|
||||
self._stop_internal(
|
||||
reason=(
|
||||
f"refused with blocker_kind={blocker!r}: "
|
||||
+ "; ".join(result.get("reasons") or [])
|
||||
)
|
||||
)
|
||||
else:
|
||||
self.transient_failures += 1
|
||||
return result
|
||||
|
||||
# -- lifecycle --
|
||||
|
||||
def start(self) -> dict[str, Any]:
|
||||
"""Start the beat thread. Idempotent; safe to call from a tool call."""
|
||||
with self._lock:
|
||||
if self.stopped_reason is not None:
|
||||
return {
|
||||
"started": False,
|
||||
"reasons": [f"supervisor stopped: {self.stopped_reason}"],
|
||||
}
|
||||
if self._thread is not None and self._thread.is_alive():
|
||||
return {"started": True, "already_running": True, "reasons": []}
|
||||
|
||||
self._stop_event.clear()
|
||||
thread = threading.Thread(
|
||||
target=self._run,
|
||||
name=f"gitea-worker-heartbeat-{self.worker_identity}",
|
||||
daemon=True,
|
||||
)
|
||||
self._thread = thread
|
||||
self.started = True
|
||||
if not self._atexit_registered:
|
||||
# Orderly shutdown stops the heartbeat. A hard kill does not
|
||||
# run this, which is correct: the row must then go stale.
|
||||
atexit.register(self._atexit_stop)
|
||||
self._atexit_registered = True
|
||||
thread.start()
|
||||
return {"started": True, "already_running": False, "reasons": []}
|
||||
|
||||
def _run(self) -> None:
|
||||
while not self._stop_event.is_set():
|
||||
# Wait first: registration already stamped a fresh heartbeat, so an
|
||||
# immediate beat would be a redundant write on every server launch.
|
||||
if self._stop_event.wait(self.interval_seconds):
|
||||
return
|
||||
try:
|
||||
self.beat_once()
|
||||
except Exception:
|
||||
# beat_once is already total; this is the last-resort guard that
|
||||
# keeps a supervisor thread from dying silently.
|
||||
self.transient_failures += 1
|
||||
if self.stopped_reason is not None:
|
||||
return
|
||||
|
||||
def stop(self, reason: str = "stopped") -> dict[str, Any]:
|
||||
"""Stop beating. Idempotent, and prompt because the loop waits on an Event."""
|
||||
self._stop_internal(reason=reason)
|
||||
thread = self._thread
|
||||
if thread is not None and thread.is_alive():
|
||||
thread.join(timeout=max(1.0, min(5.0, self.interval_seconds)))
|
||||
return {
|
||||
"stopped": True,
|
||||
"reason": self.stopped_reason,
|
||||
"thread_alive": bool(thread is not None and thread.is_alive()),
|
||||
}
|
||||
|
||||
def _stop_internal(self, *, reason: str) -> None:
|
||||
if self.stopped_reason is None:
|
||||
self.stopped_reason = reason
|
||||
self._stop_event.set()
|
||||
if self._atexit_registered:
|
||||
try:
|
||||
atexit.unregister(self._atexit_stop)
|
||||
except Exception:
|
||||
pass
|
||||
self._atexit_registered = False
|
||||
|
||||
def _atexit_stop(self) -> None:
|
||||
try:
|
||||
self.stop(reason="process exit")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# -- observability --
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
"""Read-only observability payload; safe to embed in a tool result."""
|
||||
thread = self._thread
|
||||
return {
|
||||
"supervised": True,
|
||||
"worker_identity": self.worker_identity,
|
||||
"fencing_epoch": self.fencing_epoch,
|
||||
"session_id": self.session_id,
|
||||
"generation_id": self.generation_id,
|
||||
"client_name": self.client_name,
|
||||
"pid": self.pid,
|
||||
"heartbeat_ttl_seconds": self.ttl_seconds,
|
||||
"heartbeat_interval_seconds": self.interval_seconds,
|
||||
"started": self.started,
|
||||
"running": bool(
|
||||
thread is not None
|
||||
and thread.is_alive()
|
||||
and self.stopped_reason is None
|
||||
),
|
||||
"stopped_reason": self.stopped_reason,
|
||||
"beats_attempted": self.beats_attempted,
|
||||
"beats_renewed": self.beats_renewed,
|
||||
"transient_failures": self.transient_failures,
|
||||
"last_heartbeat_at": self.last_beat_at,
|
||||
"last_blocker_kind": (self.last_result or {}).get("blocker_kind"),
|
||||
}
|
||||
|
||||
@@ -0,0 +1,759 @@
|
||||
"""Client-identity env recognition and the production heartbeat lifecycle (#975).
|
||||
|
||||
Two defects left behind by #948, covered here against the issue's acceptance
|
||||
criteria:
|
||||
|
||||
1. ``GITEA_MCP_CLIENT`` / ``GITEA_MCP_CLIENT_INSTANCE`` /
|
||||
``GITEA_MCP_CLIENT_SESSION`` were consumed by production while absent from
|
||||
the recognized-key allowlist, so the peer-env scan classified them as
|
||||
unsupported overrides and the capability resolver refused every mutation
|
||||
fleet-wide.
|
||||
2. ``WorkerRegistry.heartbeat()`` had no production caller, so
|
||||
``last_heartbeat_at`` never left ``started_at`` and ``heartbeat_ttl_seconds``
|
||||
became a hard cap on attachment rather than a liveness window.
|
||||
|
||||
Every TTL assertion here uses an injected clock (AC13). No test waits for a
|
||||
real TTL. The one test that exercises the real thread loop uses a
|
||||
millisecond-scale interval and polls with a bounded timeout, never a TTL.
|
||||
|
||||
All client, session, and generation identifiers are synthetic.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import gitea_config
|
||||
import mcp_worker_identity as mwi
|
||||
|
||||
|
||||
NOW = datetime(2026, 7, 30, 0, 40, 9, tzinfo=timezone.utc)
|
||||
|
||||
TTL = 900.0
|
||||
|
||||
#: The three keys production reads as client identity.
|
||||
CLIENT_IDENTITY_ENV_KEYS = (
|
||||
"GITEA_MCP_CLIENT",
|
||||
"GITEA_MCP_CLIENT_INSTANCE",
|
||||
"GITEA_MCP_CLIENT_SESSION",
|
||||
)
|
||||
|
||||
|
||||
def _registry() -> mwi.WorkerRegistry:
|
||||
"""A registry on a throwaway path; never the operator's real one."""
|
||||
handle, path = tempfile.mkstemp(suffix=".sqlite3")
|
||||
os.close(handle)
|
||||
os.unlink(path)
|
||||
return mwi.WorkerRegistry(path)
|
||||
|
||||
|
||||
def _attach(
|
||||
registry: mwi.WorkerRegistry,
|
||||
*,
|
||||
client: str = "claude_code",
|
||||
session: str = "sess-0001",
|
||||
generation: str = "gen-0001",
|
||||
profile: str = "prgs-author",
|
||||
role: str = "author",
|
||||
pid: int = 4242,
|
||||
now: datetime = NOW,
|
||||
ttl: float = TTL,
|
||||
) -> dict:
|
||||
"""Register one synthetic worker and return the outcome plus its identity."""
|
||||
identity = mwi.generate_worker_identity(client, session, now=now)
|
||||
outcome = registry.register(
|
||||
worker_identity=identity,
|
||||
client_name=client,
|
||||
client_instance_id=f"inst-{session}",
|
||||
session_id=session,
|
||||
generation_id=generation,
|
||||
role=role,
|
||||
profile=profile,
|
||||
pid=pid,
|
||||
heartbeat_ttl_seconds=ttl,
|
||||
now=now,
|
||||
)
|
||||
outcome["identity"] = identity
|
||||
return outcome
|
||||
|
||||
|
||||
def _supervisor(
|
||||
registry: mwi.WorkerRegistry,
|
||||
outcome: dict,
|
||||
*,
|
||||
client: str = "claude_code",
|
||||
session: str = "sess-0001",
|
||||
generation: str = "gen-0001",
|
||||
pid: int = 4242,
|
||||
ttl: float = TTL,
|
||||
interval_seconds: float | None = None,
|
||||
clock=None,
|
||||
) -> mwi.WorkerHeartbeatSupervisor:
|
||||
"""A supervisor bound to *outcome*'s registration. Never started here."""
|
||||
return mwi.WorkerHeartbeatSupervisor(
|
||||
registry,
|
||||
worker_identity=outcome["identity"],
|
||||
fencing_epoch=outcome["fencing_epoch"],
|
||||
session_id=session,
|
||||
generation_id=generation,
|
||||
client_name=client,
|
||||
pid=pid,
|
||||
ttl_seconds=ttl,
|
||||
interval_seconds=interval_seconds,
|
||||
clock=clock,
|
||||
)
|
||||
|
||||
|
||||
class _ManualClock:
|
||||
"""Controlled time. Nothing in this suite sleeps to reach a TTL."""
|
||||
|
||||
def __init__(self, start: datetime = NOW) -> None:
|
||||
self.now = start
|
||||
|
||||
def __call__(self) -> datetime:
|
||||
return self.now
|
||||
|
||||
def advance(self, seconds: float) -> datetime:
|
||||
self.now = self.now + timedelta(seconds=seconds)
|
||||
return self.now
|
||||
|
||||
|
||||
# --- AC1-AC3: environment recognition ------------------------------------
|
||||
|
||||
|
||||
class ClientIdentityEnvRecognitionTests(unittest.TestCase):
|
||||
"""AC1-AC3: the keys production consumes are the keys the gate accepts."""
|
||||
|
||||
def test_each_client_identity_key_is_recognized(self):
|
||||
# AC1: individually asserted, so a partial fix cannot pass.
|
||||
for key in CLIENT_IDENTITY_ENV_KEYS:
|
||||
with self.subTest(key=key):
|
||||
self.assertIn(
|
||||
key,
|
||||
gitea_config.RECOGNIZED_GITEA_ENV_KEYS,
|
||||
f"{key} is read by the server but not recognized by the gate",
|
||||
)
|
||||
|
||||
def test_recognized_client_identity_keys_are_not_unconsumed(self):
|
||||
# AC3 at the classification layer the resolver reads: with only these
|
||||
# keys set there is no unsupported override to raise a blocker from.
|
||||
env = {
|
||||
"GITEA_MCP_CLIENT": "claude_code",
|
||||
"GITEA_MCP_CLIENT_INSTANCE": "inst-0001",
|
||||
"GITEA_MCP_CLIENT_SESSION": "sess-0001",
|
||||
}
|
||||
self.assertEqual(gitea_config.get_unconsumed_gitea_env_overrides(env), {})
|
||||
|
||||
def test_unknown_gitea_key_is_still_unsupported(self):
|
||||
# AC2: the gate is not broadened. An unrecognised override is refused
|
||||
# exactly as before, and the recognised siblings do not shield it.
|
||||
env = {
|
||||
"GITEA_MCP_CLIENT": "claude_code",
|
||||
"GITEA_TOTALLY_INVENTED_OVERRIDE": "1",
|
||||
}
|
||||
unconsumed = gitea_config.get_unconsumed_gitea_env_overrides(env)
|
||||
self.assertIn("GITEA_TOTALLY_INVENTED_OVERRIDE", unconsumed)
|
||||
self.assertNotIn("GITEA_MCP_CLIENT", unconsumed)
|
||||
|
||||
def test_no_new_prefix_admits_arbitrary_client_keys(self):
|
||||
# A prefix would have been the lazy fix; it would admit anything
|
||||
# beginning GITEA_MCP_CLIENT, including typos and future unknowns.
|
||||
env = {"GITEA_MCP_CLIENT_SOMETHING_ELSE": "x"}
|
||||
self.assertIn(
|
||||
"GITEA_MCP_CLIENT_SOMETHING_ELSE",
|
||||
gitea_config.get_unconsumed_gitea_env_overrides(env),
|
||||
)
|
||||
|
||||
def test_server_consumes_exactly_the_recognized_key_names(self):
|
||||
# Guards the actual defect class: the constants production reads and the
|
||||
# allowlist drifting apart again.
|
||||
import gitea_mcp_server as mcp_server
|
||||
|
||||
for name in (
|
||||
mcp_server.CLIENT_NAME_ENV,
|
||||
mcp_server.CLIENT_INSTANCE_ENV,
|
||||
mcp_server.CLIENT_SESSION_ENV,
|
||||
):
|
||||
with self.subTest(env_key=name):
|
||||
self.assertIn(name, gitea_config.RECOGNIZED_GITEA_ENV_KEYS)
|
||||
|
||||
|
||||
class RuntimeDiagnosticReasonPropagationTests(unittest.TestCase):
|
||||
"""`unsupported-env:` must not be swallowed where `stale-runtime:` is not."""
|
||||
|
||||
def test_both_reason_families_are_hard_prefixes(self):
|
||||
import gitea_mcp_server as mcp_server
|
||||
|
||||
self.assertIn("stale-runtime:", mcp_server.RUNTIME_DIAGNOSTIC_HARD_PREFIXES)
|
||||
self.assertIn("unsupported-env:", mcp_server.RUNTIME_DIAGNOSTIC_HARD_PREFIXES)
|
||||
|
||||
def test_unsupported_env_message_is_recognized_as_hard(self):
|
||||
import gitea_mcp_server as mcp_server
|
||||
|
||||
message = (
|
||||
"unsupported-env: Unsupported GITEA_* environment variable override(s) "
|
||||
"detected: GITEA_INVENTED=1. Unknown env overrides are unsupported."
|
||||
)
|
||||
self.assertTrue(
|
||||
any(p in message for p in mcp_server.RUNTIME_DIAGNOSTIC_HARD_PREFIXES)
|
||||
)
|
||||
|
||||
|
||||
# --- Interval policy ------------------------------------------------------
|
||||
|
||||
|
||||
class HeartbeatIntervalPolicyTests(unittest.TestCase):
|
||||
"""The interval must always leave room for lost beats inside the TTL."""
|
||||
|
||||
def test_default_interval_is_one_third_of_ttl(self):
|
||||
self.assertAlmostEqual(mwi.heartbeat_interval_for(900.0), 300.0)
|
||||
|
||||
def test_interval_is_capped_at_half_the_ttl(self):
|
||||
# An override may not exceed the ceiling, or it would expire the very
|
||||
# registration it exists to renew.
|
||||
self.assertAlmostEqual(mwi.heartbeat_interval_for(900.0, 5000), 450.0)
|
||||
|
||||
def test_override_below_the_ceiling_is_honoured(self):
|
||||
self.assertAlmostEqual(mwi.heartbeat_interval_for(900.0, 120), 120.0)
|
||||
|
||||
def test_unparsable_or_nonpositive_override_falls_back_to_policy(self):
|
||||
for bad in ("", "abc", "0", "-5", None):
|
||||
with self.subTest(override=bad):
|
||||
self.assertAlmostEqual(mwi.heartbeat_interval_for(900.0, bad), 300.0)
|
||||
|
||||
def test_interval_is_always_strictly_below_the_ttl(self):
|
||||
for ttl in (1.0, 30.0, 900.0, 14400.0):
|
||||
with self.subTest(ttl=ttl):
|
||||
self.assertLess(mwi.heartbeat_interval_for(ttl), ttl)
|
||||
|
||||
|
||||
# --- AC4-AC6: a healthy worker stays owned --------------------------------
|
||||
|
||||
|
||||
class HealthyWorkerRemainsOwnedTests(unittest.TestCase):
|
||||
"""AC4-AC6: the TTL becomes a liveness window instead of an attachment cap."""
|
||||
|
||||
def test_unsupervised_worker_goes_unowned_at_the_ttl(self):
|
||||
# The regression itself, asserted so the fix cannot be mistaken for a
|
||||
# TTL change: without a beater the row still expires exactly as before.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
record = registry.get(outcome["identity"])
|
||||
self.assertEqual(record["started_at"], record["last_heartbeat_at"])
|
||||
|
||||
after_ttl = NOW + timedelta(seconds=TTL + 1)
|
||||
liveness = registry.is_live(record, now=after_ttl, pid_alive=True)
|
||||
self.assertFalse(liveness["live"])
|
||||
|
||||
assessment = mwi.assess_provenance(
|
||||
registry=registry,
|
||||
worker_identity=outcome["identity"],
|
||||
generation_id="gen-0001",
|
||||
env={"GITEA_CLIENT_MANAGED": "1"},
|
||||
native_transport_bound=True,
|
||||
now=after_ttl,
|
||||
pid_alive_probe=lambda _pid: True,
|
||||
)
|
||||
self.assertEqual(assessment["session_ownership"], "unowned")
|
||||
self.assertEqual(assessment["blocker_kind"], mwi.BLOCKER_NO_ATTACHMENT)
|
||||
|
||||
def test_supervised_worker_stays_owned_past_two_full_ttls(self):
|
||||
# AC4: beyond two complete TTL periods of simulated time.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
clock = _ManualClock()
|
||||
supervisor = _supervisor(registry, outcome, clock=clock)
|
||||
|
||||
interval = supervisor.interval_seconds
|
||||
self.assertLess(interval, TTL)
|
||||
|
||||
elapsed = 0.0
|
||||
target = TTL * 2 + interval
|
||||
while elapsed < target:
|
||||
clock.advance(interval)
|
||||
elapsed += interval
|
||||
result = supervisor.beat_once()
|
||||
self.assertTrue(result["renewed"], result.get("reasons"))
|
||||
|
||||
assessment = mwi.assess_provenance(
|
||||
registry=registry,
|
||||
worker_identity=outcome["identity"],
|
||||
generation_id="gen-0001",
|
||||
env={"GITEA_CLIENT_MANAGED": "1"},
|
||||
native_transport_bound=True,
|
||||
now=clock.now,
|
||||
pid_alive_probe=lambda _pid: True,
|
||||
)
|
||||
self.assertEqual(
|
||||
assessment["session_ownership"],
|
||||
"owned",
|
||||
f"ownership lost after {elapsed:.0f}s of simulated time",
|
||||
)
|
||||
|
||||
self.assertGreater(elapsed, TTL * 2)
|
||||
|
||||
def test_last_heartbeat_advances_and_diverges_from_started_at(self):
|
||||
# AC5.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
clock = _ManualClock()
|
||||
supervisor = _supervisor(registry, outcome, clock=clock)
|
||||
|
||||
before = registry.get(outcome["identity"])
|
||||
clock.advance(300)
|
||||
supervisor.beat_once()
|
||||
after = registry.get(outcome["identity"])
|
||||
|
||||
self.assertEqual(before["started_at"], after["started_at"])
|
||||
self.assertNotEqual(after["started_at"], after["last_heartbeat_at"])
|
||||
self.assertGreater(after["last_heartbeat_at"], before["last_heartbeat_at"])
|
||||
|
||||
def test_repeated_heartbeats_create_no_duplicate_rows(self):
|
||||
# AC6: heartbeating is not a second registration path.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
clock = _ManualClock()
|
||||
supervisor = _supervisor(registry, outcome, clock=clock)
|
||||
|
||||
for _ in range(5):
|
||||
clock.advance(300)
|
||||
self.assertTrue(supervisor.beat_once()["renewed"])
|
||||
|
||||
workers = registry.list_workers()
|
||||
self.assertEqual(len(workers), 1)
|
||||
self.assertEqual(workers[0]["worker_identity"], outcome["identity"])
|
||||
self.assertEqual(supervisor.status()["beats_renewed"], 5)
|
||||
|
||||
|
||||
# --- AC7, AC8, AC12: only the owner may renew -----------------------------
|
||||
|
||||
|
||||
class HeartbeatOwnershipFencingTests(unittest.TestCase):
|
||||
"""AC7/AC8/AC12: one worker can never refresh another worker's row."""
|
||||
|
||||
def test_two_client_names_stay_distinct_rows_and_identities(self):
|
||||
# AC7.
|
||||
registry = _registry()
|
||||
first = _attach(registry, client="claude_code", session="sess-a")
|
||||
second = _attach(registry, client="gemini", session="sess-b")
|
||||
|
||||
self.assertNotEqual(first["identity"], second["identity"])
|
||||
rows = {w["worker_identity"]: w for w in registry.list_workers()}
|
||||
self.assertEqual(len(rows), 2)
|
||||
self.assertEqual(rows[first["identity"]]["client_name"], "claude_code")
|
||||
self.assertEqual(rows[second["identity"]]["client_name"], "gemini")
|
||||
|
||||
def test_wrong_session_cannot_refresh_the_row(self):
|
||||
# AC8, session dimension.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry, session="sess-real")
|
||||
result = registry.heartbeat(
|
||||
worker_identity=outcome["identity"],
|
||||
fencing_epoch=outcome["fencing_epoch"],
|
||||
expected_session_id="sess-impostor",
|
||||
)
|
||||
self.assertFalse(result["renewed"])
|
||||
self.assertEqual(result["blocker_kind"], mwi.BLOCKER_FENCED)
|
||||
self.assertIn(
|
||||
"session_id", [field for field, _r, _p in result["expectation_drift"]]
|
||||
)
|
||||
|
||||
def test_wrong_generation_cannot_refresh_the_row(self):
|
||||
# AC8, generation dimension.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry, generation="gen-real")
|
||||
result = registry.heartbeat(
|
||||
worker_identity=outcome["identity"],
|
||||
fencing_epoch=outcome["fencing_epoch"],
|
||||
expected_generation_id="gen-other",
|
||||
)
|
||||
self.assertFalse(result["renewed"])
|
||||
self.assertEqual(result["blocker_kind"], mwi.BLOCKER_FENCED)
|
||||
|
||||
def test_wrong_client_name_cannot_refresh_the_row(self):
|
||||
registry = _registry()
|
||||
outcome = _attach(registry, client="claude_code")
|
||||
result = registry.heartbeat(
|
||||
worker_identity=outcome["identity"],
|
||||
fencing_epoch=outcome["fencing_epoch"],
|
||||
expected_client_name="gemini",
|
||||
)
|
||||
self.assertFalse(result["renewed"])
|
||||
self.assertEqual(result["blocker_kind"], mwi.BLOCKER_FENCED)
|
||||
|
||||
def test_wrong_pid_cannot_refresh_the_row(self):
|
||||
registry = _registry()
|
||||
outcome = _attach(registry, pid=4242)
|
||||
result = registry.heartbeat(
|
||||
worker_identity=outcome["identity"],
|
||||
fencing_epoch=outcome["fencing_epoch"],
|
||||
expected_pid=9999,
|
||||
)
|
||||
self.assertFalse(result["renewed"])
|
||||
self.assertEqual(result["blocker_kind"], mwi.BLOCKER_FENCED)
|
||||
|
||||
def test_namespaces_of_one_application_share_a_client_name(self):
|
||||
# Several namespaces of one application must stay one client, while
|
||||
# separate applications stay distinct. Aliases normalise, and the
|
||||
# expectation check compares normalised values, so an author and a
|
||||
# reviewer namespace of the same product are not treated as impostors.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry, client="claude")
|
||||
self.assertEqual(
|
||||
registry.get(outcome["identity"])["client_name"], "claude_code"
|
||||
)
|
||||
result = registry.heartbeat(
|
||||
worker_identity=outcome["identity"],
|
||||
fencing_epoch=outcome["fencing_epoch"],
|
||||
expected_client_name="claude_desktop",
|
||||
)
|
||||
self.assertTrue(result["renewed"], result.get("reasons"))
|
||||
|
||||
def test_matching_expectations_renew_normally(self):
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
result = registry.heartbeat(
|
||||
worker_identity=outcome["identity"],
|
||||
fencing_epoch=outcome["fencing_epoch"],
|
||||
expected_session_id="sess-0001",
|
||||
expected_generation_id="gen-0001",
|
||||
expected_client_name="claude_code",
|
||||
expected_pid=4242,
|
||||
)
|
||||
self.assertTrue(result["renewed"], result.get("reasons"))
|
||||
|
||||
def test_omitted_expectations_preserve_pre_fix_behaviour(self):
|
||||
# AC12: the #948 contract is unchanged for callers presenting nothing.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
self.assertTrue(
|
||||
registry.heartbeat(
|
||||
worker_identity=outcome["identity"],
|
||||
fencing_epoch=outcome["fencing_epoch"],
|
||||
)["renewed"]
|
||||
)
|
||||
|
||||
def test_stale_fencing_epoch_is_still_refused(self):
|
||||
# AC12: the pre-existing epoch fence is untouched.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
result = registry.heartbeat(
|
||||
worker_identity=outcome["identity"],
|
||||
fencing_epoch=int(outcome["fencing_epoch"]) + 1,
|
||||
)
|
||||
self.assertFalse(result["renewed"])
|
||||
self.assertEqual(result["blocker_kind"], mwi.BLOCKER_FENCED)
|
||||
|
||||
def test_fenced_supervisor_stops_permanently(self):
|
||||
# A fenced session must never beat its way back into ownership.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
supervisor = mwi.WorkerHeartbeatSupervisor(
|
||||
registry,
|
||||
worker_identity=outcome["identity"],
|
||||
fencing_epoch=int(outcome["fencing_epoch"]) + 1,
|
||||
session_id="sess-0001",
|
||||
generation_id="gen-0001",
|
||||
client_name="claude_code",
|
||||
pid=4242,
|
||||
ttl_seconds=TTL,
|
||||
clock=_ManualClock(),
|
||||
)
|
||||
first = supervisor.beat_once()
|
||||
self.assertFalse(first["renewed"])
|
||||
self.assertIsNotNone(supervisor.stopped_reason)
|
||||
|
||||
# A second attempt does not even reach the registry.
|
||||
second = supervisor.beat_once()
|
||||
self.assertFalse(second["renewed"])
|
||||
self.assertFalse(second["beat_attempted"])
|
||||
self.assertEqual(supervisor.status()["beats_attempted"], 1)
|
||||
|
||||
def test_missing_registration_stops_the_supervisor(self):
|
||||
registry = _registry()
|
||||
supervisor = mwi.WorkerHeartbeatSupervisor(
|
||||
registry,
|
||||
worker_identity=mwi.generate_worker_identity("grok", "sess-gone", now=NOW),
|
||||
fencing_epoch=1,
|
||||
ttl_seconds=TTL,
|
||||
clock=_ManualClock(),
|
||||
)
|
||||
result = supervisor.beat_once()
|
||||
self.assertEqual(result["blocker_kind"], mwi.BLOCKER_NO_ATTACHMENT)
|
||||
self.assertIsNotNone(supervisor.stopped_reason)
|
||||
|
||||
|
||||
# --- AC9, AC10, AC11: stopping and expiry ---------------------------------
|
||||
|
||||
|
||||
class ShutdownAndExpiryTests(unittest.TestCase):
|
||||
"""AC9-AC11: orderly shutdown stops beating; dead workers still expire."""
|
||||
|
||||
def test_stop_ends_the_heartbeat(self):
|
||||
# AC9.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
clock = _ManualClock()
|
||||
supervisor = _supervisor(registry, outcome, clock=clock)
|
||||
|
||||
clock.advance(300)
|
||||
self.assertTrue(supervisor.beat_once()["renewed"])
|
||||
supervisor.stop(reason="orderly shutdown")
|
||||
|
||||
self.assertFalse(supervisor.status()["running"])
|
||||
self.assertEqual(supervisor.stopped_reason, "orderly shutdown")
|
||||
|
||||
clock.advance(300)
|
||||
after_stop = supervisor.beat_once()
|
||||
self.assertFalse(after_stop["renewed"])
|
||||
self.assertFalse(after_stop["beat_attempted"])
|
||||
|
||||
def test_stop_is_idempotent_and_keeps_the_first_reason(self):
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
supervisor = _supervisor(registry, outcome, clock=_ManualClock())
|
||||
supervisor.stop(reason="orderly shutdown")
|
||||
supervisor.stop(reason="second call")
|
||||
self.assertEqual(supervisor.stopped_reason, "orderly shutdown")
|
||||
|
||||
def test_stopped_worker_becomes_unowned_once_the_ttl_elapses(self):
|
||||
# AC10: stopping beating restores normal expiration; the fix does not
|
||||
# make a worker immortal.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
clock = _ManualClock()
|
||||
supervisor = _supervisor(registry, outcome, clock=clock)
|
||||
|
||||
clock.advance(300)
|
||||
supervisor.beat_once()
|
||||
supervisor.stop(reason="orderly shutdown")
|
||||
|
||||
last_beat = clock.now
|
||||
clock.advance(TTL + 1)
|
||||
record = registry.get(outcome["identity"])
|
||||
liveness = registry.is_live(record, now=clock.now, pid_alive=True)
|
||||
self.assertFalse(liveness["live"])
|
||||
self.assertGreater(liveness["heartbeat_age_seconds"], TTL)
|
||||
|
||||
assessment = mwi.assess_provenance(
|
||||
registry=registry,
|
||||
worker_identity=outcome["identity"],
|
||||
generation_id="gen-0001",
|
||||
env={"GITEA_CLIENT_MANAGED": "1"},
|
||||
native_transport_bound=True,
|
||||
now=clock.now,
|
||||
pid_alive_probe=lambda _pid: True,
|
||||
)
|
||||
self.assertEqual(assessment["session_ownership"], "unowned")
|
||||
self.assertLess(last_beat, clock.now)
|
||||
|
||||
def test_dead_worker_expires_even_with_a_live_pid(self):
|
||||
# A recycled or long-lived pid must not keep a silent worker owned.
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
record = registry.get(outcome["identity"])
|
||||
liveness = registry.is_live(
|
||||
record, now=NOW + timedelta(seconds=TTL + 1), pid_alive=True
|
||||
)
|
||||
self.assertFalse(liveness["live"])
|
||||
self.assertFalse(liveness["heartbeat_fresh"])
|
||||
|
||||
def test_one_worker_expiring_does_not_stale_healthy_workers(self):
|
||||
# AC11.
|
||||
registry = _registry()
|
||||
healthy = _attach(registry, client="claude_code", session="sess-live")
|
||||
abandoned = _attach(registry, client="gemini", session="sess-dead")
|
||||
|
||||
clock = _ManualClock()
|
||||
supervisor = _supervisor(
|
||||
registry,
|
||||
healthy,
|
||||
client="claude_code",
|
||||
session="sess-live",
|
||||
clock=clock,
|
||||
)
|
||||
|
||||
elapsed = 0.0
|
||||
while elapsed < TTL + 300:
|
||||
clock.advance(supervisor.interval_seconds)
|
||||
elapsed += supervisor.interval_seconds
|
||||
self.assertTrue(supervisor.beat_once()["renewed"])
|
||||
|
||||
healthy_live = registry.is_live(
|
||||
registry.get(healthy["identity"]), now=clock.now, pid_alive=True
|
||||
)
|
||||
abandoned_live = registry.is_live(
|
||||
registry.get(abandoned["identity"]), now=clock.now, pid_alive=True
|
||||
)
|
||||
self.assertTrue(healthy_live["live"])
|
||||
self.assertFalse(abandoned_live["live"])
|
||||
|
||||
|
||||
# --- Failure handling and observability -----------------------------------
|
||||
|
||||
|
||||
class HeartbeatFailureHandlingTests(unittest.TestCase):
|
||||
"""A heartbeat failure never grants ownership and never crashes a caller."""
|
||||
|
||||
def test_transient_exception_is_counted_and_beating_continues(self):
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
clock = _ManualClock()
|
||||
supervisor = _supervisor(registry, outcome, clock=clock)
|
||||
|
||||
calls = {"n": 0}
|
||||
real = registry.heartbeat
|
||||
|
||||
def flaky(**kwargs):
|
||||
calls["n"] += 1
|
||||
if calls["n"] == 1:
|
||||
raise RuntimeError("database is locked")
|
||||
return real(**kwargs)
|
||||
|
||||
registry.heartbeat = flaky # type: ignore[method-assign]
|
||||
try:
|
||||
clock.advance(300)
|
||||
first = supervisor.beat_once()
|
||||
self.assertFalse(first["renewed"])
|
||||
self.assertTrue(first["transient"])
|
||||
self.assertIsNone(supervisor.stopped_reason)
|
||||
|
||||
clock.advance(300)
|
||||
second = supervisor.beat_once()
|
||||
self.assertTrue(second["renewed"], second.get("reasons"))
|
||||
finally:
|
||||
registry.heartbeat = real # type: ignore[method-assign]
|
||||
|
||||
status = supervisor.status()
|
||||
self.assertEqual(status["transient_failures"], 1)
|
||||
self.assertEqual(status["beats_renewed"], 1)
|
||||
self.assertTrue(status["supervised"])
|
||||
|
||||
def test_status_reports_an_unstarted_supervisor_honestly(self):
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
status = _supervisor(registry, outcome, clock=_ManualClock()).status()
|
||||
self.assertTrue(status["supervised"])
|
||||
self.assertFalse(status["started"])
|
||||
self.assertFalse(status["running"])
|
||||
self.assertEqual(status["beats_renewed"], 0)
|
||||
self.assertIsNone(status["stopped_reason"])
|
||||
|
||||
def test_runtime_context_reports_unsupervised_when_none_attached(self):
|
||||
# The observability surface must distinguish "no supervisor" from
|
||||
# "supervisor with no beats yet", or the regression is invisible again.
|
||||
import gitea_mcp_server as mcp_server
|
||||
|
||||
saved = mcp_server._WORKER_HEARTBEAT_SUPERVISOR
|
||||
mcp_server._WORKER_HEARTBEAT_SUPERVISOR = None
|
||||
try:
|
||||
status = mcp_server._worker_heartbeat_status()
|
||||
finally:
|
||||
mcp_server._WORKER_HEARTBEAT_SUPERVISOR = saved
|
||||
self.assertFalse(status["supervised"])
|
||||
self.assertFalse(status["running"])
|
||||
self.assertTrue(status["reasons"])
|
||||
|
||||
def test_runtime_context_surfaces_an_attached_supervisor(self):
|
||||
import gitea_mcp_server as mcp_server
|
||||
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
supervisor = _supervisor(registry, outcome, clock=_ManualClock())
|
||||
|
||||
saved = mcp_server._WORKER_HEARTBEAT_SUPERVISOR
|
||||
mcp_server._WORKER_HEARTBEAT_SUPERVISOR = supervisor
|
||||
try:
|
||||
status = mcp_server._worker_heartbeat_status()
|
||||
finally:
|
||||
mcp_server._WORKER_HEARTBEAT_SUPERVISOR = saved
|
||||
|
||||
self.assertTrue(status["supervised"])
|
||||
self.assertEqual(status["worker_identity"], outcome["identity"])
|
||||
self.assertEqual(status["heartbeat_ttl_seconds"], TTL)
|
||||
self.assertLess(status["heartbeat_interval_seconds"], TTL)
|
||||
|
||||
def test_status_failure_degrades_instead_of_raising(self):
|
||||
import gitea_mcp_server as mcp_server
|
||||
|
||||
class Exploding:
|
||||
def status(self):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
saved = mcp_server._WORKER_HEARTBEAT_SUPERVISOR
|
||||
mcp_server._WORKER_HEARTBEAT_SUPERVISOR = Exploding()
|
||||
try:
|
||||
status = mcp_server._worker_heartbeat_status()
|
||||
finally:
|
||||
mcp_server._WORKER_HEARTBEAT_SUPERVISOR = saved
|
||||
|
||||
self.assertTrue(status["supervised"])
|
||||
self.assertFalse(status["running"])
|
||||
self.assertTrue(status["reasons"])
|
||||
|
||||
|
||||
class SupervisorThreadLoopTests(unittest.TestCase):
|
||||
"""The real thread, on a millisecond interval. Never waits for a TTL."""
|
||||
|
||||
def test_thread_beats_repeatedly_then_stops_on_request(self):
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
beats = threading.Event()
|
||||
seen = {"n": 0}
|
||||
real = registry.heartbeat
|
||||
|
||||
def counting(**kwargs):
|
||||
result = real(**kwargs)
|
||||
seen["n"] += 1
|
||||
if seen["n"] >= 3:
|
||||
beats.set()
|
||||
return result
|
||||
|
||||
registry.heartbeat = counting # type: ignore[method-assign]
|
||||
supervisor = _supervisor(registry, outcome, ttl=1.0, interval_seconds=0.01)
|
||||
try:
|
||||
started = supervisor.start()
|
||||
self.assertTrue(started["started"])
|
||||
self.assertTrue(
|
||||
beats.wait(timeout=10.0), "heartbeat thread produced no beats"
|
||||
)
|
||||
self.assertTrue(supervisor.status()["running"])
|
||||
finally:
|
||||
supervisor.stop(reason="test teardown")
|
||||
registry.heartbeat = real # type: ignore[method-assign]
|
||||
|
||||
self.assertGreaterEqual(supervisor.status()["beats_renewed"], 3)
|
||||
self.assertFalse(supervisor.status()["running"])
|
||||
|
||||
# Stopping is prompt: the loop waits on an Event, not a sleep.
|
||||
deadline = time.monotonic() + 5.0
|
||||
while supervisor._thread is not None and supervisor._thread.is_alive():
|
||||
if time.monotonic() > deadline:
|
||||
self.fail("heartbeat thread did not exit after stop()")
|
||||
time.sleep(0.01)
|
||||
|
||||
def test_start_is_idempotent(self):
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
supervisor = _supervisor(registry, outcome, ttl=1.0, interval_seconds=0.05)
|
||||
try:
|
||||
self.assertFalse(supervisor.start()["already_running"])
|
||||
self.assertTrue(supervisor.start()["already_running"])
|
||||
finally:
|
||||
supervisor.stop(reason="test teardown")
|
||||
|
||||
def test_start_refuses_after_a_terminal_stop(self):
|
||||
registry = _registry()
|
||||
outcome = _attach(registry)
|
||||
supervisor = _supervisor(registry, outcome, clock=_ManualClock())
|
||||
supervisor.stop(reason="orderly shutdown")
|
||||
self.assertFalse(supervisor.start()["started"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user