aethergraph 0.1.0a1__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

aethergraph/core/runtime/base_service.py

@@ -23,10 +23,34 @@ __all__ = [
 
 class _ServiceHandle:
     """
-    A callable, transparent handle around a bound service.
-    - Attribute access delegates to the underlying service.
-    - Calling with no args returns the service (ergonomic parity with built-ins).
-    - Calling with args forwards to service.__call__ if present.
+    This class provides ergonomic access to a service instance, allowing attribute access,
+    direct retrieval, and forwarding of calls if the underlying service is callable.
+    Examples:
+        Accessing service attributes:
+        ```python
+        handle = _ServiceHandle("my_service", service_instance)
+        value = handle.some_attribute
+        ```
+        Retrieving the service instance:
+        ```python
+        svc = handle()
+        ```
+        Forwarding calls to a callable service:
+        ```python
+        result = handle(arg1, arg2)
+        ```
+    Args:
+        name: The name of the service for identification and error reporting.
+        bound_service: The actual service instance to be wrapped.
+    Attribute Access:
+        All attribute lookups are delegated to the underlying service instance.
+    Calling:
+        - If called with no arguments, returns the underlying service instance.
+        - If called with arguments and the service is callable, forwards the call.
+        - Raises TypeError if the service is not callable and called with arguments.
+    Returns:
+        The result of the underlying service's __call__ method, or the service instance itself
+        if called with no arguments.
     """
 
     __slots__ = ("_name", "_svc")

aethergraph/core/runtime/execution_context.py

@@ -7,6 +7,9 @@ from datetime import datetime
 import importlib
 from typing import TYPE_CHECKING, Any
 
+from aethergraph.api.v1.deps import RequestIdentity
+from aethergraph.services.scope.scope import Scope
+
 if TYPE_CHECKING:
     from aethergraph.core.graph.task_node import TaskNodeRuntime
 
@@ -24,6 +27,10 @@ from .node_services import NodeServices
 class ExecutionContext:
     run_id: str
     graph_id: str | None
+    session_id: str | None
+    agent_id: str | None
+    app_id: str | None
+    identity: RequestIdentity | None
     graph_inputs: dict[str, Any]
     outputs_by_node: dict[str, dict[str, Any]]
     services: NodeServices
@@ -32,7 +39,7 @@ class ExecutionContext:
     resume_payload: dict[str, Any] | None = None
     should_run_fn: Callable[[], bool] | None = None
     resume_router: ResumeRouter | None = None  # ResumeRouter
-
+    scope: Scope | None = None  # Node Scope
     # Back-compat shim
     bound_memory: BoundMemoryAdapter | None = None
 
@@ -40,27 +47,18 @@ class ExecutionContext:
         return NodeContext(
             run_id=self.run_id,
             graph_id=self.graph_id or "",
+            session_id=self.session_id,
             node_id=node.node_id,
             services=self.services,
+            identity=self.identity,
             resume_payload=self.resume_payload,
+            scope=self.scope,
+            agent_id=self.agent_id,
+            app_id=self.app_id,
             # back-compat for old ctx.mem()
             bound_memory=self.bound_memory,
         )
 
-    # def as_node_context(self, ad) -> "NodeContext":
-    #     """ Create a NodeContext representing this execution context itself as a node.
-    #         Useful for ad-hoc contexts that don't have real nodes.
-    #     """
-    #     return NodeContext(
-    #         run_id=self.run_id,
-    #         graph_id=self.graph_id or "",
-    #         node_id="ad_",
-    #         services=self.services,
-    #         resume_payload=self.resume_payload,
-    #         # back-compat for old ctx.mem()
-    #         bound_memory=self.bound_memory,
-    #     )
-
     # ----- helpers used by step forward() -----
     def now(self) -> datetime:
         return self.clock.now()

aethergraph/core/runtime/graph_runner.py

@@ -5,10 +5,12 @@ import threading
 from typing import Any
 import uuid
 
+from aethergraph.api.v1.deps import RequestIdentity
 from aethergraph.contracts.errors.errors import GraphHasPendingWaits
 from aethergraph.contracts.services.state_stores import GraphSnapshot
+from aethergraph.core.graph.task_graph import TaskGraph
 from aethergraph.core.runtime.recovery import hash_spec, recover_graph_run
-from aethergraph.services.container.default_container import build_default_container
+from aethergraph.services.container.default_container import build_default_container  # adjust path
 from aethergraph.services.state_stores.graph_observer import PersistenceObserver
 from aethergraph.services.state_stores.resume_policy import (
     assert_snapshot_json_only,
@@ -20,6 +22,7 @@ from ..execution.retry_policy import RetryPolicy
 from ..graph.graph_fn import GraphFunction
 from ..graph.graph_refs import resolve_any as _resolve_any
 from ..runtime.runtime_env import RuntimeEnv
+from ..runtime.runtime_metering import current_meter_context
 from ..runtime.runtime_services import ensure_services_installed
 from .run_registration import RunRegistrationGuard
 
@@ -51,7 +54,7 @@ async def _attach_persistence(graph, env, spec, snapshot_every=1) -> Persistence
 
 
 async def _build_env(
-    owner, inputs: dict[str, Any], **rt_overrides
+    owner, inputs: dict[str, Any], identity: RequestIdentity | None = None, **rt_overrides
 ) -> tuple[RuntimeEnv, RetryPolicy, int]:
     container = _get_container()
     # apply optional overrides onto the container instance
@@ -60,12 +63,36 @@ async def _build_env(
             setattr(container, k, v)
 
     run_id = rt_overrides.get("run_id") or f"run-{uuid.uuid4().hex[:8]}"
+    session_id = rt_overrides.get("session_id")
+    graph_id = getattr(owner, "graph_id", None) or getattr(owner, "name", None)
+
+    # Prefer runtime overrides (from RunRecord) over static graph metadata -- UI provenance
+    agent_id = (
+        rt_overrides.get("agent_id")
+        or getattr(owner, "agent_id", None)
+        or getattr(getattr(owner, "spec", None) or {}, "agent_id", None)
+        or (getattr(owner, "spec", None) or {}).get("agent_id")  # for dict-like spec
+    )
+
+    app_id = (
+        rt_overrides.get("app_id")
+        or getattr(owner, "app_id", None)
+        or getattr(getattr(owner, "spec", None) or {}, "app_id", None)
+        or (getattr(owner, "spec", None) or {}).get("app_id")
+    )
+
     env = RuntimeEnv(
         run_id=run_id,
+        graph_id=graph_id,
+        session_id=session_id,
+        identity=identity,
         graph_inputs=inputs,
         outputs_by_node={},
         container=container,
+        agent_id=agent_id,
+        app_id=app_id,
     )
+
     retry = rt_overrides.get("retry") or RetryPolicy()
     max_conc = rt_overrides.get("max_concurrency", getattr(owner, "max_concurrency", 4))
     return env, retry, max_conc
@@ -139,7 +166,7 @@ def _resolve_graph_outputs(
             continuations=continuations,
         ) from e
 
-    return next(iter(result.values())) if len(result) == 1 else result
+    return result
 
 
 def _resolve_graph_outputs_or_waits(graph, inputs, env, *, raise_on_waits: bool = True):
@@ -187,6 +214,7 @@ async def load_latest_snapshot_json(store, run_id: str) -> dict[str, Any] | None
     # JsonGraphStateStore serializes GraphSnapshot via snap.__dict__
     # load_latest_snapshot already returns a GraphSnapshot(**jsondict).
     # Convert back to plain JSON-ish dict:
+
     return {
         "run_id": snap.run_id,
         "graph_id": snap.graph_id,
@@ -196,22 +224,124 @@ async def load_latest_snapshot_json(store, run_id: str) -> dict[str, Any] | None
     }
 
 
+def _register_metering_context(
+    env: RuntimeEnv, target: GraphFunction | TaskGraph | Any
+) -> dict[str, Any]:
+    """
+    Build a metering context dict from the RuntimeEnv.
+    """
+    run_id = env.run_id
+    # derive graph_id from target if possible - GraphFunction.name or TaskGraph.spec.graph_id
+    graph_id = getattr(target, "name", None) or getattr(
+        getattr(target, "spec", None), "graph_id", None
+    )
+
+    # user id etc through auth context if available
+    user_id = getattr(getattr(env, "identity", None), "user_id", None)
+    org_id = getattr(getattr(env, "identity", None), "org_id", None)
+
+    token = current_meter_context.set(
+        {
+            "run_id": run_id,
+            "graph_id": graph_id,
+            "user_id": user_id,
+            "org_id": org_id,
+        }
+    )
+    return token
+
+
 # ---------- public API ----------
-async def run_async(target, inputs: dict[str, Any] | None = None, **rt_overrides):
+async def run_async(
+    target,
+    inputs: dict[str, Any] | None = None,
+    identity: RequestIdentity | None = None,
+    **rt_overrides,
+):
     """
-    Generic async runner for TaskGraph or GraphFunction.
-    - GraphFunction → delegates to gf.run(env=..., **inputs)
-    - TaskGraph/builder → schedules and resolves graph-level outputs
+    Execute a TaskGraph or GraphFunction asynchronously with optional persistence and resumability.
+
+    This method handles environment setup, cold-resume from persisted state (if available),
+    input validation, scheduling, and output resolution. It supports both fresh runs and
+    resuming incomplete runs, automatically wiring up persistence observers and enforcing
+    snapshot policies.
+
+    Examples:
+        Running a graph function:
+        ```python
+        result = await run_async(my_graph_fn, {"x": 1, "y": 2})
+        ```
+
+        Running a TaskGraph with custom run ID and identity:
+        ```python
+        result = await run_async(
+            my_task_graph,
+            {"input": 42},
+            run_id="custom-run-123",
+            identity=my_identity,     # Only used with API requests. Ignored when running locally.
+            max_concurrency=8
+        )
+        ```
+
+    Args:
+        target: The TaskGraph, GraphFunction, or builder to execute.
+        inputs: Dictionary of input values for the graph.
+        identity: Optional RequestIdentity for user/session context.
+        **rt_overrides: Optional runtime overrides for environment and execution. Recognized runtime overrides include:
+
+            - run_id (str): Custom run identifier.
+            - session_id (str): Session identifier for grouping runs.
+            - agent_id (str): Agent identifier for provenance.
+            - app_id (str): Application identifier for provenance.
+            - retry (RetryPolicy): Custom retry policy.
+            - max_concurrency (int): Maximum number of concurrent tasks.
+            - Any additional container attributes supported by your environment.
+
+    Returns:
+        dict: The resolved outputs of the graph, or a status dict if waiting on continuations.
+
+    Raises:
+        GraphHasPendingWaits: If the graph is waiting on external events and outputs are not ready.
+        TypeError: If the target is not a valid TaskGraph or GraphFunction.
+
+    Notes:
+        - Speficially for GraphFunctions, you can directly use `await graph_fn(**inputs)` without needing `run_async`.
+        - `graph_fn` is not resumable; use TaskGraphs for persistence and recovery features.
+        - when using `graph` for persistence/resumability, ensure your outputs are JSON-serializable, for examples:
+            - primitive types (str, int, float, bool, None)
+            - lists/dicts of primitive types
+
+            - graph that can be resumed with JSON-serializable outputs:
+            ```python
+            @graphify(...)
+            def my_graph(...):
+                ...
+                return {"result": 42, "data": [1, 2, 3], "info": None} # valid JSON-serializable output
+            ```
+            - graph that cannot be resumed due to non-JSON-serializable outputs:
+            ```python
+            @graphify(...)
+            def my_graph(...):
+                ...
+                return {"chekpoint": torch.pt, "file": open("data.bin", "rb")} # invalid outputs for resuming (but valid for fresh runs)
+            ```
+            - Despite this, you can still use `graph` without persistence features; just avoid resuming such graphs.
+
     """
     inputs = inputs or {}
     # GraphFunction path
     if isinstance(target, GraphFunction):
-        env, retry, max_conc = await _build_env(target, inputs, **rt_overrides)
-        return await target.run(env=env, max_concurrency=max_conc, **inputs)
+        env, retry, max_conc = await _build_env(target, inputs, identity=identity, **rt_overrides)
+        token = _register_metering_context(env, target)  # set metering context
+        try:
+            return await target.run(env=env, max_concurrency=max_conc, **inputs)
+        finally:
+            # reset metering context
+            current_meter_context.reset(token)
 
     # TaskGraph path
     graph = _materialize_task_graph(target)
-    env, retry, max_conc = await _build_env(graph, inputs, **rt_overrides)
+    env, retry, max_conc = await _build_env(graph, inputs, identity=identity, **rt_overrides)
 
     # Extract spec for run/recovery ...
     spec = getattr(graph, "spec", None) or getattr(graph, "get_spec", lambda: None)()
@@ -266,33 +396,44 @@ async def run_async(target, inputs: dict[str, Any] | None = None, **rt_overrides
     )
 
     # Register for resumes and run
-    with RunRegistrationGuard(run_id=env.run_id, scheduler=sched, container=env.container):
-        try:
-            await sched.run()
-        except asyncio.CancelledError:
-            raise
-        finally:
-            # FINAL SNAPSHOT on normal or cancelled exit (if store exists)
-            if store and obs:
-                artifacts = getattr(env.container, "artifacts", None)
-                snap = await snapshot_from_graph(
-                    run_id=graph.state.run_id or env.run_id,
-                    graph_id=graph.graph_id,
-                    rev=graph.state.rev,
-                    spec_hash=hash_spec(spec),
-                    state_obj=graph.state,
-                    artifacts=artifacts,
-                    allow_externalize=False,  # FIXME: artifact writer async loop error; set False to *avoid* writing artifacts during snapshot
-                    include_wait_spec=True,
-                )
-                await store.save_snapshot(snap)
-
-    # Resolve graph-level outputs (will raise GraphHasPendingWaits if waits)
-    return _resolve_graph_outputs_or_waits(graph, inputs, env, raise_on_waits=True)
+    token = _register_metering_context(env, target)  # set metering context
+    try:
+        with RunRegistrationGuard(run_id=env.run_id, scheduler=sched, container=env.container):
+            try:
+                await sched.run()
+            except asyncio.CancelledError:
+                raise
+            finally:
+                # FINAL SNAPSHOT on normal or cancelled exit (if store exists)
+                if store and obs:
+                    artifacts = getattr(env.container, "artifacts", None)
+                    snap = await snapshot_from_graph(
+                        run_id=graph.state.run_id or env.run_id,
+                        graph_id=graph.graph_id,
+                        rev=graph.state.rev,
+                        spec_hash=hash_spec(spec),
+                        state_obj=graph.state,
+                        artifacts=artifacts,
+                        allow_externalize=False,  # FIXME: artifact writer async loop error; set False to *avoid* writing artifacts during snapshot
+                        include_wait_spec=True,
+                    )
+                    await store.save_snapshot(snap)
+
+        # Resolve graph-level outputs (will raise  if waits)
+        return _resolve_graph_outputs_or_waits(graph, inputs, env, raise_on_waits=True)
+    finally:
+        # reset metering context
+        current_meter_context.reset(token)
 
 
 async def run_or_resume_async(
-    target, inputs: dict[str, Any], *, run_id: str | None = None, **rt_overrides
+    target,
+    inputs: dict[str, Any],
+    *,
+    run_id: str | None = None,
+    session_id: str | None = None,
+    identity: RequestIdentity | None = None,
+    **rt_overrides,
 ):
     """
     If state exists for run_id → cold resume, else fresh run.
@@ -300,7 +441,9 @@ async def run_or_resume_async(
     """
     if run_id is not None:
         rt_overrides = dict(rt_overrides or {}, run_id=run_id)
-    return await run_async(target, inputs, **rt_overrides)
+    if session_id is not None:
+        rt_overrides = dict(rt_overrides or {}, session_id=session_id)
+    return await run_async(target, inputs, identity=identity, **rt_overrides)
 
 
 # sync adapter (optional, safe in notebooks/servers)
@@ -344,6 +487,48 @@ class _LoopThread:
 _LOOP = _LoopThread()
 
 
-def run(target, inputs: dict[str, Any] | None = None, **rt_overrides):
+def run(
+    target,
+    inputs: dict[str, Any] | None = None,
+    identity: RequestIdentity | None = None,
+    **rt_overrides,
+):
+    """
+    Execute a target graph node synchronously with the provided inputs.
+
+    This function submits the execution of a target node (or graph) to the event loop,
+    allowing for asynchronous execution while providing a synchronous interface.
+    Runtime configuration overrides can be supplied as keyword arguments.
+
+    Examples:
+        Running a graph node with default inputs:
+        ```python
+        future = run(my_node)
+        result = future.result()
+        ```
+
+        Running with custom inputs and runtime overrides:
+        ```python
+        future = run(my_node, inputs={"x": 42}, timeout=10)
+        output = future.result()
+        ```
+
+    Args:
+        target: The graph node or callable to execute.
+        inputs: Optional dictionary of input values to pass to the node.
+        identity: Optional RequestIdentity for user/session context.
+        **rt_overrides: Additional keyword arguments to override runtime configuration.
+
+    Returns:
+        concurrent.futures.Future: A future representing the asynchronous execution of the node.
+
+    Notes:
+        - This function is suitable for use in synchronous contexts where asynchronous execution is desired.
+        - It is recommended to use asynchronous execution directly when possible for better performance and responsiveness.
+
+    Warnings:
+        - KeyboardInterrupt handling may not be perfect; consider using an async main function when possible.
+        - This function blocks the calling thread until the execution is complete.
+    """
     inputs = inputs or {}
-    return _LOOP.submit(run_async(target, inputs, **rt_overrides))
+    return _LOOP.submit(run_async(target, inputs, identity=identity, **rt_overrides))