core-lead/molecule-core
1
0
Fork 0
forked from molecule-ai/molecule-core
Code Pull Requests Activity
1d231ed295
molecule-core/workspace/heartbeat.py
Hongming Wang 1d231ed295 refactor(wedge): extract claude_sdk_executor wedge state into runtime_wedge module 
Prerequisite for the universal-runtime refactor (task #87) to move
claude_sdk_executor.py out of molecule-runtime into the claude-code
template repo. heartbeat.py had a hard import:

    from claude_sdk_executor import is_wedged, wedge_reason

which would break the moment the executor moves out of the runtime
package — the heartbeat would lose access to the wedge state used to
flip workspace status to degraded.

Extract the wedge state to a runtime-side module that the heartbeat
can keep importing regardless of which adapter executor is wedged:

  - workspace/runtime_wedge.py — single-flag state + mark_wedged /
    clear_wedge / is_wedged / wedge_reason / reset_for_test. Same
    semantics as the original claude_sdk_executor implementation
    (sticky first-write-wins, auto-clear on observed success). 100
    LOC of pure stateless helpers; lock-free ok because there's one
    executor per workspace process today.

  - workspace/claude_sdk_executor.py — drops the in-file definitions;
    re-exports the same names from runtime_wedge as a backwards-compat
    shim. Any third-party adapter that imported is_wedged / wedge_reason
    / _mark_sdk_wedged from claude_sdk_executor keeps working for one
    release cycle while they migrate to runtime_wedge.

  - workspace/heartbeat.py — _runtime_state_payload() now imports
    from runtime_wedge instead of claude_sdk_executor. Lazy-import
    pattern preserved; the docstring updated to explain the new
    cross-cutting source-of-truth.

Tests (10 new in test_runtime_wedge.py):
  - Default state (unwedged), mark sets flag, first-write-wins,
    clear restores healthy, clear-when-not-wedged is no-op,
    re-marking after clear is allowed
  - Re-export shim: each old name in claude_sdk_executor IS the
    runtime_wedge function (identity check), state is shared
    (marking via the executor shim is observable via runtime_wedge
    and vice versa)

Verification:
  - 1251/1251 workspace pytest pass (was 1241 after orphan deletion;
    +10 = exactly the new test_runtime_wedge.py cases)
  - All existing test_claude_sdk_executor.py cases (which call
    _mark_sdk_wedged via the shim) still pass

After this lands + the claude-code template image rebuilds with the
local claude_sdk_executor.py copy (template PR #13), the molecule-
core deletion of workspace/claude_sdk_executor.py becomes safe (the
shim deletion comes alongside the file deletion, since runtime_wedge
is the new public API).

See project memory `project_runtime_native_pluggable.md`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-27 00:08:53 -07:00

404 lines
18 KiB
Python
Raw Blame History

"""Heartbeat loop — alive signal + delegation status checker.
Every 30 seconds:
1. Send heartbeat to platform (alive signal with current_task, error_rate)
2. Check pending delegations — any results back?
3. Store completed delegation results for the agent to pick up
Resilient: recreates HTTP client on failure, auto-restarts on crash.
"""
import asyncio
import json
import logging
import os
import time
from pathlib import Path
import httpx
from platform_auth import auth_headers, refresh_cache, self_source_headers
def _runtime_state_payload() -> dict:
"""Build the {runtime_state, sample_error} portion of the heartbeat
body when SOME adapter executor has marked itself wedged. Returns
an empty dict when the runtime is healthy so the heartbeat payload
doesn't grow fields the platform doesn't need.
Source of truth is runtime_wedge (lives in molecule-runtime,
independent of any specific adapter). Pre task #87 this imported
from claude_sdk_executor — that worked because the executor was
bundled into molecule-runtime, but blocked moving it to the
claude-code template repo. The runtime_wedge module is now the
cross-cutting wedge-state holder; adapters mark/clear via it,
heartbeat reads it.
Imported lazily so a workspace whose runtime image somehow ships
without runtime_wedge (corrupt install, mid-rolling-deploy state)
keeps heartbeating — a missing import means "no wedge info; assume
healthy."
"""
try:
from runtime_wedge import is_wedged, wedge_reason
except Exception:
return {}
if not is_wedged():
return {}
return {
"runtime_state": "wedged",
# sample_error doubles as the human-readable banner text on the
# canvas's degraded card — keep it short and actionable.
"sample_error": wedge_reason(),
}
def _runtime_metadata_payload() -> dict:
"""Build the {runtime_metadata} portion of the heartbeat body —
adapter-declared capabilities + per-capability override values
(idle timeout, etc.). The platform reads this to route capabilities
to the right owner: native (adapter) vs fallback (platform).
Returns an empty dict if the adapter can't be loaded or introspected.
Heartbeat must NEVER fail because of capability discovery — observability
is more important than capability accuracy. The platform falls through
to its own defaults when fields are missing.
See project memory `project_runtime_native_pluggable.md` and
workspace/adapter_base.py:RuntimeCapabilities.
"""
try:
from adapters import get_adapter
# ADAPTER_MODULE wins over the runtime arg in get_adapter — pass
# an empty string to force the env-var path.
adapter_cls = get_adapter("")
adapter = adapter_cls()
caps = adapter.capabilities()
meta: dict = {"capabilities": caps.to_dict()}
idle = adapter.idle_timeout_override()
# Only include the override when it's a positive integer. None /
# zero / negative falls through to the platform's global default
# (env A2A_IDLE_TIMEOUT_SECONDS, default 5min) — that "absent
# field = use default" contract is what keeps the wire small.
if isinstance(idle, int) and idle > 0:
meta["idle_timeout_seconds"] = idle
return {"runtime_metadata": meta}
except Exception as e:
# debug-level: missing ADAPTER_MODULE in dev / test envs is normal
logger.debug("runtime_metadata: failed to read adapter caps: %s", e)
return {}
logger = logging.getLogger(__name__)
HEARTBEAT_INTERVAL = 30 # seconds
MAX_CONSECUTIVE_FAILURES = 10
MAX_SEEN_DELEGATION_IDS = 200
SELF_MESSAGE_COOLDOWN = 60 # seconds — minimum between self-messages to prevent loops
# Shared path — also used by cli_executor._read_delegation_results()
DELEGATION_RESULTS_FILE = os.environ.get("DELEGATION_RESULTS_FILE", "/tmp/delegation_results.jsonl")
class HeartbeatLoop:
def __init__(self, platform_url: str, workspace_id: str):
self.platform_url = platform_url
self.workspace_id = workspace_id
self.start_time = time.time()
self.error_count = 0
self.request_count = 0
self.active_tasks = 0
self.current_task = ""
self.sample_error = ""
self._task = None
self._consecutive_failures = 0
self._seen_delegation_ids: set[str] = set()
self._last_self_message_time = 0.0
self._parent_name: str | None = None # Cached after first lookup
@property
def error_rate(self) -> float:
if self.request_count == 0:
return 0.0
return self.error_count / self.request_count
def record_error(self, error: str):
self.error_count += 1
self.request_count += 1
self.sample_error = error
def record_success(self):
self.request_count += 1
def start(self):
self._task = asyncio.create_task(self._loop())
self._task.add_done_callback(self._on_done)
def _on_done(self, task):
if not task.cancelled() and task.exception():
logger.error("Heartbeat loop died: %s — restarting", task.exception())
self._task = asyncio.create_task(self._loop())
self._task.add_done_callback(self._on_done)
async def stop(self):
if self._task:
self._task.cancel()
try:
await self._task
except asyncio.CancelledError:
pass
async def _loop(self):
while True:
client = None
try:
client = httpx.AsyncClient(timeout=10.0)
while True:
# 1. Send heartbeat (Phase 30.1: include auth header if token known)
try:
body = {
"workspace_id": self.workspace_id,
"error_rate": self.error_rate,
"sample_error": self.sample_error,
"active_tasks": self.active_tasks,
"current_task": self.current_task,
"uptime_seconds": int(time.time() - self.start_time),
}
# Layer the runtime-wedge fields on top so a
# non-empty sample_error from the wedge wins
# over the (typically empty) heartbeat
# sample_error field. The platform reads
# runtime_state to flip status → degraded.
body.update(_runtime_state_payload())
body.update(_runtime_metadata_payload())
await client.post(
f"{self.platform_url}/registry/heartbeat",
json=body,
headers=auth_headers(),
)
self.error_count = 0
self.request_count = 0
self._consecutive_failures = 0
except Exception as e:
self._consecutive_failures += 1
# Issue #1877: if heartbeat 401'd, re-read the token from disk
# and retry once. This handles the platform's token-rotation race
# where WriteFilesToContainer hasn't finished writing the new
# token before the runtime boots and caches the old value.
is_401 = False
if isinstance(e, httpx.HTTPStatusError) and e.response.status_code == 401:
is_401 = True
if is_401:
logger.warning("Heartbeat 401 for %s — refreshing token cache and retrying once", self.workspace_id)
refresh_cache()
try:
retry_body = {
"workspace_id": self.workspace_id,
"error_rate": self.error_rate,
"sample_error": self.sample_error,
"active_tasks": self.active_tasks,
"current_task": self.current_task,
"uptime_seconds": int(time.time() - self.start_time),
}
retry_body.update(_runtime_state_payload())
await client.post(
f"{self.platform_url}/registry/heartbeat",
json=retry_body,
headers=auth_headers(),
)
self._consecutive_failures = 0
self.request_count += 1
except Exception:
# Retry also failed — fall through to the normal
# failure tracking below.
pass
if self._consecutive_failures <= 3 or self._consecutive_failures % MAX_CONSECUTIVE_FAILURES == 0:
logger.warning("Heartbeat failed (%d consecutive): %s", self._consecutive_failures, e)
if self._consecutive_failures >= MAX_CONSECUTIVE_FAILURES:
logger.info("Heartbeat: recreating HTTP client after %d failures", self._consecutive_failures)
try:
await client.aclose()
except Exception:
pass
break
# 2. Check delegation status
try:
await self._check_delegations(client)
except Exception as e:
logger.debug("Delegation check failed: %s", e)
await asyncio.sleep(HEARTBEAT_INTERVAL)
except asyncio.CancelledError:
raise
except Exception as e:
logger.error("Heartbeat loop error: %s — retrying in 30s", e)
await asyncio.sleep(HEARTBEAT_INTERVAL)
finally:
if client:
try:
await client.aclose()
except Exception:
pass
async def _check_delegations(self, client: httpx.AsyncClient):
"""Check for completed delegations and store results for the agent."""
try:
resp = await client.get(
f"{self.platform_url}/workspaces/{self.workspace_id}/delegations",
headers=auth_headers(),
)
if resp.status_code != 200:
return
delegations = resp.json()
if not isinstance(delegations, list):
return
new_results = []
for d in delegations:
did = d.get("delegation_id", "")
status = d.get("status", "")
if not did or did in self._seen_delegation_ids:
continue
if status in ("completed", "failed"):
# Fix B (Cycle 5): validate source_id before accepting delegation
# results. Only process delegations that THIS workspace created
# (source_id == self.workspace_id). Attacker-crafted delegation
# records with a foreign source_id cannot inject instructions.
source_id = d.get("source_id", "")
if source_id != self.workspace_id:
logger.warning(
"Heartbeat: skipping delegation %s — source_id %r does not "
"match this workspace %r; possible injection attempt",
did, source_id, self.workspace_id,
)
self._seen_delegation_ids.add(did) # mark seen so we don't warn again
continue
self._seen_delegation_ids.add(did)
new_results.append({
"delegation_id": did,
"target_id": d.get("target_id", ""),
"source_id": source_id,
"status": status,
"summary": d.get("summary", ""),
"response_preview": d.get("response_preview", ""),
"error": d.get("error", ""),
"timestamp": time.time(),
})
# Evict old seen IDs if over limit
if len(self._seen_delegation_ids) > MAX_SEEN_DELEGATION_IDS:
# Keep most recent half
self._seen_delegation_ids = set(list(self._seen_delegation_ids)[MAX_SEEN_DELEGATION_IDS // 2:])
if new_results:
# Append to results file for context injection on next message
with open(DELEGATION_RESULTS_FILE, "a") as f:
for r in new_results:
f.write(json.dumps(r) + "\n")
logger.info("Heartbeat: %d new delegation results — triggering self-message", len(new_results))
# Build a summary message for the agent.
# Fix B (Cycle 5): do NOT embed raw response_preview text in
# user-role A2A messages — that is the prompt-injection vector.
# Instead reference only the delegation ID and status; the agent
# reads full content from DELEGATION_RESULTS_FILE which was
# written above from trusted platform data.
summary_lines = []
for r in new_results:
line = f"- [{r['status']}] Delegation {r['delegation_id'][:8]}: {r['summary'][:80]}"
if r.get("error"):
line += f"\n Error: {r['error'][:100]}"
summary_lines.append(line)
# Look up parent workspace (cached after first call)
if self._parent_name is None:
try:
parent_resp = await client.get(
f"{self.platform_url}/workspaces/{self.workspace_id}",
headers=auth_headers(),
)
if parent_resp.status_code == 200:
parent_id = parent_resp.json().get("parent_id", "")
if parent_id:
parent_info = await client.get(
f"{self.platform_url}/workspaces/{parent_id}",
headers=auth_headers(),
)
if parent_info.status_code == 200:
self._parent_name = parent_info.json().get("name", "")
if self._parent_name is None:
self._parent_name = "" # No parent — cache empty
except Exception:
pass # Will retry next cycle
parent_name = self._parent_name or ""
report_instruction = ""
if parent_name:
report_instruction = (
f"\n\nIMPORTANT: Report these results back to your parent '{parent_name}' "
f"by delegating a summary to them. Use delegate_task or delegate_task_async "
f"with a concise status report. Also use send_message_to_user to notify the user."
)
else:
report_instruction = (
"\n\nReport results using send_message_to_user to notify the user."
)
trigger_msg = (
"Delegation results are ready. Review them and take appropriate action:\n"
+ "\n".join(summary_lines)
+ report_instruction
)
# Send A2A self-message to wake the agent.
# Minimum 60s between self-messages to avoid spam, but always send
# when there are genuinely NEW results to process.
now = time.time()
if now - self._last_self_message_time < SELF_MESSAGE_COOLDOWN:
logger.debug("Heartbeat: self-message cooldown (60s), will retry next cycle")
else:
self._last_self_message_time = now
try:
# self_source_headers() adds X-Workspace-ID so the
# platform tags this row source=agent, not canvas
# — see platform_auth.py for the full rationale.
await client.post(
f"{self.platform_url}/workspaces/{self.workspace_id}/a2a",
json={
"method": "message/send",
"params": {
"message": {
"role": "user",
"parts": [{"type": "text", "text": trigger_msg}],
},
},
},
headers=self_source_headers(self.workspace_id),
timeout=120.0,
)
logger.info("Heartbeat: self-message sent to process delegation results")
except Exception as e:
logger.warning("Heartbeat: failed to send self-message: %s", e)
# Also push notification to user via canvas
for r in new_results:
try:
msg = f"Delegation {r['status']}: {r['summary'][:100]}"
if r.get("response_preview"):
msg += f"\nResult: {r['response_preview'][:200]}"
await client.post(
f"{self.platform_url}/workspaces/{self.workspace_id}/notify",
json={"message": msg, "type": "delegation_result"},
headers=auth_headers(),
)
except Exception:
pass
except Exception as e:
logger.debug("Delegation check error: %s", e)
View Git Blame Copy Permalink