diff --git a/pmoves/services/agent-zero/main.py b/pmoves/services/agent-zero/main.py index 1d33b81e7a..c95dc0c36a 100644 --- a/pmoves/services/agent-zero/main.py +++ b/pmoves/services/agent-zero/main.py @@ -627,7 +627,7 @@ async def lifespan(app: FastAPI): for sig in (signal.SIGTERM, signal.SIGINT): try: loop.add_signal_handler( - sig, lambda s=sig: asyncio.create_task(process_manager.stop()) + sig, lambda: asyncio.create_task(process_manager.stop()) ) except (NotImplementedError, RuntimeError, ValueError): # Tests run event loops in worker threads where signal handlers are unsupported. @@ -643,8 +643,10 @@ async def lifespan(app: FastAPI): await _controller_task _controller_task = None if event_controller.is_started: - with contextlib.suppress(Exception): + try: await event_controller.stop() + except Exception: + logger.debug("Exception during event controller shutdown", exc_info=True) app = FastAPI(title="Agent Zero Supervisor", lifespan=lifespan) diff --git a/pmoves/services/agent_zero/controller.py b/pmoves/services/agent_zero/controller.py index fed382667b..c1fcd95c8a 100644 --- a/pmoves/services/agent_zero/controller.py +++ b/pmoves/services/agent_zero/controller.py @@ -671,8 +671,9 @@ async def _fallback_to_core_nats(self) -> None: for sub in self._subscriptions: try: await sub.unsubscribe() - except Exception: - pass + except Exception as exc: + sub_info = getattr(sub, "subject", sub) + logger.warning("Failed to unsubscribe from %s: %s", sub_info, exc) self._subscriptions.clear() # Disable JS for this controller and create core subscriptions self.settings.use_jetstream = False diff --git a/pmoves/services/comfy-watcher/watcher.py b/pmoves/services/comfy-watcher/watcher.py index 385f565a30..52b04affac 100644 --- a/pmoves/services/comfy-watcher/watcher.py +++ b/pmoves/services/comfy-watcher/watcher.py @@ -1,9 +1,12 @@ import asyncio import datetime +from datetime import timedelta import hashlib import json +import logging import os import pathlib +import shutil import time from minio import Minio @@ -11,34 +14,107 @@ from services.common.events import envelope +logger = logging.getLogger("comfy-watcher") + + +def _parse_int_env(key: str, default: int) -> int: + """Parse integer from environment with validation. + + Args: + key: Environment variable name. + default: Default value if missing or invalid. + + Returns: + Parsed integer value or default. + """ + value = os.environ.get(key) + if value is None: + return default + try: + return int(value) + except ValueError: + logger.warning(f"Invalid {key}={value!r}, using default={default}") + return default + MINIO_ENDPOINT = os.environ.get("MINIO_ENDPOINT", "minio:9000") MINIO_USE_SSL = os.environ.get("MINIO_USE_SSL", "false").lower() == "true" MINIO_ACCESS_KEY = os.environ.get("MINIO_ACCESS_KEY") MINIO_SECRET_KEY = os.environ.get("MINIO_SECRET_KEY") BUCKET = os.environ.get("MINIO_BUCKET", "pmoves-comfyui") PUBLIC_BASE_URL = os.environ.get("PUBLIC_BASE_URL", "http://minio:9000") -PRESIGN_HOURS = int(os.environ.get("PRESIGN_EXPIRES_HOURS", "24")) +PRESIGN_HOURS = _parse_int_env("PRESIGN_EXPIRES_HOURS", 24) OUTPUT_DIR = os.environ.get("COMFY_OUTPUT_DIR", "/data/output") STATE_PATH = os.environ.get("COMFY_WATCHER_STATE_PATH", "/state/state.json") NATS_URL = os.environ.get("NATS_URL", "nats://nats:4222") -POLL_SECONDS = int(os.environ.get("COMFY_WATCHER_POLL_SECONDS", "5")) +POLL_SECONDS = _parse_int_env("COMFY_WATCHER_POLL_SECONDS", 5) def load_state() -> dict: + """Load state from disk with proper error handling. + + Returns the uploaded file registry. Handles missing state files, + corrupted JSON, and other I/O errors gracefully with logging. + """ try: with open(STATE_PATH) as f: return json.load(f) - except Exception: + except FileNotFoundError: + logger.info(f"State file not found, starting fresh: {STATE_PATH}") + return {"uploaded": {}} + except json.JSONDecodeError as exc: + logger.error(f"Corrupted state file at {STATE_PATH}, backing up and starting fresh: {exc}") + # Backup corrupted file for investigation + backup_path = f"{STATE_PATH}.corrupted.{int(time.time())}" + try: + shutil.copy(STATE_PATH, backup_path) + logger.info(f"Backed up corrupted state to: {backup_path}") + except Exception as backup_exc: + logger.warning(f"Could not backup corrupted state file: {backup_exc}") + return {"uploaded": {}} + except Exception as exc: + logger.error(f"Unexpected error loading state from {STATE_PATH}: {exc}") return {"uploaded": {}} def save_state(state: dict) -> None: + """Save the watcher state to disk. + + Persists the current state dictionary to JSON storage, ensuring the + target directory exists. Creates parent directories as needed. + + Args: + state: A dictionary containing the watcher state, typically with + an "uploaded" key mapping file hashes to their metadata + (key, timestamp, size). + + Raises: + OSError: If the state file cannot be written due to permission + or disk space issues. + TypeError: If the state dictionary contains non-serializable objects. + """ pathlib.Path(os.path.dirname(STATE_PATH)).mkdir(parents=True, exist_ok=True) with open(STATE_PATH, "w") as f: json.dump(state, f) def file_hash(path: str) -> str: + """Calculate the SHA-256 hash of a file. + + Computes a cryptographic hash by reading the file in 1 MB chunks, + allowing efficient hashing of large files without loading them + entirely into memory. + + Args: + path: The absolute or relative path to the file to hash. + + Returns: + A hexadecimal string representing the SHA-256 hash of the + file's contents. + + Raises: + FileNotFoundError: If the specified file does not exist. + PermissionError: If the file cannot be read due to permissions. + """ h = hashlib.sha256() with open(path, "rb") as f: while True: @@ -50,37 +126,72 @@ def file_hash(path: str) -> str: async def run() -> None: + """Main watcher loop for ComfyUI output files. + + Continuously monitors the configured output directory for new files, + uploads them to MinIO object storage, and publishes NATS events + with artifact metadata and presigned URLs. + + The function: + - Validates MinIO credentials are configured + - Loads the upload state to avoid re-uploading files + - Ensures the target MinIO bucket exists + - Polls the output directory at configured intervals + - Uploads new files with timestamped S3-style keys + - Generates presigned URLs for temporary access + - Publishes completion events to NATS for downstream processing + + Raises: + RuntimeError: If MinIO credentials are not set or bucket + creation fails. + + The loop runs indefinitely until the process is terminated. + """ if not MINIO_ACCESS_KEY or not MINIO_SECRET_KEY: raise RuntimeError("MINIO_ACCESS_KEY and MINIO_SECRET_KEY must be set") state = load_state() - client = Minio(MINIO_ENDPOINT, access_key=MINIO_ACCESS_KEY, secret_key=MINIO_SECRET_KEY, secure=MINIO_USE_SSL) - + client = Minio( + MINIO_ENDPOINT, + access_key=MINIO_ACCESS_KEY, + secret_key=MINIO_SECRET_KEY, + secure=MINIO_USE_SSL, + ) nc = NATS() await nc.connect(servers=[NATS_URL]) + # Ensure MinIO bucket exists with proper error handling try: if not client.bucket_exists(BUCKET): client.make_bucket(BUCKET) - except Exception as e: - print("Bucket check/make:", e) + logger.info(f"Created MinIO bucket: {BUCKET}") + except Exception as exc: + logger.error(f"Failed to verify/create MinIO bucket '{BUCKET}': {exc}") + raise RuntimeError(f"MinIO bucket setup failed for '{BUCKET}': {exc}") from exc while True: for root, _, files in os.walk(OUTPUT_DIR): - for fn in files: - path = os.path.join(root, fn) - try: - size = os.path.getsize(path) - if size == 0: - continue + for fname in files: + if not fname.endswith((".png", ".jpg", ".jpeg", ".webp", ".gif")): + continue + path = os.path.join(root, fname) + h = file_hash(path) + + if h in state.get("uploaded", {}): + continue - h = file_hash(path) - if state["uploaded"].get(h): - continue + size_bytes = os.path.getsize(path) + timestamp = int(time.time()) + key = f"comfyui/{timestamp}/{fname}" - key = datetime.datetime.now(datetime.timezone.utc).strftime("comfyui/%Y/%m/%d/") + fn - client.fput_object(BUCKET, key, path, content_type="image/png") - state["uploaded"][h] = {"key": key, "ts": time.time(), "size": size} + try: + with open(path, "rb") as f: + client.put_object(BUCKET, key, f, length=size_bytes) + state.setdefault("uploaded", {})[h] = { + "key": key, + "timestamp": timestamp, + "size": size_bytes, + } save_state(state) public_url = ( @@ -88,18 +199,20 @@ async def run() -> None: ) payload = {"artifact_uri": f"s3://{BUCKET}/{key}", "meta": {"public_url": public_url}} - try: - from datetime import timedelta + # Generate presigned URL with error handling + try: payload["meta"]["presigned_url"] = client.presigned_get_object(BUCKET, key, expires=timedelta(hours=PRESIGN_HOURS)) - except Exception: - pass + except Exception as presign_exc: + logger.warning(f"Failed to generate presigned URL for {key}: {presign_exc}") + # Continue without presigned URL - not critical for all workflows env = envelope("gen.image.result.v1", payload, source="comfy-watcher") await nc.publish("gen.image.result.v1", json.dumps(env).encode()) - print("Uploaded and announced:", key) - except Exception as e: - print("Error processing", path, e) + logger.info(f"Uploaded and announced: {key}") + except Exception as exc: + logger.error(f"Error processing file {path}: {exc}", exc_info=True) + # Continue to next file instead of aborting entire loop await asyncio.sleep(POLL_SECONDS) diff --git a/pmoves/services/gateway-agent/app.py b/pmoves/services/gateway-agent/app.py index d1ea2b8e79..285a13d990 100644 --- a/pmoves/services/gateway-agent/app.py +++ b/pmoves/services/gateway-agent/app.py @@ -113,8 +113,9 @@ async def lifespan(app: FastAPI): app = FastAPI( title="PMOVES Gateway Agent", description="Orchestrates 100+ MCP tools with Cipher memory integration", - version="1.0.0" -, lifespan=lifespan) + version="1.0.0", + lifespan=lifespan, +) # ============================================================================ diff --git a/pmoves/services/hi-rag-gateway-v2/app.py b/pmoves/services/hi-rag-gateway-v2/app.py index 2ce2156b0b..74bee3feb3 100644 --- a/pmoves/services/hi-rag-gateway-v2/app.py +++ b/pmoves/services/hi-rag-gateway-v2/app.py @@ -30,6 +30,11 @@ _RERANK_CONFIG_ERRORS: List[str] = [] _RERANK_CONFIG_WARNINGS: List[str] = [] +# Configure logging early for use in helper functions +logging.basicConfig(level=logging.INFO) +logging.getLogger().setLevel(logging.INFO) +logger = logging.getLogger("hirag.gateway.v2") + def _parse_bool(name: str, default: bool) -> bool: raw = os.environ.get(name) @@ -61,6 +66,17 @@ def _parse_positive_int(name: str, default: int) -> int: return value +def _parse_int_env(name: str, default: int) -> int: + """Parse integer environment variable with validation and fallback.""" + raw = os.environ.get(name, str(default)) + try: + return int(raw) + except (TypeError, ValueError) as exc: + # Use logger for general config warnings, not rerank-specific list + logger.warning("%s value %r invalid, using default %s: %s", name, raw, default, exc) + return default + + def _parse_optional_bool(name: str) -> Optional[bool]: raw = os.environ.get(name) if raw is None or str(raw).strip() == "": @@ -167,10 +183,10 @@ def get_rerank_status() -> Dict[str, Any]: NEO4J_URL = (os.environ.get("NEO4J_URL","bolt://neo4j:7687") or "").strip() NEO4J_USER = os.environ.get("NEO4J_USER","neo4j") NEO4J_PASSWORD = os.environ.get("NEO4J_PASSWORD","neo4j") -NEO4J_DICT_REFRESH_SEC = int(os.environ.get("NEO4J_DICT_REFRESH_SEC","60")) -NEO4J_DICT_LIMIT = int(os.environ.get("NEO4J_DICT_LIMIT","50000")) -ENTITY_CACHE_TTL = int(os.environ.get("ENTITY_CACHE_TTL","60")) -ENTITY_CACHE_MAX = int(os.environ.get("ENTITY_CACHE_MAX","1000")) +NEO4J_DICT_REFRESH_SEC = _parse_int_env("NEO4J_DICT_REFRESH_SEC", 60) +NEO4J_DICT_LIMIT = _parse_int_env("NEO4J_DICT_LIMIT", 50000) +ENTITY_CACHE_TTL = _parse_int_env("ENTITY_CACHE_TTL", 60) +ENTITY_CACHE_MAX = _parse_int_env("ENTITY_CACHE_MAX", 1000) SUPABASE_REST_URL = os.environ.get("SUPA_REST_URL") or os.environ.get("SUPABASE_REST_URL") # Optional internal REST base explicitly pointing at host-gateway; prefer for derivations @@ -203,7 +219,7 @@ def _is_jwt(key: str) -> bool: SUPABASE_ANON_KEY if _is_jwt(SUPABASE_ANON_KEY) else _jwt_service_key or _explicit_realtime_key or SUPABASE_ANON_KEY ) -GEOMETRY_CACHE_WARM_LIMIT = int(os.environ.get("GEOMETRY_CACHE_WARM_LIMIT", "64")) +GEOMETRY_CACHE_WARM_LIMIT = _parse_int_env("GEOMETRY_CACHE_WARM_LIMIT", 64) # Realtime connection retry configuration GEOMETRY_REALTIME_BACKOFF = float(os.environ.get("GEOMETRY_REALTIME_BACKOFF", "5.0")) # Initial retry delay (seconds) GEOMETRY_REALTIME_MAX_BACKOFF = float(os.environ.get("GEOMETRY_REALTIME_MAX_BACKOFF", "60.0")) # Max retry delay cap (seconds) @@ -214,13 +230,9 @@ def _is_jwt(key: str) -> bool: TAILSCALE_CIDRS = [c.strip() for c in os.environ.get("TAILSCALE_CIDRS","100.64.0.0/10").split(",") if c.strip()] TRUSTED_PROXY_SOURCES = [c.strip() for c in os.environ.get("HIRAG_TRUSTED_PROXIES", "").split(",") if c.strip()] -HTTP_PORT = int(os.environ.get("HIRAG_HTTP_PORT","8086")) +HTTP_PORT = _parse_int_env("HIRAG_HTTP_PORT", 8086) NAMESPACE_DEFAULT = os.environ.get("INDEXER_NAMESPACE","pmoves") -logging.basicConfig(level=logging.INFO) -logging.getLogger().setLevel(logging.INFO) -logger = logging.getLogger("hirag.gateway.v2") - if RERANK_ENABLE: logger.info( "Rerank enabled with model=%s provider=%s topn=%s k=%s fusion=%s device=%s", @@ -1006,7 +1018,7 @@ async def _geometry_realtime_worker(ws_url: str, api_key: str) -> None: GAN_SIDECAR_ENABLED = os.environ.get("GAN_SIDECAR_ENABLED", "false").lower()=="true" PGHOST = os.environ.get("PGHOST") -PGPORT = int(os.environ.get("PGPORT","5432")) +PGPORT = _parse_int_env("PGPORT", 5432) PGUSER = os.environ.get("PGUSER") PGPASSWORD = os.environ.get("PGPASSWORD") PGDATABASE = os.environ.get("PGDATABASE") diff --git a/pmoves/services/pdf-ingest/app.py b/pmoves/services/pdf-ingest/app.py index b49c63699a..be8b7da066 100644 --- a/pmoves/services/pdf-ingest/app.py +++ b/pmoves/services/pdf-ingest/app.py @@ -127,7 +127,6 @@ async def lifespan(app: FastAPI): NATS_URL = os.environ.get("NATS_URL", "nats://nats:4222") _nc: Optional[NATS] = None -_nats_supervisor: Optional[asyncio.Task[None]] = None async def _nats_connect_loop() -> None: diff --git a/pmoves/services/pmoves-yt/yt.py b/pmoves/services/pmoves-yt/yt.py index a505f77ab1..776058dc9a 100644 --- a/pmoves/services/pmoves-yt/yt.py +++ b/pmoves/services/pmoves-yt/yt.py @@ -2418,69 +2418,6 @@ async def attempt_ingest() -> Dict[str, Any]: if job_id: _job_update(job_id, 'failed' if any_failures else 'completed', None if not any_failures else 'one or more items failed') return {'ok': not any_failures, 'job_id': job_id, 'count': len(results), 'results': results} - try: - concurrency = int(os.environ.get('YT_CONCURRENCY', str(YT_CONCURRENCY))) - except Exception: - concurrency = YT_CONCURRENCY - concurrency = max(1, concurrency) - - async def _ingest_entry(entry: Dict[str, Any]) -> Dict[str, Any]: - vid = entry['id'] - vid_url = f"https://www.youtube.com/watch?v={vid}" if len(vid) == 11 else vid - if job_id: - _item_upsert(job_id, vid, 'running', None, {'title': entry.get('title')}) - result: Dict[str, Any] - status = 'failed' - error_msg: Optional[str] = None - try: - result = await asyncio.to_thread(_ingest_one, vid_url, ns, bucket) - status = 'completed' if result.get('ok') else 'failed' - error_msg = result.get('error') - except Exception as exc: - error_msg = str(exc) - result = {'ok': False, 'error': error_msg} - if job_id: - _item_upsert(job_id, vid, status, error_msg) - return {'id': vid, **result} - - results: List[Dict[str, Any]] = [] - - if concurrency <= 1: - for entry in entries: - res = await _ingest_entry(entry) - results.append(res) - if rate_limit > 0: - await asyncio.sleep(rate_limit) - else: - semaphore = asyncio.Semaphore(concurrency) - rate_lock = asyncio.Lock() - next_available = 0.0 - results_buffer: List[Optional[Dict[str, Any]]] = [None] * len(entries) - - async def _wait_rate_limit() -> None: - nonlocal next_available - if rate_limit <= 0: - return - async with rate_lock: - now = time.monotonic() - if now < next_available: - await asyncio.sleep(next_available - now) - now = next_available - next_available = now + rate_limit - - async def _worker(idx: int, entry: Dict[str, Any]) -> None: - async with semaphore: - await _wait_rate_limit() - res = await _ingest_entry(entry) - results_buffer[idx] = res - - tasks = [asyncio.create_task(_worker(idx, entry)) for idx, entry in enumerate(entries)] - await asyncio.gather(*tasks) - results = [res for res in results_buffer if res is not None] - - if job_id: - _job_update(job_id, 'completed') - return {'ok': True, 'job_id': job_id, 'count': len(results), 'results': results} @app.post('/yt/channel') async def yt_channel(body: Dict[str,Any] = Body(...)): diff --git a/pmoves/services/publisher-discord/main.py b/pmoves/services/publisher-discord/main.py index 37cefc91c4..8177e48e20 100644 --- a/pmoves/services/publisher-discord/main.py +++ b/pmoves/services/publisher-discord/main.py @@ -1133,34 +1133,6 @@ async def _handle_claude_session_end(payload: Dict[str, Any]) -> None: logger.info("claude_session_ended", extra={"session_id": session_id, "end_reason": end_reason}) - global _nats_loop_task - # Validate critical environment configuration - if not DISCORD_WEBHOOK_URL: - logger.warning( - "discord_webhook_url_not_configured", - extra={"event": "discord_webhook_url_not_configured"}, - ) - else: - logger.info( - "discord_webhook_url_configured", - extra={"event": "discord_webhook_url_configured", "domain": _extract_webhook_domain(DISCORD_WEBHOOK_URL)}, - ) - if _nats_loop_task and _nats_loop_task.done(): - try: - _nats_loop_task.result() - except Exception as exc: # pragma: no cover - startup diagnostics - logger.warning( - "nats_loop_previous_failure", - extra={"event": "nats_loop_previous_failure", "error": str(exc)}, - ) - if _nats_loop_task is None or _nats_loop_task.done(): - logger.info( - "nats_loop_start", - extra={"event": "nats_loop_start", "servers": [NATS_URL]}, - ) - _nats_loop_task = asyncio.create_task(_nats_resilience_loop()) - - @app.post("/publish") async def publish_test(body: Dict[str, Any] = Body(...)): """Test endpoint for publishing messages to Discord webhook.""" diff --git a/pmoves/services/session-context-worker/main.py b/pmoves/services/session-context-worker/main.py index 2a8505c518..f4467d8aea 100644 --- a/pmoves/services/session-context-worker/main.py +++ b/pmoves/services/session-context-worker/main.py @@ -2,32 +2,8 @@ """ Session Context Worker - Transforms Claude Code session context to Hi-RAG knowledge base entries. -This module provides a NATS-based worker service that subscribes to Claude Code session -context events, transforms them into searchable knowledge base entries, and publishes -them to the Hi-RAG knowledge base for retrieval and analysis. - -NATS Subjects: - - Subscribes to: claude.code.session.context.v1 - - Publishes to: kb.upsert.request.v1 - -Environment Variables: - NATS_URL: NATS server connection URL (default: "nats://nats:4222") - HEALTH_PORT: Port for health check endpoint (default: 8100) - -Prometheus Metrics: - - session_context_worker_messages_received_total: Messages received by subject - - session_context_worker_messages_processed_total: Successfully processed messages - - session_context_worker_messages_failed_total: Failed message processing - - session_context_worker_kb_upserts_published_total: KB upsert requests published - - session_context_worker_processing_duration_seconds: Processing time histogram - -Example: - To run this service:: - - NATS_URL=nats://localhost:4222 python main.py - - This will start the FastAPI server with health and metrics endpoints, - connect to NATS, and begin processing session context messages. +Subscribes to: claude.code.session.context.v1 +Publishes to: kb.upsert.request.v1 """ import asyncio @@ -43,7 +19,10 @@ from fastapi import FastAPI from contextlib import asynccontextmanager from nats.aio.client import Client as NATS +from nats.aio.msg import Msg from prometheus_client import Counter, Histogram, generate_latest, CONTENT_TYPE_LATEST, REGISTRY +from jsonschema import validate, ValidationError as JsonSchemaValidationError +from services.common.events import load_schema # Prometheus metrics messages_received = Counter( @@ -79,9 +58,32 @@ ) logger = logging.getLogger("session_context_worker") +# Load schemas for payload validation (follows coding guidelines) +# "Validate payloads against schemas before publishing events using services/common/events.py" +try: + _SESSION_CONTEXT_SCHEMA = load_schema("claude.code.session.context.v1") + _KB_UPSERT_SCHEMA = load_schema("kb.upsert.request.v1") + logger.info("Loaded schemas: claude.code.session.context.v1, kb.upsert.request.v1") +except Exception as exc: + logger.error("Failed to load schemas: %s", exc, exc_info=True) + # Fallback: allow startup but validation will fail + _SESSION_CONTEXT_SCHEMA = None + _KB_UPSERT_SCHEMA = None + + +def _parse_int_env(env_var: str, default: int) -> int: + """Parse integer environment variable with validation and fallback.""" + value = os.environ.get(env_var, str(default)) + try: + return int(value) + except ValueError: + logger.warning("Invalid %s value %r, using default %s", env_var, value, default) + return default + + # Environment variables NATS_URL = os.environ.get("NATS_URL", "nats://nats:4222") -HEALTH_PORT = int(os.environ.get("HEALTH_PORT", "8100")) +HEALTH_PORT = _parse_int_env("HEALTH_PORT", 8100) SESSION_CONTEXT_SUBJECT = "claude.code.session.context.v1" KB_UPSERT_SUBJECT = "kb.upsert.request.v1" @@ -89,50 +91,46 @@ _nc: Optional[NATS] = None _nats_loop_task: Optional[asyncio.Task] = None + +def _nats_loop_done(task: asyncio.Task) -> None: + """Callback for NATS resilience loop task completion.""" + if not task.cancelled(): + exc = task.exception() + if exc: + logger.error("NATS resilience loop crashed unexpectedly: %s", exc, exc_info=True) + + # FastAPI app for health endpoint @asynccontextmanager async def lifespan(app: FastAPI): - """Manage application lifespan with NATS connection lifecycle. - - This async context manager handles the startup and shutdown of the NATS - connection resilience loop. On startup, it creates the NATS connection task. - On shutdown, it gracefully cancels the task and closes the connection. - - Args: - app (FastAPI): The FastAPI application instance. This parameter is required - by FastAPI's lifespan interface but is not used directly. - - Yields: - None: This context manager yields control back to FastAPI during the - application's lifetime. - - Notes: - - The NATS resilience loop runs in a background task and handles automatic - reconnection with exponential backoff. - - On shutdown, the task is cancelled and the NATS connection is closed - gracefully. - - Exceptions during shutdown are silently caught to ensure clean exit. - """ + """Manage application lifespan.""" global _nats_loop_task, _nc - # Startup: Start NATS connection loop + # Startup if _nats_loop_task is None or _nats_loop_task.done(): logger.info("Starting NATS resilience loop") _nats_loop_task = asyncio.create_task(_nats_resilience_loop()) + _nats_loop_task.add_done_callback(_nats_loop_done) yield - # Shutdown: Clean shutdown of NATS connection + # Shutdown + if _nats_loop_task: _nats_loop_task.cancel() try: await _nats_loop_task - except Exception: - pass + except asyncio.CancelledError: + logger.debug("NATS resilience loop cancelled successfully") + except Exception as e: + logger.warning("Unexpected error during NATS loop shutdown: %s", e) _nats_loop_task = None if _nc: try: await _nc.close() - except Exception: - pass + logger.info("NATS connection closed cleanly") + except asyncio.CancelledError: + raise + except Exception as e: + logger.warning("NATS close error during shutdown: %s", e) _nc = None @@ -145,33 +143,11 @@ async def lifespan(app: FastAPI): def _extract_searchable_content(context: Dict[str, Any]) -> str: - """Extract searchable text content from session context. - - This function combines summary, task descriptions, decisions, and other relevant - information from a Claude Code session context into a single searchable text block - suitable for indexing in the Hi-RAG knowledge base. - - Args: - context (Dict[str, Any]): The session context dictionary containing: - - summary (str, optional): Overall session summary. - - repository (str, optional): Git repository name. - - branch (str, optional): Git branch name. - - pending_tasks (List[Dict], optional): List of pending/completed tasks. - - decisions (List[Dict], optional): List of Q&A decisions. - - active_files (List[Dict], optional): List of active file contexts. - - tool_executions (List[Dict], optional): List of tool execution records. - - agent_spawns (List[Dict], optional): List of agent spawn events. - - Returns: - str: A formatted string containing all searchable content sections - separated by double newlines. Sections include summary, repository/branch, - tasks, decisions, active files, tool executions, and agent spawns. - - Notes: - - Task descriptions are prefixed with their status (e.g., "[completed]"). - - Only the first 10 active files are included to avoid excessive length. - - Only the first 5 tool executions are included to avoid excessive length. - - Empty sections are omitted from the output. + """ + Extract searchable text content from session context. + + Combines summary, task descriptions, decisions, and other relevant information + into a single searchable text block. """ parts = [] @@ -248,45 +224,10 @@ def _extract_searchable_content(context: Dict[str, Any]) -> str: def _build_metadata(context: Dict[str, Any]) -> Dict[str, Any]: - """Build metadata object for Hi-RAG knowledge base entry. - - Constructs a metadata dictionary for the knowledge base entry including session - tracking information, repository details, and context type classification. - This metadata enables efficient filtering and retrieval of session contexts. - - Args: - context (Dict[str, Any]): The session context dictionary containing: - - session_id (str, optional): Unique session identifier. - - context_type (str, optional): Type of context (e.g., "completion"). - - timestamp (str, optional): ISO format timestamp of the context. - - worktree (str, optional): Git worktree path. - - branch (str, optional): Git branch name. - - repository (str, optional): Git repository identifier. - - working_directory (str, optional): Current working directory path. - - parent_session_id (str, optional): Parent session ID if applicable. - - pending_tasks (List[Dict], optional): List of tasks for count metadata. - - active_files (List[Dict], optional): List of files for count metadata. - - decisions (List[Dict], optional): List of decisions for count metadata. - - Returns: - Dict[str, Any]: A metadata dictionary containing: - - source (str): Always "claude-code". - - session_id (str): Session identifier. - - context_type (str): Context type classification. - - timestamp (str): ISO format timestamp. - - worktree (str, optional): Git worktree path if present. - - branch (str, optional): Git branch name if present. - - repository (str, optional): Git repository identifier if present. - - working_directory (str, optional): Working directory path if present. - - parent_session_id (str, optional): Parent session ID if present. - - task_count (int, optional): Total number of tasks. - - completed_task_count (int, optional): Number of completed tasks. - - active_file_count (int, optional): Number of active files. - - decision_count (int, optional): Number of decisions made. - - Notes: - - Optional fields are only included if present in the input context. - - Count fields are only included if the corresponding lists are non-empty. + """ + Build metadata object for Hi-RAG knowledge base entry. + + Includes session tracking, repository info, and context type classification. """ metadata = { "source": "claude-code", @@ -332,34 +273,10 @@ def _build_metadata(context: Dict[str, Any]) -> Dict[str, Any]: def _transform_to_kb_upsert(context: Dict[str, Any]) -> Dict[str, Any]: - """Transform session context to kb.upsert.request.v1 payload. - - Creates a knowledge base upsert request that transforms a Claude Code session - context into a format suitable for the Hi-RAG knowledge base. The resulting - entry can be searched and retrieved for future reference. - - Args: - context (Dict[str, Any]): The session context dictionary containing all - session information including summary, tasks, decisions, files, etc. - - Returns: - Dict[str, Any]: A kb.upsert.request.v1 compatible payload with the structure: - - items (List[Dict]): List containing a single knowledge base item: - - id (str): Unique identifier generated from session_id, context_type, - and timestamp. - - text (str): Searchable text content extracted from the context. - - metadata (Dict[str, Any]): Metadata fields for filtering. - - namespace (str): Always "claude-code-sessions". - - meta (Dict[str, Any]): Processing metadata including: - - worker (str): Always "session-context-worker". - - version (str): Worker version string. - - processed_at (str): ISO format timestamp of processing. - - Notes: - - The KB entry ID is deterministic: "claude-session-{session_id}-{context_type}-{timestamp}". - - All entries are stored in the "claude-code-sessions" namespace. - - The processed_at timestamp reflects when the transformation occurred, not - the original session timestamp. + """ + Transform session context to kb.upsert.request.v1 payload. + + Creates a knowledge base entry that can be searched and retrieved later. """ session_id = context.get("session_id", "unknown") context_type = context.get("context_type", "unknown") @@ -394,33 +311,11 @@ def _transform_to_kb_upsert(context: Dict[str, Any]) -> Dict[str, Any]: return kb_upsert -async def _handle_session_context(msg: nats.aio.msg.Msg) -> None: - """Handle incoming session context messages from NATS. - - This is the message handler for the claude.code.session.context.v1 subject. - It parses the incoming session context, transforms it into a knowledge base - upsert request, and publishes it to the kb.upsert.request.v1 subject. - - Args: - msg (nats.aio.msg.Msg): The NATS message object containing: - - data (bytes): JSON-encoded session context data. - - subject (str): The NATS subject the message was published to. - - reply (str, optional): Optional reply subject for responses. - - Returns: - None: This function publishes results to NATS and updates Prometheus - metrics but does not return a value. - - Raises: - json.JSONDecodeError: If the message data cannot be decoded as JSON. - Exception: For any other processing errors (logged but not raised). +async def _handle_session_context(msg: Msg) -> None: + """ + Handle incoming session context messages. - Notes: - - Updates Prometheus metrics for messages received, processed, and failed. - - Tracks processing duration as a histogram metric. - - Logs warnings for invalid message formats and skips processing. - - Logs errors and continues processing other messages if transformation fails. - - Requires the global NATS client (_nc) to be connected for publishing. + Transforms the context and publishes to kb.upsert.request.v1. """ messages_received.labels(SESSION_CONTEXT_SUBJECT).inc() start_time = time.time() @@ -431,16 +326,31 @@ async def _handle_session_context(msg: nats.aio.msg.Msg) -> None: data = json.loads(msg.data.decode("utf-8")) if not isinstance(data, dict): - logger.warning(f"Invalid message format: expected dict, got {type(data)}") + logger.warning("Invalid message format: expected dict, got %s", type(data)) messages_failed.labels("invalid_format").inc() processing_duration.labels("unknown").observe(time.time() - start_time) return + # Validate incoming payload against schema (prevents schema drift) + if _SESSION_CONTEXT_SCHEMA: + try: + validate(instance=data, schema=_SESSION_CONTEXT_SCHEMA) + except JsonSchemaValidationError as exc: + logger.warning( + "Session context payload validation failed: %s", + exc.message, + extra={"session_id": data.get("session_id", "unknown"), "validation_error": exc.message} + ) + messages_failed.labels("schema_validation").inc() + processing_duration.labels("unknown").observe(time.time() - start_time) + return + session_id = data.get("session_id", "unknown") context_type = data.get("context_type", "unknown") logger.info( - f"Processing session context: session_id={session_id}, type={context_type}", + "Processing session context: session_id=%s, type=%s", + session_id, context_type, extra={ "session_id": session_id, "context_type": context_type, @@ -450,6 +360,20 @@ async def _handle_session_context(msg: nats.aio.msg.Msg) -> None: # Transform to kb.upsert format kb_upsert = _transform_to_kb_upsert(data) + # Validate outgoing payload against schema (prevents schema drift) + if _KB_UPSERT_SCHEMA: + try: + validate(instance=kb_upsert, schema=_KB_UPSERT_SCHEMA) + except JsonSchemaValidationError as exc: + logger.error( + "KB upsert payload validation failed: %s", + exc.message, + extra={"session_id": session_id, "validation_error": exc.message} + ) + messages_failed.labels("kb_schema_validation").inc() + processing_duration.labels(context_type).observe(time.time() - start_time) + return + # Publish to kb.upsert.request.v1 if _nc: await _nc.publish( @@ -458,7 +382,8 @@ async def _handle_session_context(msg: nats.aio.msg.Msg) -> None: ) kb_upserts_published.labels("claude-code-sessions").inc() logger.info( - f"Published KB upsert for session {session_id}", + "Published KB upsert for session %s", + session_id, extra={ "session_id": session_id, "kb_id": kb_upsert["items"][0]["id"], @@ -474,72 +399,38 @@ async def _handle_session_context(msg: nats.aio.msg.Msg) -> None: messages_processed.labels(context_type).inc() processing_duration.labels(context_type).observe(time.time() - start_time) - except json.JSONDecodeError as e: + except json.JSONDecodeError: logger.exception("Failed to decode JSON") messages_failed.labels("json_decode_error").inc() processing_duration.labels(context_type).observe(time.time() - start_time) - except Exception as e: - logger.error(f"Error processing session context: {e}", exc_info=True) + except Exception as exc: + logger.error("Error processing session context: %s", exc, exc_info=True) messages_failed.labels("processing_error").inc() processing_duration.labels(context_type).observe(time.time() - start_time) async def _register_nats_subscriptions(nc: NATS) -> None: - """Register NATS subscriptions for the session context worker. - - Subscribes to the claude.code.session.context.v1 subject to receive - session context events for processing and transformation into knowledge - base entries. - - Args: - nc (NATS): The connected NATS client instance to register subscriptions on. - - Returns: - None: This function registers subscriptions directly on the NATS client. - - Raises: - Exception: If subscription registration fails. Errors are logged with - details about the subject and exception. - - Notes: - - Only one subscription is registered: claude.code.session.context.v1. - - The message handler (_handle_session_context) is called for each message. - - Subscription failures are logged but do not crash the worker. - """ + """Register NATS subscriptions.""" try: await nc.subscribe(SESSION_CONTEXT_SUBJECT, cb=_handle_session_context) logger.info( - f"Subscribed to {SESSION_CONTEXT_SUBJECT}", + "Subscribed to %s", + SESSION_CONTEXT_SUBJECT, extra={"subject": SESSION_CONTEXT_SUBJECT} ) except Exception as exc: logger.error( - f"Failed to subscribe to {SESSION_CONTEXT_SUBJECT}: {exc}", + "Failed to subscribe to %s: %s", + SESSION_CONTEXT_SUBJECT, exc, exc_info=True ) async def _nats_resilience_loop() -> None: - """NATS connection resilience loop with automatic reconnection. - - This function runs in a background task and maintains a persistent connection - to the NATS server. It implements exponential backoff reconnection logic to - handle network failures and server restarts gracefully. - - Returns: - None: This function runs indefinitely until cancelled or raises - asyncio.CancelledError. - - Raises: - asyncio.CancelledError: When the task is cancelled during shutdown. - - Notes: - - Implements exponential backoff starting at 1 second, maxing out at 30 seconds. - - Backoff resets to 1 second after a successful connection. - - Registers subscriptions after each successful connection. - - Waits indefinitely on a disconnect event for the lifetime of the connection. - - Cleanly closes the connection on shutdown. - - Updates the global _nc variable with the active connection. + """ + NATS connection resilience loop with automatic reconnection. + + Maintains persistent NATS connection with exponential backoff on failures. """ global _nc backoff = 1.0 @@ -589,7 +480,7 @@ async def _closed_cb(): # Connection successful _nc = nc backoff = 1.0 - logger.info(f"NATS connected: {NATS_URL}", extra={"servers": [NATS_URL]}) + logger.info("NATS connected: %s", NATS_URL, extra={"servers": [NATS_URL]}) # Register subscriptions await _register_nats_subscriptions(nc) @@ -600,8 +491,8 @@ async def _closed_cb(): except asyncio.CancelledError: try: await nc.close() - except Exception: - pass + except Exception as e: + logger.debug("Error closing NATS connection during cancellation: %s", e) if _nc is nc: _nc = None raise @@ -609,52 +500,22 @@ async def _closed_cb(): # Clean up connection try: await nc.close() - except Exception: - pass + except Exception as e: + logger.debug("Error closing NATS connection during cleanup: %s", e) @app.get("/healthz") async def healthz(): - """Health check endpoint for monitoring service status. - - Provides a simple health check that indicates whether the service is running - and if the NATS connection is currently active. - - Returns: - Dict[str, Any]: A health status dictionary containing: - - ok (bool): Always True, indicating the service is running. - - nats_connected (bool): True if NATS client is connected, False otherwise. - - Notes: - - This endpoint is typically used by orchestrators (Kubernetes, Docker Compose) - for health checks. - - The NATS connection status reflects the current state and may change - as the resilience loop reconnects. - """ + """Health check endpoint.""" return { "ok": True, - "nats_connected": _nc is not None and not getattr(_nc, "_is_closed", True) + "nats_connected": _nc is not None, } @app.get("/metrics") async def metrics(): - """Prometheus-compatible metrics endpoint for observability. - - Exposes Prometheus metrics in the standard text format for scraping by - Prometheus or compatible monitoring systems. - - Returns: - Response: A FastAPI Response object containing: - - body (bytes): Prometheus metrics in text exposition format. - - media_type (str): CONTENT_TYPE_LATEST (text/plain; version=0.0.4). - - Notes: - - Metrics include counters for messages received/processed/failed and KB upserts. - - Processing duration histogram with context_type label. - - Typically scraped by Prometheus at /metrics every 15-60 seconds. - - The REGISTRY contains all prometheus_client metrics defined at module level. - """ + """Prometheus-compatible metrics endpoint.""" from fastapi.responses import Response return Response(generate_latest(REGISTRY), media_type=CONTENT_TYPE_LATEST) diff --git a/pmoves/services/tokenism-simulator/api/simulation.py b/pmoves/services/tokenism-simulator/api/simulation.py index 6228836d19..b1fc80252d 100644 --- a/pmoves/services/tokenism-simulator/api/simulation.py +++ b/pmoves/services/tokenism-simulator/api/simulation.py @@ -264,9 +264,10 @@ def _run_simulation_background( except Exception as e: # Acquire locks together for consistency with success path - with _status_lock, _results_lock: - _simulation_statuses[simulation_id] = "failed" + # CRITICAL: Must use same lock order as success path: _results_lock, _status_lock + with _results_lock, _status_lock: _simulation_results[simulation_id] = {"error": str(e)} + _simulation_statuses[simulation_id] = "failed" _evict_old_results() simulation_requests.labels( scenario=scenario.value, @@ -662,32 +663,34 @@ def get_simulation_status(simulation_id: str): For polling, implement exponential backoff (e.g., 1s, 2s, 4s, 8s...) to avoid overwhelming the service with frequent requests. """ - status = _simulation_statuses.get(simulation_id, 'unknown') + # CRITICAL: Lock order must match _run_simulation_background: _results_lock, _status_lock + with _results_lock, _status_lock: + status = _simulation_statuses.get(simulation_id, 'unknown') - if status == 'complete': - result = _simulation_results.get(simulation_id, {}) - return jsonify({ - 'simulation_id': simulation_id, - 'status': status, - 'result': result, - }), 200 - elif status == 'failed': - result = _simulation_results.get(simulation_id, {}) - return jsonify({ - 'simulation_id': simulation_id, - 'status': status, - 'error': result.get('error', 'Unknown error'), - }), 500 - elif status == 'running': - return jsonify({ - 'simulation_id': simulation_id, - 'status': status, - 'message': 'Simulation still in progress', - }), 202 - else: - return jsonify({ - 'error': f'Simulation {simulation_id} not found', - }), 404 + if status == 'complete': + result = _simulation_results.get(simulation_id, {}) + return jsonify({ + 'simulation_id': simulation_id, + 'status': status, + 'result': result, + }), 200 + elif status == 'failed': + result = _simulation_results.get(simulation_id, {}) + return jsonify({ + 'simulation_id': simulation_id, + 'status': status, + 'error': result.get('error', 'Unknown error'), + }), 500 + elif status == 'running': + return jsonify({ + 'simulation_id': simulation_id, + 'status': status, + 'message': 'Simulation still in progress', + }), 202 + else: + return jsonify({ + 'error': f'Simulation {simulation_id} not found', + }), 404 @simulation_bp.route('/api/v1/scenarios', methods=['GET']) diff --git a/pmoves/services/tokenism-simulator/app.py b/pmoves/services/tokenism-simulator/app.py index d0955c9e83..cfd314f2e2 100644 --- a/pmoves/services/tokenism-simulator/app.py +++ b/pmoves/services/tokenism-simulator/app.py @@ -62,14 +62,11 @@ def create_app() -> Flask: app.config['SECRET_KEY'] = config.secret_key # Enable CORS with configurable allowed origins (strip whitespace) - allowed_origins = [ - origin.strip() - for origin in os.getenv( - 'ALLOWED_ORIGINS', - 'http://localhost:3000,http://localhost:8080,http://localhost:4000' - ).split(',') - if origin.strip() - ] + raw_origins = os.getenv( + 'ALLOWED_ORIGINS', + 'http://localhost:3000,http://localhost:8080,http://localhost:4000' + ).split(',') + allowed_origins = [origin for origin in (o.strip() for o in raw_origins) if origin] CORS(app, resources={r"/*": {"origins": allowed_origins}}) # Register blueprints diff --git a/pmoves/tools/consciousness_harvester.py b/pmoves/tools/consciousness_harvester.py index a11eb57fb0..70e0b0f49f 100644 --- a/pmoves/tools/consciousness_harvester.py +++ b/pmoves/tools/consciousness_harvester.py @@ -149,7 +149,7 @@ async def _handle_crawl_result(self, msg): links=data.get("links", []), title=data.get("title", ""), error=data.get("error"), - crawled_at=data.get("crawled_at", datetime.now(timezone.utc).isoformat() + "Z"), + crawled_at=data.get("crawled_at", datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")), ) self.pending_results[request_id].set_result(result) logger.debug(f"Received crawl result for {request_id}") @@ -227,7 +227,7 @@ async def crawl_url_http(self, url: str) -> CrawlResult: html=data.get("html", ""), links=data.get("links", []), title=data.get("title", ""), - crawled_at=datetime.now(timezone.utc).isoformat() + "Z", + crawled_at=datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), ) except httpx.HTTPError as e: @@ -301,7 +301,7 @@ def extract_theories(self, result: CrawlResult) -> List[TheoryEntry]: description=description[:1000], # Truncate long descriptions proponents=proponents, source_url=result.url, - extracted_at=datetime.now(timezone.utc).isoformat() + "Z", + extracted_at=datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), ) ) @@ -441,7 +441,7 @@ async def harvest_urls( # Save summary summary_path = output_dir / "harvest-summary.json" - results_summary["harvested_at"] = datetime.now(timezone.utc).isoformat() + "Z" + results_summary["harvested_at"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") summary_path.write_text( json.dumps(results_summary, indent=2), encoding="utf-8",