aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/core/runtime/node_context.py

@@ -0,0 +1,203 @@
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Any
+
+from aethergraph.contracts.services.llm import LLMClientProtocol
+from aethergraph.core.runtime.runtime_services import get_ext_context_service
+from aethergraph.services.channel.session import ChannelSession
+from aethergraph.services.continuations.continuation import Continuation
+from aethergraph.services.llm.providers import Provider
+from aethergraph.services.memory.facade import MemoryFacade
+
+from .base_service import _ServiceHandle
+from .bound_memory import BoundMemoryAdapter
+from .node_services import NodeServices
+
+
+@dataclass
+class NodeContext:
+    run_id: str
+    graph_id: str
+    node_id: str
+    services: NodeServices
+    resume_payload: dict[str, Any] | None = None
+    bound_memory: BoundMemoryAdapter | None = None  # back-compat
+
+    # --- accessors (compatible names) ---
+    def runtime(self) -> NodeServices:
+        return self.services
+
+    def logger(self):
+        return self.services.logger.for_node_ctx(
+            run_id=self.run_id, node_id=self.node_id, graph_id=self.graph_id
+        )
+
+    def channel(self, channel_key: str | None = None):
+        return ChannelSession(self, channel_key)
+
+    # New way: prefer memory_facade directly
+    def memory(self) -> MemoryFacade:
+        if not self.services.memory_facade:
+            raise RuntimeError("MemoryFacade not bound")
+        return self.services.memory_facade
+
+    # Back-compat: old ctx.mem() -> To be deprecated
+    def mem(self) -> BoundMemoryAdapter:
+        if not self.bound_memory:
+            raise RuntimeError("BoundMemory adapter not available")
+        return self.bound_memory
+
+    # Artifacts / index
+    def artifacts(self):
+        return self.services.artifact_store
+
+    def kv(self):
+        if not self.services.kv:
+            raise RuntimeError("KV not available")
+        return self.services.kv
+
+    def llm(
+        self,
+        profile: str = "default",
+        *,
+        provider: Provider | None = None,
+        model: str | None = None,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        azure_deployment: str | None = None,
+        timeout: float | None = None,
+    ) -> LLMClientProtocol:
+        """
+        Get an LLM client by profile.
+        - If no overrides are provided, just return existing profile.
+        - If overrides are provided, create/update that profile at runtime.
+        """
+        svc = self.services.llm
+
+        if (
+            provider is None
+            and model is None
+            and base_url is None
+            and api_key is None
+            and azure_deployment is None
+            and timeout is None
+        ):
+            return svc.get(profile)
+
+        return svc.configure_profile(
+            profile=profile,
+            provider=provider,
+            model=model,
+            base_url=base_url,
+            api_key=api_key,
+            azure_deployment=azure_deployment,
+            timeout=timeout,
+        )
+
+    def llm_set_key(self, provider: str, model: str, api_key: str, profile: str = "default"):
+        """
+        Quickly configure or override the provider/key for a profile.
+        """
+        svc = self.services.llm
+        svc.set_key(provider=provider, model=model, api_key=api_key, profile=profile)
+
+    def rag(self):
+        if not self.services.rag:
+            raise RuntimeError("RAGService not available")
+        return self.services.rag
+
+    def mcp(self, name):
+        if not self.services.mcp:
+            raise RuntimeError("MCPService not available")
+        return self.services.mcp.get(name)
+
+    def continuations(self):
+        return self.services.continuation_store
+
+    def prepare_wait_for_resume(self, token: str):
+        # creates and registers a Future for this token without awaiting
+        return self.services.wait_registry.register(token)
+
+    def clock(self):
+        if not self.services.clock:
+            raise RuntimeError("Clock service not available")
+        return self.services.clock
+
+    def svc(self, name: str) -> Any:
+        # generic accessor for external context services
+        raw = get_ext_context_service(name)
+        if raw is None:
+            raise KeyError(f"Service '{name}' not registered")
+        # bind the service to the context
+        bind = getattr(raw, "bind", None)
+        if callable(bind):
+            return raw.bind(context=self)
+        return raw
+
+    def __getattr__(self, name: str) -> Any:
+        # Try to resolve as an external context service
+        try:
+            bound = self.svc(name)
+        except KeyError:
+            # Fall back to normal attribute error for anything else
+            raise AttributeError(f"NodeContext has no attribute '{name}'") from None
+        # Return a callable handle that behaves like the bound service
+        return _ServiceHandle(name, bound)
+
+    def _now(self):
+        if self.services.clock:
+            return self.services.clock.now()
+        else:
+            from datetime import datetime
+
+            return datetime.utcnow()
+
+    # ---- continuation helpers ----
+    async def create_continuation(
+        self,
+        *,
+        kind: str,
+        payload: dict | None,
+        channel: str | None,
+        deadline_s: int | None = None,
+        poll: dict | None = None,
+        attempts: int = 0,
+    ) -> Continuation:
+        """Create and store a continuation for this node in the continuation store."""
+        token = await self.services.continuation_store.mint_token(
+            self.run_id, self.node_id, attempts=attempts
+        )
+        deadline = None
+        if deadline_s:
+            deadline = self._now() + timedelta(seconds=deadline_s)
+
+        continuation = Continuation(
+            run_id=self.run_id,
+            node_id=self.node_id,
+            kind=kind,
+            token=token,
+            prompt=payload.get("prompt") if payload else None,
+            resume_schema=payload.get("resume_schema") if payload else None,
+            channel=channel,
+            deadline=deadline,
+            poll=poll,
+            next_wakeup_at=deadline,
+            created_at=self._now(),
+            attempts=attempts,
+            payload=payload,
+        )
+        await self.services.continuation_store.save(continuation)
+        return continuation
+
+    async def wait_for_resume(self, token: str) -> dict:
+        """Wait for a continuation to be resumed, and return the payload.
+        This will register the wait in the wait registry, and suspend until resumed.
+        Useful for nodes that need to pause and wait for short-term external events.
+        For long-term waits, use DualStage Tools instead.
+        """
+        waits = self.services.wait_registry
+        if not waits:
+            raise RuntimeError("WaitRegistry missing on context/runtime")
+        fut = waits.register(token)
+        payload = await fut
+        return payload

aethergraph/core/runtime/node_services.py

@@ -0,0 +1,30 @@
+from dataclasses import dataclass
+from typing import Any
+
+from aethergraph.contracts.services.llm import LLMClientProtocol
+from aethergraph.services.channel.channel_bus import ChannelBus
+from aethergraph.services.clock.clock import SystemClock
+from aethergraph.services.continuations.stores.fs_store import FSContinuationStore
+from aethergraph.services.logger.std import StdLoggerService
+from aethergraph.services.mcp.service import MCPService
+from aethergraph.services.memory.facade import MemoryFacade
+from aethergraph.services.rag.facade import RAGFacade
+from aethergraph.services.waits.wait_registry import WaitRegistry
+
+
+@dataclass
+class NodeServices:
+    channels: ChannelBus
+    continuation_store: FSContinuationStore
+    artifact_store: Any  # e.g., ArtifactFacadeAsync
+    wait_registry: WaitRegistry | None = None
+    clock: SystemClock | None = None
+    logger: StdLoggerService | None = (
+        None  # StdLoggerService.for_node_ctx() will be used in NodeContext
+    )
+    kv: Any | None = None
+    memory: Any | None = None  # MemoryFactory (for cross-session needs)
+    memory_facade: MemoryFacade | None = None  # bound memory for this node
+    llm: LLMClientProtocol | None = None  # LLMService
+    rag: RAGFacade | None = None  # RAGService
+    mcp: MCPService | None = None  # MCPService

aethergraph/core/runtime/recovery.py

@@ -0,0 +1,159 @@
+# aethergraph/runtime/recovery.py
+from __future__ import annotations
+
+import datetime
+import hashlib
+import time
+from typing import Any
+
+from aethergraph.contracts.services.state_stores import GraphStateStore
+
+from ..graph.node_state import NodeStatus
+from ..graph.task_graph import TaskGraph, TaskGraphSpec
+
+
+def hash_spec(spec: TaskGraphSpec) -> str:
+    import json
+
+    # stable hash of the immutable parts
+    raw = json.dumps(
+        {
+            "graph_id": spec.graph_id,
+            "version": spec.version,
+            "nodes": {
+                nid: {
+                    "type": ns.type,
+                    "dependencies": ns.dependencies,
+                    "logic": ns.logic if isinstance(ns.logic, str) else str(ns.logic),
+                    "metadata": ns.metadata,
+                }
+                for nid, ns in spec.nodes.items()
+            },
+            "io": {
+                "required": sorted(list(spec.io.required.keys())),
+                "optional": sorted(list(spec.io.optional.keys())),
+                "outputs": sorted(list(spec.io.outputs.keys())),
+            },
+        },
+        sort_keys=True,
+    )
+    return hashlib.sha256(raw.encode("utf-8")).hexdigest()
+
+
+async def recover_graph_run(
+    *,
+    spec: TaskGraphSpec,
+    run_id: str,
+    store: GraphStateStore,
+) -> TaskGraph:
+    snap = await store.load_latest_snapshot(run_id)
+    g = TaskGraph.from_spec(spec=spec, state=None)
+    g.state.run_id = run_id
+    # If no snapshot, we're starting fresh.
+    if not snap:
+        return g
+
+    # Basic drift guard (optional: warn if different)
+    want = hash_spec(spec)
+    if snap.spec_hash != want:
+        # Soft warning; TODO: raise if later want strictness.
+        import logging
+
+        logger = logging.getLogger("aethergraph.core.runtime.recovery")
+        logger.warning(
+            f"[recover_graph_run] Spec hash mismatch for run {run_id}: snapshot has {snap.spec_hash[:8]}..., want {want[:8]}... This typically means the graph definition changed since the snapshot was taken. It is not a problem if you created the graph differently on resume."
+        )
+
+    # Apply snapshot state
+    _hydrate_state_from_json(g, snap.state)
+
+    return g
+
+
+def _hydrate_state_from_json(graph, j: dict[str, Any]) -> None:
+    graph.state.rev = j.get("rev", 0)
+    graph.state._bound_inputs = j.get("_bound_inputs")
+    for nid, ns_json in j.get("nodes", {}).items():
+        ns = graph.state.nodes.setdefault(nid, graph.state.nodes.get(nid))
+        status_name = ns_json.get("status", "PENDING")
+        status = getattr(NodeStatus, status_name, NodeStatus.PENDING)
+        if status == NodeStatus.RUNNING:
+            status = NodeStatus.PENDING
+        ns.status = status
+
+        outs = ns_json.get("outputs") or {}
+        # Keep as-is; resume_policy already blocked non-JSON/ref earlier
+        ns.outputs = outs
+
+
+async def rearm_waits_if_needed(graph, env, *, ttl_s: int = 3600):
+    store = env.container.cont_store
+    bus = env.container.channels
+    now = time.time()
+
+    for nid, ns in graph.state.nodes.items():
+        if getattr(ns, "status", None) not in (
+            NodeStatus.WAITING_HUMAN,
+            getattr(NodeStatus, "WAITING_EXTERNAL", "WAITING_EXTERNAL"),
+        ):
+            continue
+
+        cont = await store.get(run_id=env.run_id, node_id=nid)
+        # Normalize deadline to a numeric timestamp to avoid comparing datetime with float
+        deadline = getattr(cont, "deadline", None)
+        deadline_ts = deadline.timestamp() if isinstance(deadline, datetime.datetime) else deadline
+        expired = (not cont) or (deadline_ts is not None and deadline_ts < now)
+
+        if not expired:
+            continue  # still valid
+
+        # Rebuild OutEvent from saved wait_spec
+        ws = getattr(ns, "wait_spec", None)
+        if not ws:
+            # No spec → safest fallback is to keep waiting but log it
+            env.container.logger.for_run().warning(
+                f"[rearm] missing wait_spec for {env.run_id}:{nid}; staying WAITING"
+            )
+            continue
+
+        # Mint a new continuation token
+        new_deadline = now + ttl_s
+        token = store.mint(
+            run_id=env.run_id,
+            node_id=nid,
+            kind=ws["kind"],
+            channel=ws.get("channel"),
+            deadline=new_deadline,
+            meta=ws.get("meta") or {},
+        )
+        # Build + send OutEvent
+        out = {
+            "type": "session.need_input"
+            if ws["kind"] == "text"
+            else "session.need_approval"
+            if ws["kind"] == "approval"
+            else "session.need_input",  # default
+            "channel": ws.get("channel"),
+            "text": ws.get("prompt"),
+            "buttons": [{"label": o} for o in (ws.get("options") or [])],
+            "meta": ws.get("meta") or {},
+        }
+        payload = await bus.send(out)  # may inline-resume for console/web
+
+        # If adapter returned a payload immediately → deliver inline
+        if payload and "payload" in payload:
+            # inline path (same as in _enter_wait)
+            await env.container.resume_bus.deliver_inline(
+                run_id=env.run_id, node_id=nid, payload=payload["payload"]
+            )
+        else:
+            # Persist (replace/insert) the new continuation
+            store.save_for_node(
+                run_id=env.run_id,
+                node_id=nid,
+                token=token,
+                kind=ws["kind"],
+                channel=ws.get("channel"),
+                deadline=new_deadline,
+                meta=ws.get("meta") or {},
+            )

aethergraph/core/runtime/run_registration.py

@@ -0,0 +1,33 @@
+from contextlib import AbstractContextManager
+
+
+class RunRegistrationGuard(AbstractContextManager):
+    """Context manager to register and unregister a scheduler for a run. Primarily for resume handling.
+    On enter, registers the scheduler with the container's scheduler registry.
+    On exit, unregisters the scheduler.
+    """
+
+    def __init__(self, *, run_id: str, scheduler, container):
+        self.run_id = run_id
+        self.scheduler = scheduler
+        self.container = container
+        self._did_reg = False
+
+    def __enter__(self):
+        reg = self.container.sched_registry
+        existing = reg.get(self.run_id)
+        if existing is not None and existing is not self.scheduler:
+            # Be explicit to avoid silent clobbering
+            raise RuntimeError(f"Scheduler already registered for run_id={self.run_id}")
+        reg.register(self.run_id, self.scheduler)
+        self._did_reg = True
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        if self._did_reg:
+            try:
+                self.container.sched_registry.unregister(self.run_id)
+            finally:
+                self._did_reg = False
+        # Return False to propagate any exception (important so callers can detect failures)
+        return False

aethergraph/core/runtime/runtime_env.py

@@ -0,0 +1,157 @@
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+# ---- artifact services ----
+from aethergraph.services.artifacts.fs_store import FSArtifactStore  # AsyncArtifactStore
+from aethergraph.services.artifacts.jsonl_index import JsonlArtifactIndex  # AsyncArtifactIndex
+
+# ---- channel services ----
+from aethergraph.services.channel.channel_bus import ChannelBus
+from aethergraph.services.clock.clock import SystemClock
+from aethergraph.services.container.default_container import DefaultContainer, get_container
+from aethergraph.services.continuations.stores.fs_store import (
+    FSContinuationStore,  # AsyncContinuationStore
+)
+
+# ---- memory services ----
+from aethergraph.services.memory.facade import MemoryFacade
+from aethergraph.services.resume.router import ResumeRouter
+from aethergraph.services.waits.wait_registry import WaitRegistry
+
+from ..graph.task_node import TaskNodeRuntime
+from .bound_memory import BoundMemoryAdapter
+from .execution_context import ExecutionContext
+from .node_services import NodeServices
+
+
+@dataclass
+class RuntimeEnv:
+    """Unified runtime env that is built from DefaultContainer and can spawn NodeContexts."""
+
+    run_id: str
+    graph_id: str | None = None
+    graph_inputs: dict[str, Any] = field(default_factory=dict)
+    outputs_by_node: dict[str, dict[str, Any]] = field(default_factory=dict)
+
+    # container (DI)
+    container: DefaultContainer = field(default_factory=get_container)
+
+    # optional predicate to skip execution
+    should_run_fn: Callable[[], bool] | None = None
+
+    # --- convenience projections of commonly used services ---
+    @property
+    def schedulers(self) -> dict[str, Any]:
+        return self.container.schedulers
+
+    @property
+    def registry(self):
+        return self.container.registry
+
+    @property
+    def logger_factory(self):
+        return self.container.logger
+
+    @property
+    def clock(self) -> SystemClock:
+        return self.container.clock
+
+    @property
+    def channels(self) -> ChannelBus:
+        return self.container.channels
+
+    @property
+    def continuation_store(self) -> FSContinuationStore:
+        return self.container.cont_store
+
+    @property
+    def wait_registry(self) -> WaitRegistry:
+        return self.container.wait_registry
+
+    @property
+    def artifacts(self) -> FSArtifactStore:
+        return self.container.artifacts
+
+    @property
+    def artifact_index(self) -> JsonlArtifactIndex:
+        return self.container.artifact_index
+
+    @property
+    def memory_factory(self):
+        return self.container.memory_factory
+
+    @property
+    def llm_service(self):
+        return self.container.llm
+
+    @property
+    def rag_facade(self):
+        return self.container.rag
+
+    @property
+    def mcp_service(self):
+        return self.container.mcp
+
+    @property
+    def resume_router(self) -> ResumeRouter:
+        return self.container.resume_router
+
+    def make_ctx(
+        self, *, node: "TaskNodeRuntime", resume_payload: dict[str, Any] | None = None
+    ) -> Any:
+        defaults = {
+            "run_id": self.run_id,
+            "graph_id": self.graph_id,
+            "node_id": node.node_id,
+            "agent_id": getattr(node, "tool_name", None),
+            "tags": [],
+            "entities": [],
+        }
+        mem: MemoryFacade = self.memory_factory.for_session(
+            run_id=self.run_id,
+            graph_id=self.graph_id,
+            node_id=node.node_id,
+            agent_id=defaults["agent_id"],
+        )
+
+        from aethergraph.services.artifacts.facade import ArtifactFacade
+
+        artifact_facade = ArtifactFacade(
+            run_id=self.run_id,
+            graph_id=self.graph_id or "",
+            node_id=node.node_id,
+            tool_name=node.tool_name,
+            tool_version=node.tool_version,  # to be filled from node if available
+            store=self.artifacts,
+            index=self.artifact_index,
+        )
+
+        services = NodeServices(
+            channels=self.channels,
+            continuation_store=self.continuation_store,
+            artifact_store=artifact_facade,
+            wait_registry=self.wait_registry,
+            clock=self.clock,
+            logger=self.logger_factory,
+            kv=self.container.kv_hot,  # keep using hot kv for ephemeral
+            memory=self.memory_factory,  # factory (for other sessions if needed)
+            memory_facade=mem,  # bound memory for this run/node
+            llm=self.llm_service,  # LLMService
+            rag=self.rag_facade,  # RAGService
+            mcp=self.mcp_service,  # MCPService
+        )
+        return ExecutionContext(
+            run_id=self.run_id,
+            graph_id=self.graph_id,
+            graph_inputs=self.graph_inputs,
+            outputs_by_node=self.outputs_by_node,
+            services=services,
+            logger_factory=self.logger_factory,
+            clock=self.clock,
+            resume_payload=resume_payload,
+            should_run_fn=self.should_run_fn,
+            # Back-compat shim for old ctx.mem()
+            bound_memory=BoundMemoryAdapter(mem, defaults),
+            resume_router=self.resume_router,
+        )

aethergraph/core/runtime/runtime_registry.py

@@ -0,0 +1,32 @@
+from contextvars import ContextVar
+
+from aethergraph.services.registry.unified_registry import UnifiedRegistry
+
+__singleton_registry: UnifiedRegistry = UnifiedRegistry()
+_current_registry: ContextVar[UnifiedRegistry | None] = ContextVar("ag_registry", default=None)
+
+
+def set_current_registry(reg: UnifiedRegistry):
+    """Set the current registry in contextvar."""
+    _current_registry.set(reg)
+
+
+def current_registry() -> UnifiedRegistry:
+    """Get the current registry from contextvar, or raise if not set."""
+    # first try if services has a registry set
+    from .runtime_services import current_services
+
+    svc = None
+    try:
+        # get current services and registry from there
+        svc = current_services()
+        if hasattr(svc, "registry") and svc.registry is not None:
+            return svc.registry
+    except Exception:
+        pass
+
+    # otherwise use contextvar
+    reg = _current_registry.get()
+    if reg is None:
+        return __singleton_registry  # fallback to singleton if not set in local context
+    return reg