state-mesh 0.1.0__py3-none-any.whl

state_mesh/__init__.py

@@ -0,0 +1,26 @@
+from state_mesh.core.context import Context, Flag
+from state_mesh.core.step import Step, Branch, RetryConfig, step
+from state_mesh.core.pipeline import Pipeline, PipelineResult
+from state_mesh.guardrails.base import Guard, GuardResult
+from state_mesh.guardrails.schema import SchemaGuard
+from state_mesh.guardrails.content import PIIGuard
+from state_mesh.output.contract import OutputContract
+from state_mesh.mcp.server_manager import MCPServerConfig, MCPServerManager
+
+__all__ = [
+    "step",
+    "Step",
+    "Pipeline",
+    "Context",
+    "RetryConfig",
+    "Branch",
+    "PipelineResult",
+    "Guard",
+    "GuardResult",
+    "SchemaGuard",
+    "PIIGuard",
+    "OutputContract",
+    "Flag",
+    "MCPServerConfig",
+    "MCPServerManager",
+]

state_mesh/core/context.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Generic, TypeVar, Literal, TYPE_CHECKING
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    from bus import MCPBus
+
+StateSchema = TypeVar("StateSchema", bound=BaseModel)
+
+
+class Flag(BaseModel):
+    step_name:str 
+    severity: Literal["info","warn","error"]
+    timestamp:datetime=Field(default_factory=lambda:datetime.now(timezone.utc))
+    flag_type: str
+    reason:str
+    payload:dict[str,Any]=Field(default_factory=dict)
+
+class ContextMutationError(Exception):
+    pass
+    
+class Context(Generic[StateSchema]):
+    """
+    Context is the main object that is passed to all steps.
+    It contains the state, run_id, trace_id, pipeline_name and flags.
+    """
+    def __init__(self,state,run_id=None,trace_id=None,pipeline_name="unnamed",tools:MCPBus|None=None):
+        self._state = state
+        self._run_id = run_id or str(uuid.uuid4())
+        self._trace_id = trace_id or str(uuid.uuid4())
+        self._pipeline_name = pipeline_name
+        self._flags:list[Flag] = []
+        self._extras:dict[str,Any]={}
+        self._tools=tools
+        self._step_name=None
+       
+    @property
+    def state(self) -> StateSchema:
+        return self._state
+    
+    @property
+    def tools(self) -> MCPBus|None:
+        return self._tools
+        
+    @property
+    def run_id(self) -> str:
+        return self._run_id
+    
+    @property
+    def trace_id(self) -> str:
+        return self._trace_id
+    
+    @property
+    def pipeline_name(self) -> str:
+        return self._pipeline_name 
+
+    @property
+    def flags(self) -> list[Flag]:
+        flag=self._flags.copy() # sp the user wont be able directly modify this
+        return flag
+    
+    def set(self,key:str,value:Any)->None:
+        if key in self._extras:
+            raise ContextMutationError(f"Key {key} already exists in extras")
+        self._extras[key]=value
+    
+    def get(self,key:str,default:Any=None)->Any:
+        return self._extras.get(key,default)
+    
+    def emit_flag(self,flag_type:str,reason:str,severity:Literal["info","warn","error"]="warn",payload:dict[str,Any]|None=None)->None:
+        self._flags.append(Flag(
+            step_name=self._step_name,
+            severity=severity,
+            flag_type=flag_type,
+            reason=reason,
+            payload=payload or {}
+        ))
+        
+    def _set_current_step(self, step_name: str) -> None:
+        self._step_name = step_name
+
+    def _replace_state(self, new_state: StateSchema) -> None:
+        self._state = new_state
+    
+    #serialize the context for pipeline resumption
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "state": self._state.model_dump() if isinstance(self._state, BaseModel) else self._state,
+            "run_id": self._run_id,
+            "trace_id": self._trace_id,
+            "pipeline_name": self._pipeline_name,
+            "extras": self._extras,
+            "flags": [f.model_dump() for f in self._flags],
+            "last_completed_step": self._step_name,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any], state_schema: type[BaseModel]) -> Context:
+        ctx = cls(
+            state=state_schema(**data["state"]),
+            run_id=data["run_id"],
+            trace_id=data["trace_id"],
+            pipeline_name=data["pipeline_name"],
+        )
+        ctx._extras = data["extras"]
+        ctx._flags = [Flag(**f) for f in data["flags"]]
+        ctx._step_name = data.get("last_completed_step")
+        return ctx
+    
+    def __repr__(self) -> str:
+        return (
+        f"Context(run_id={self._run_id!r}, "
+        f"step={self._step_name!r}, "
+        f"flags={len(self._flags)})"
+    )
+        
+        
+
+        

state_mesh/core/human_input.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+
+class HumanInput:
+    def __init__(
+        self,
+        prompt: str,
+        response_key: str,
+        timeout_seconds: float = 300.0,
+        fallback_step: str | None = None,
+    ):
+        self.prompt = prompt
+        self.response_key = response_key
+        self.timeout_seconds = timeout_seconds
+        self.fallback_step = fallback_step

state_mesh/core/parallel.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import asyncio
+
+from state_mesh.core.context import Context
+from state_mesh.core.step import Step, StepResult
+
+
+class Parallel:
+    def __init__(self, steps: list[Step], name: str = "parallel"):
+        self.name = name
+        self.steps = steps
+
+    async def execute(self, ctx: Context) -> list[StepResult]:
+        results: list[StepResult] = list(
+            await asyncio.gather(*[step.execute(ctx) for step in self.steps], return_exceptions=True)
+        )
+
+        for step, result in zip(self.steps, results):
+            if not isinstance(result, BaseException) and result.status == "success":
+                ctx.set(step.name, result.output)
+
+        return results

state_mesh/core/pipeline.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+from state_mesh.core.context import Context, Flag
+from state_mesh.core.step import StepResult, Step, Branch
+from state_mesh.core.human_input import HumanInput
+from state_mesh.core.parallel import Parallel
+from state_mesh.observability.events import EventEmitter
+from state_mesh.observability.tracer import get_tracer
+from state_mesh.observability.flags import attach_flags_to_span
+from state_mesh.mcp.bus import MCPBus
+
+
+class PipelineResult(BaseModel):
+    output: Any
+    run_id: str
+    trace_id: str
+    duration_ms: float
+    status: Literal["success", "failed", "timed_out", "guarded"]
+    step_results: list[StepResult]
+    flags: list[Flag]
+
+
+class Pipeline:
+    def __init__(self, steps: list[Step], name: str = None, state_backend: Any = None, mcp_servers: list[str] | None = None):
+        self.name = name
+        self.steps = steps
+        self.state_backend = state_backend
+        self.mcp_servers = mcp_servers
+        self._step_map: dict[str, Step] = {s.name: s for s in steps}
+        self._emitter = EventEmitter()
+
+    async def run(self, ctx: Context, event_callback=None) -> PipelineResult:
+        step_results = []
+        start = time.perf_counter()
+
+        bus = None
+        if self.mcp_servers:
+            bus = MCPBus(self.mcp_servers)
+            await bus.start()
+            ctx._tools = bus
+
+        try:
+            tracer = get_tracer()
+            with tracer.start_as_current_span("pipeline.run") as pipeline_span:
+                pipeline_span.set_attribute("pipeline.name", self.name or "unnamed")
+                pipeline_span.set_attribute("run_id", ctx.run_id)
+                pipeline_span.set_attribute("trace_id", ctx.trace_id)
+
+                self._emitter.pipeline_started(ctx, self.name)
+                if event_callback:
+                    await event_callback("pipeline_started", {"pipeline_name": self.name, "run_id": ctx.run_id})
+
+                for step in self.steps:
+                    if isinstance(step, Parallel):
+                        parallel_results = await step.execute(ctx)
+                        step_results.extend(parallel_results)
+                        failed = next((r for r in parallel_results if isinstance(r, BaseException) or r.status != "success"), None)
+                        if failed:
+                            duration = (time.perf_counter() - start) * 1000
+                            status = failed.status if isinstance(failed, StepResult) else "failed"
+                            self._emitter.pipeline_finished(ctx, self.name, duration, status)
+                            if event_callback:
+                                await event_callback("pipeline_finished", {"run_id": ctx.run_id, "status": status, "duration_ms": duration})
+                            return PipelineResult(output=None, run_id=ctx.run_id, trace_id=ctx.trace_id, duration_ms=duration, status=status, step_results=step_results, flags=ctx.flags)
+                        continue
+
+                    self._emitter.step_started(ctx, step.name)
+                    if event_callback:
+                        await event_callback("step_started", {"step_name": step.name, "run_id": ctx.run_id})
+                    with tracer.start_as_current_span(f"step.{step.name}") as step_span:
+                        step_span.set_attribute("step.name", step.name)
+                        step_span.set_attribute("run_id", ctx.run_id)
+                        result = await step.execute(ctx)
+                        step_span.set_attribute("step.status", result.status)
+                        step_span.set_attribute("step.attempts", result.attempts)
+                        step_span.set_attribute("step.duration_ms", result.duration_ms)
+                        attach_flags_to_span(result.flags)
+
+                    self._emitter.step_finished(ctx, step.name, result.duration_ms, result.status)
+                    if event_callback:
+                        await event_callback("step_finished", {"step_name": step.name, "run_id": ctx.run_id, "status": result.status, "duration_ms": result.duration_ms})
+                    step_results.append(result)
+
+                    if result.status != "success":
+                        duration = (time.perf_counter() - start) * 1000
+                        self._emitter.pipeline_finished(ctx, self.name, duration, result.status)
+                        if event_callback:
+                            await event_callback("pipeline_finished", {"run_id": ctx.run_id, "status": result.status, "duration_ms": duration})
+                        return PipelineResult(output=result.output, run_id=ctx.run_id, trace_id=ctx.trace_id, duration_ms=duration, status=result.status, step_results=step_results, flags=ctx.flags)
+
+                    if isinstance(result.output, HumanInput):
+                        hi = result.output
+                        if event_callback:
+                            await event_callback("human_input_requested", {"run_id": ctx.run_id, "prompt": hi.prompt, "response_key": hi.response_key})
+                        from state_mesh.integrations.fastapi import _pending_inputs
+                        loop = asyncio.get_event_loop()
+                        future: asyncio.Future = loop.create_future()
+                        _pending_inputs[ctx.run_id] = future
+                        try:
+                            response = await asyncio.wait_for(asyncio.shield(future), timeout=hi.timeout_seconds)
+                            ctx.set(hi.response_key, response)
+                        except asyncio.TimeoutError:
+                            _pending_inputs.pop(ctx.run_id, None)
+                            if hi.fallback_step and hi.fallback_step in self._step_map:
+                                result = await self._step_map[hi.fallback_step].execute(ctx)
+                                step_results.append(result)
+                            duration = (time.perf_counter() - start) * 1000
+                            self._emitter.pipeline_finished(ctx, self.name, duration, "timed_out")
+                            if event_callback:
+                                await event_callback("pipeline_finished", {"run_id": ctx.run_id, "status": "timed_out", "duration_ms": duration})
+                            return PipelineResult(output=None, run_id=ctx.run_id, trace_id=ctx.trace_id, duration_ms=duration, status="timed_out", step_results=step_results, flags=ctx.flags)
+                        finally:
+                            _pending_inputs.pop(ctx.run_id, None)
+                        continue
+
+                    if isinstance(result.output, Branch):
+                        branch = result.output
+                        ctx = branch.context
+                        if branch.to not in self._step_map:
+                            raise ValueError(f"Branch target to {branch.to} not found in pipeline steps")
+                        self._emitter.step_started(ctx, branch.to)
+                        if event_callback:
+                            await event_callback("step_started", {"step_name": branch.to, "run_id": ctx.run_id})
+                        with tracer.start_as_current_span(f"step.{branch.to}") as step_span:
+                            step_span.set_attribute("step.name", branch.to)
+                            step_span.set_attribute("run_id", ctx.run_id)
+                            result = await self._step_map[branch.to].execute(ctx)
+                            step_span.set_attribute("step.status", result.status)
+                            step_span.set_attribute("step.attempts", result.attempts)
+                            step_span.set_attribute("step.duration_ms", result.duration_ms)
+                            attach_flags_to_span(result.flags)
+
+                        self._emitter.step_finished(ctx, branch.to, result.duration_ms, result.status)
+                        if event_callback:
+                            await event_callback("step_finished", {"step_name": branch.to, "run_id": ctx.run_id, "status": result.status, "duration_ms": result.duration_ms})
+                        step_results.append(result)
+                        if result.status != "success":
+                            duration = (time.perf_counter() - start) * 1000
+                            self._emitter.pipeline_finished(ctx, self.name, duration, result.status)
+                            if event_callback:
+                                await event_callback("pipeline_finished", {"run_id": ctx.run_id, "status": result.status, "duration_ms": duration})
+                            return PipelineResult(output=result.output, run_id=ctx.run_id, trace_id=ctx.trace_id, duration_ms=duration, status=result.status, step_results=step_results, flags=ctx.flags)
+                        continue
+
+                    ctx._replace_state(result.output)
+                    if self.state_backend:
+                        await self.state_backend.save(ctx.run_id, ctx.to_dict())
+
+                duration = (time.perf_counter() - start) * 1000
+                self._emitter.pipeline_finished(ctx, self.name, duration, "success")
+                if event_callback:
+                    await event_callback("pipeline_finished", {"run_id": ctx.run_id, "status": "success", "duration_ms": duration})
+                if self.state_backend:
+                    await self.state_backend.delete(ctx.run_id)
+                return PipelineResult(output=step_results[-1].output if step_results else None, run_id=ctx.run_id, trace_id=ctx.trace_id, duration_ms=duration, status="success", step_results=step_results, flags=ctx.flags)
+        finally:
+            if bus:
+                await bus.stop()
+
+    async def resume(self, run_id: str, state_schema: type[BaseModel]) -> PipelineResult:
+        if not self.state_backend:
+            raise RuntimeError("Cannot resume — no state_backend configured")
+        data = await self.state_backend.load(run_id)
+        ctx = Context.from_dict(data, state_schema)
+
+        last = ctx._step_name
+        step_names = [s.name for s in self.steps]
+        start_index = step_names.index(last) + 1 if last in step_names else 0
+
+        remaining = self.steps[start_index:]
+        pipeline = Pipeline(steps=remaining, name=self.name, state_backend=self.state_backend, mcp_servers=self.mcp_servers)
+        return await pipeline.run(ctx)

state_mesh/core/step.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from dataclasses import dataclass
+from typing import Any, Literal, Optional, Callable
+
+from pydantic import BaseModel, Field
+
+from state_mesh.core.context import Context, Flag
+from state_mesh.guardrails.runner import run_guards
+from state_mesh.output.retry import run_with_retry
+from state_mesh.output.parser import Parser
+from state_mesh.output.contract import OutputContract
+
+
+class StepResult(BaseModel):
+    status: Literal["success", "failed", "timed_out", "guarded"]
+    step_name: str
+    output: Any
+    duration_ms: float
+    attempts: int
+    flags: list[Flag] = Field(default_factory=list)
+    error: Optional[str] = None
+
+
+@dataclass
+class RetryConfig:
+    max_attempts: int = 3
+    backoff_base: float = 1.0
+    backoff_multiplier: float = 2.0
+
+
+class Branch:
+    def __init__(self, to: str, context: Context):
+        self.to = to
+        self.context = context
+
+
+class Step:
+    def __init__(self, fn: Callable, name: str | None = None, retry_config: RetryConfig = None,
+                 timeout_seconds: Optional[float] = None, guard_before: list = None, guard_after: list = None,
+                 tags: Optional[list[str]] = None, output_contract: Optional[OutputContract] = None):
+        self.name = name or fn.__name__
+        self.fn = fn
+        self.retry_config = retry_config or RetryConfig()
+        self.timeout_seconds = timeout_seconds
+        self.guard_before = guard_before or []
+        self.guard_after = guard_after or []
+        self.tags = tags or []
+        self.output_contract = output_contract
+
+    async def execute(self, ctx: Context, prompt: str | None = None) -> StepResult:
+        start = time.perf_counter()
+        last_error = None
+        ctx._set_current_step(self.name)
+
+        if self.guard_before:
+            guard_result = await run_guards(self.guard_before, ctx, ctx.state)
+            if not guard_result.passed:
+                elapsed = (time.perf_counter() - start) * 1000
+                return StepResult(status="guarded", flags=ctx.flags, step_name=self.name, output=None, duration_ms=elapsed, attempts=0, error=guard_result.reason)
+
+        for attempt in range(1, self.retry_config.max_attempts + 1):
+            try:
+                if self.output_contract is not None:
+                    coro = run_with_retry(self.fn, prompt, self.output_contract, Parser())
+                    if self.timeout_seconds:
+                        result = await asyncio.wait_for(coro, timeout=self.timeout_seconds)
+                    else:
+                        result = await coro
+                elif self.timeout_seconds:
+                    result = await asyncio.wait_for(self.fn(ctx), timeout=self.timeout_seconds)
+                else:
+                    result = await self.fn(ctx)
+
+                if self.guard_after:
+                    guard_result = await run_guards(self.guard_after, ctx, result)
+                    if not guard_result.passed:
+                        elapsed = (time.perf_counter() - start) * 1000
+                        return StepResult(status="guarded", flags=ctx.flags, step_name=self.name, output=None, duration_ms=elapsed, attempts=attempt, error=guard_result.reason)
+
+                elapsed = (time.perf_counter() - start) * 1000
+
+                if isinstance(result, Branch):
+                    return StepResult(status="success", flags=ctx.flags, step_name=self.name, output=result, duration_ms=elapsed, attempts=attempt)
+
+                return StepResult(status="success", flags=ctx.flags, step_name=self.name, output=result, duration_ms=elapsed, attempts=attempt)
+            except asyncio.TimeoutError:
+                elapsed = (time.perf_counter() - start) * 1000
+                return StepResult(
+                    status="timed_out",
+                    flags=ctx.flags,
+                    step_name=self.name,
+                    output=None,
+                    duration_ms=elapsed,
+                    attempts=attempt,
+                    error=f"Step timed out after {self.timeout_seconds} seconds"
+                )
+            except Exception as e:
+                last_error = e
+                if attempt < self.retry_config.max_attempts:
+                    wait = self.retry_config.backoff_base * (self.retry_config.backoff_multiplier ** (attempt - 1))
+                    await asyncio.sleep(wait)
+                continue
+
+        elapsed = (time.perf_counter() - start) * 1000
+        return StepResult(status="failed", flags=ctx.flags, step_name=self.name, output=None, duration_ms=elapsed, attempts=self.retry_config.max_attempts, error=str(last_error))
+
+
+def step(fn=None, *, name=None, timeout_seconds=None, retry_config=None, guard_before=None, guard_after=None, tags=None, output_contract=None):
+    if fn is not None:
+        return Step(fn=fn)
+    def wrap(f):
+        return Step(fn=f, name=name, timeout_seconds=timeout_seconds, retry_config=retry_config, guard_before=guard_before, guard_after=guard_after, tags=tags, output_contract=output_contract)
+    return wrap

state_mesh/guardrails/base.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal
+from abc import ABC, abstractmethod
+
+from state_mesh.core.context import Context
+
+
+@dataclass
+class GuardResult:
+    passed: bool
+    reason: str
+    severity: Literal["warn", "block"] = "warn"
+
+
+class Guard(ABC):
+    @abstractmethod
+    async def check(self, ctx: Context, data: Any) -> GuardResult:
+        pass

state_mesh/guardrails/content.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import re
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+from state_mesh.guardrails.base import Guard, GuardResult
+from state_mesh.core.context import Context
+
+try:
+    from detoxify import Detoxify
+except ImportError:  # pragma: no cover - optional dependency
+    class Detoxify:  # type: ignore[no-redef]
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            raise ImportError("detoxify is required to use ToxicityGuard")
+
+EMAIL_PATTERN = re.compile(r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}")
+
+
+class PIIGuard(Guard):
+    def __init__(
+        self,
+        sensitive_fields: list[str] | None = None,
+        ignore_patterns: list[str] | None = None,
+        severity: Literal["warn", "block"] = "block",
+    ):
+        self.sensitive_fields = sensitive_fields
+        self.ignore_patterns = [re.compile(p) for p in (ignore_patterns or [])]
+        self.severity = severity
+
+    async def check(self, ctx: Context, data: Any) -> GuardResult:
+        if self.sensitive_fields and isinstance(data, BaseModel):
+            candidates = {
+                f: str(v)
+                for f, v in data.model_dump().items()
+                if f in self.sensitive_fields
+            }
+        else:
+            candidates = {"data": data if isinstance(data, str) else str(data)}
+
+        for field, value in candidates.items():
+            if any(p.search(value) for p in self.ignore_patterns):
+                continue
+            if EMAIL_PATTERN.search(value):
+                return GuardResult(passed=False, severity=self.severity, reason=f"Email detected in field '{field}'")
+
+        return GuardResult(passed=True, reason="No PII detected")
+
+
+class ConfidenceGuard(Guard):
+    def __init__(self, min_confidence: float, field: str = "confidence", severity: Literal["warn", "block"] = "block"):
+        self.min_confidence = min_confidence
+        self.field = field
+        self.severity = severity
+
+    async def check(self, ctx: Context, data: Any) -> GuardResult:
+        if isinstance(data, BaseModel):
+            score = data.model_dump().get(self.field)
+        elif isinstance(data, dict):
+            score = data.get(self.field)
+        else:
+            return GuardResult(passed=False, severity=self.severity, reason=f"Cannot extract confidence field '{self.field}' from {type(data).__name__}")
+
+        if score is None:
+            return GuardResult(passed=False, severity=self.severity, reason=f"Confidence field '{self.field}' missing")
+
+        if score < self.min_confidence:
+            return GuardResult(passed=False, severity=self.severity, reason=f"Confidence {score:.2f} below threshold {self.min_confidence:.2f}")
+
+        return GuardResult(passed=True, reason=f"Confidence {score:.2f} meets threshold")
+
+
+class ToxicityGuard(Guard):
+    def __init__(self, threshold: float = 0.5, categories: list[str] | None = None, severity: Literal["warn", "block"] = "block"):
+        self._model = Detoxify("original")
+        self.threshold = threshold
+        self.categories = categories or ["toxicity", "severe_toxicity", "obscene", "threat", "insult", "identity_attack"]
+        self.severity = severity
+
+    async def check(self, ctx: Context, data: Any) -> GuardResult:
+        text = data if isinstance(data, str) else str(data)
+        scores = self._model.predict(text)
+
+        for category in self.categories:
+            score = scores.get(category, 0)
+            if score >= self.threshold:
+                return GuardResult(passed=False, severity=self.severity, reason=f"Toxicity detected — {category}: {score:.2f} (threshold: {self.threshold:.2f})")
+
+        return GuardResult(passed=True, reason="No toxic content detected")
+

state_mesh/guardrails/runner.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import Any
+
+from state_mesh.guardrails.base import Guard, GuardResult
+from state_mesh.core.context import Context
+
+
+async def run_guards(guards: list[Guard], ctx: Context, data: Any) -> GuardResult:
+    for guard in guards:
+        result = await guard.check(ctx, data)
+
+        if not result.passed:
+            if result.severity == "block":
+                return result
+            if result.severity == "warn":
+                ctx.emit_flag(flag_type="guard_warn", reason=result.reason, severity="warn")
+
+    return GuardResult(passed=True, reason="all guards passed")

state_mesh/guardrails/schema.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from typing import Any, Type
+
+from pydantic import BaseModel, ValidationError
+
+from state_mesh.guardrails.base import Guard, GuardResult
+from state_mesh.core.context import Context
+
+
+class SchemaGuard(Guard):
+    def __init__(self, schema: Type[BaseModel]):
+        self.schema = schema
+
+    async def check(self, ctx: Context, data: Any) -> GuardResult:
+        try:
+            self.schema.model_validate(data)
+            return GuardResult(passed=True, reason="")
+        except ValidationError as e:
+            return GuardResult(passed=False, reason=str(e))