Lifecycle Callbacks

from tenxgraph.utils.callbacks import (
    GraphLifecycleHook,
    GraphLifecycleContext,
)
from tenxgraph.core.state import AgentState, Message

from tenxgraph.utils.callbacks import GraphLifecycleHook, GraphLifecycleContext
from tenxgraph.core.state import AgentState
import uuid

class TraceStartHook(GraphLifecycleHook):
    async def on_graph_start(
        self,
        context: GraphLifecycleContext,
        state: AgentState,
    ) -> AgentState | None:
        """Initialize observability trace when graph starts."""
        trace_id = str(uuid.uuid4())
        state.execution_meta.internal_data["trace_id"] = trace_id
        print(f"[TRACE START] thread_id={context.thread_id} trace_id={trace_id}")
        return state  # Return modified state

from tenxgraph.utils.callbacks import GraphLifecycleHook, GraphLifecycleContext
from tenxgraph.core.state import AgentState, Message
import json

class MetricsEndHook(GraphLifecycleHook):
    async def on_graph_end(
        self,
        context: GraphLifecycleContext,
        final_state: AgentState,
        messages: list[Message],
        total_steps: int,
    ) -> AgentState | None:
        """Record execution metrics and send completion notification."""
        metrics = {
            "thread_id": context.thread_id,
            "run_id": context.run_id,
            "total_steps": total_steps,
            "message_count": len(messages),
            "completed_at": context.config.get("timestamp"),
        }
        print(f"[METRICS] {json.dumps(metrics)}")
        
        # Optionally notify external systems
        await self._send_webhook_notification(context, final_state)
        
        # Close OpenTelemetry span
        trace_id = final_state.execution_meta.internal_data.get("trace_id")
        if trace_id:
            print(f"[TRACE END] trace_id={trace_id} status=success")
        
        return None  # Keep state unchanged

from tenxgraph.utils.callbacks import GraphLifecycleHook, GraphLifecycleContext
from tenxgraph.core.state import AgentState, Message
import logging

logger = logging.getLogger(__name__)

class ErrorAlertHook(GraphLifecycleHook):
    async def on_graph_error(
        self,
        context: GraphLifecycleContext,
        error: Exception,
        partial_state: AgentState,
        messages: list[Message],
        step: int,
        node_name: str,
    ) -> tuple[AgentState, str] | None:
        """Alert PagerDuty on production errors and mask sensitive state."""
        logger.error(
            f"Graph failed at node={node_name} step={step}: {error}",
            extra={
                "thread_id": context.thread_id,
                "run_id": context.run_id,
            }
        )
        
        # Alert on-call engineer
        trace_id = partial_state.execution_meta.internal_data.get("trace_id")
        await self._page_duty_alert(
            title=f"Agent graph failed: {node_name}",
            description=f"Trace: {trace_id}\nError: {str(error)[:200]}",
            severity="error" if "critical" in str(error).lower() else "warning",
        )
        
        # Mask sensitive state before persistence
        partial_state.execution_meta.internal_data["masked"] = True
        masked_msg = f"Graph failed: {type(error).__name__} (details logged)"
        
        return partial_state, masked_msg

from tenxgraph.utils.callbacks import GraphLifecycleHook, GraphLifecycleContext
from tenxgraph.core.state import AgentState

class InterruptNotificationHook(GraphLifecycleHook):
    async def on_interrupt(
        self,
        context: GraphLifecycleContext,
        interrupted_node: str,
        interrupt_type: str,  # "before" | "after" | "stop" | "remote_tool"
        state: AgentState,
    ) -> AgentState | None:
        """Notify frontend and start timeout timer when paused."""
        print(f"[INTERRUPT] node={interrupted_node} type={interrupt_type}")
        
        # Mark in state for frontend to pick up
        state.execution_meta.internal_data["paused_at"] = context.config.get("timestamp")
        state.execution_meta.internal_data["interrupt_type"] = interrupt_type
        
        # Notify frontend via WebSocket or server-sent events
        await self._notify_frontend(
            thread_id=context.thread_id,
            status="waiting_input",
            node=interrupted_node,
        )
        
        # Start a timeout timer: auto-cancel if not resumed in 1 hour
        await self._start_timeout_timer(
            thread_id=context.thread_id,
            timeout_seconds=3600,
        )
        
        return state

from tenxgraph.utils.callbacks import GraphLifecycleHook, GraphLifecycleContext
from tenxgraph.core.state import AgentState
from typing import Any

class ResumeValidationHook(GraphLifecycleHook):
    async def on_resume(
        self,
        context: GraphLifecycleContext,
        resumed_node: str,
        state: AgentState,
        resume_data: dict[str, Any],
    ) -> AgentState | None:
        """Validate resume data and record who approved."""
        print(f"[RESUME] node={resumed_node} approved_by={resume_data.get('approved_by')}")
        
        # Validate resume payload
        if "approval_code" in resume_data:
            is_valid = await self._validate_approval_code(resume_data["approval_code"])
            if not is_valid:
                raise ValueError("Invalid approval code")
        
        # Record approval in audit log
        state.execution_meta.internal_data["resumed_at"] = context.config.get("timestamp")
        state.execution_meta.internal_data["resumed_by"] = resume_data.get("approved_by", "unknown")
        
        # Cancel the timeout timer from on_interrupt
        await self._cancel_timeout_timer(context.thread_id)
        
        # Update frontend status back to "thinking..."
        await self._notify_frontend(
            thread_id=context.thread_id,
            status="resuming",
        )
        
        return state

from tenxgraph.utils.callbacks import GraphLifecycleHook, GraphLifecycleContext
from tenxgraph.core.state import AgentState, Message
from typing import Sequence

class RedactionCheckpointHook(GraphLifecycleHook):
    async def on_checkpoint(
        self,
        context: GraphLifecycleContext,
        state: AgentState,
        messages: list[Message],
        is_context_trimmed: bool,
    ) -> tuple[AgentState, list[Message]] | AgentState | None:
        """Redact PII before persistence and log for compliance."""
        # Redact sensitive fields
        for msg in messages:
            if hasattr(msg.content, "__dict__"):
                if "email" in msg.content:
                    msg.content["email"] = "***@***.com"
                if "phone" in msg.content:
                    msg.content["phone"] = "***-****"
        
        # Drop transient messages (e.g., debug logs) before checkpoint
        filtered_messages = [
            m for m in messages
            if not m.metadata.get("transient", False)
        ]
        
        # Write to compliance audit log (HIPAA, SOC2, etc.)
        audit_entry = {
            "timestamp": context.config.get("timestamp"),
            "thread_id": context.thread_id,
            "run_id": context.run_id,
            "context_trimmed": is_context_trimmed,
        }
        await self._write_compliance_log(audit_entry)
        
        # Replicate to secondary cache (Redis, CDN edge)
        await self._replicate_to_cache(
            thread_id=context.thread_id,
            state=state,
            ttl=3600,
        )
        
        return state, filtered_messages

from tenxgraph.utils.callbacks import GraphLifecycleHook, GraphLifecycleContext
from tenxgraph.core.state import AgentState
import copy

class ObserveNodeTransitionHook(GraphLifecycleHook):
    async def on_state_update(
        self,
        context: GraphLifecycleContext,
        node_name: str,
        old_state: AgentState,
        new_state: AgentState,
        step: int,
    ) -> AgentState | None:
        """Log each node transition and detect infinite loops."""
        # Detect infinite loop: same node executed 5+ times in a row
        if hasattr(self, "_last_node"):
            if self._last_node == node_name:
                self._repeat_count = getattr(self, "_repeat_count", 1) + 1
                if self._repeat_count > 5:
                    print(f"[WARNING] Possible infinite loop: node {node_name} repeated {self._repeat_count} times")
                    # Could raise exception or set interrupt flag
            else:
                self._repeat_count = 1
        
        self._last_node = node_name
        
        # Log state transition for observability
        print(f"[STEP {step}] {node_name}")
        print(f"  Messages before: {len(old_state.messages)}")
        print(f"  Messages after: {len(new_state.messages)}")
        
        # Example: apply node-specific policies
        if node_name == "approval_node":
            # Verify approval was recorded
            if not new_state.execution_meta.internal_data.get("approved"):
                print(f"[WARNING] Approval node completed but no approval recorded")
        
        return None  # Keep state unchanged

class TurnLogHook(GraphLifecycleHook):
    async def on_turn_end(self, context, state, turn_index):
        print(f"turn {turn_index} ended on thread {context.thread_id}")
        return None

from tenxgraph.utils.callbacks import GraphLifecycleHook, GraphLifecycleContext
from tenxgraph.core.state import AgentState, Message
from opentelemetry import trace

class OpenTelemetryHook(GraphLifecycleHook):
    def __init__(self):
        self._tracer = trace.get_tracer(__name__)
        self._spans = {}

    async def on_graph_start(
        self, context: GraphLifecycleContext, state: AgentState
    ) -> AgentState | None:
        span = self._tracer.start_span(
            name=f"graph_{context.thread_id}",
            attributes={
                "thread_id": context.thread_id,
                "run_id": context.run_id,
            }
        )
        self._spans[context.run_id] = span
        return None

    async def on_graph_end(
        self, context: GraphLifecycleContext, final_state: AgentState,
        messages: list[Message], total_steps: int
    ) -> AgentState | None:
        span = self._spans.pop(context.run_id, None)
        if span:
            span.set_attribute("total_steps", total_steps)
            span.set_attribute("message_count", len(messages))
            span.end()
        return None

    async def on_graph_error(
        self, context: GraphLifecycleContext, error: Exception,
        partial_state: AgentState, messages: list[Message],
        step: int, node_name: str
    ) -> tuple[AgentState, str] | None:
        span = self._spans.get(context.run_id)
        if span:
            span.record_exception(error)
            span.set_attribute("error", True)
        return None

from tenxgraph.utils.callbacks import GraphLifecycleHook, GraphLifecycleContext
from tenxgraph.core.state import AgentState

class ApprovalWorkflowHook(GraphLifecycleHook):
    async def on_interrupt(
        self, context: GraphLifecycleContext, interrupted_node: str,
        interrupt_type: str, state: AgentState
    ) -> AgentState | None:
        # Store approval request in external system
        approval_id = await self._create_approval_request(
            thread_id=context.thread_id,
            node=interrupted_node,
        )
        state.execution_meta.internal_data["approval_id"] = approval_id
        return state

    async def on_resume(
        self, context: GraphLifecycleContext, resumed_node: str,
        state: AgentState, resume_data: dict
    ) -> AgentState | None:
        # Record approval decision in audit trail
        approval_id = state.execution_meta.internal_data.get("approval_id")
        await self._record_approval(
            approval_id=approval_id,
            decision=resume_data.get("decision"),
            approved_by=resume_data.get("approved_by"),
            timestamp=context.config.get("timestamp"),
        )
        return None

    async def on_checkpoint(
        self, context: GraphLifecycleContext, state: AgentState,
        messages: list, is_context_trimmed: bool
    ) -> tuple[AgentState, list] | None:
        # Ensure approval is durable before returning to caller
        approval_id = state.execution_meta.internal_data.get("approval_id")
        if approval_id:
            await self._persist_approval_to_audit_log(approval_id, state)
        return None

from tenxgraph.utils.callbacks import GraphLifecycleHook, GraphLifecycleContext
from tenxgraph.core.state import AgentState, Message
import time

class MetricsHook(GraphLifecycleHook):
    def __init__(self):
        self._start_times = {}

    async def on_graph_start(
        self, context: GraphLifecycleContext, state: AgentState
    ) -> AgentState | None:
        self._start_times[context.run_id] = time.time()
        return None

    async def on_state_update(
        self, context: GraphLifecycleContext, node_name: str,
        old_state: AgentState, new_state: AgentState, step: int
    ) -> AgentState | None:
        # Record message count change per node
        msg_delta = len(new_state.messages) - len(old_state.messages)
        self._record_metric(
            f"node.message_delta.{node_name}",
            msg_delta,
        )
        return None

    async def on_graph_end(
        self, context: GraphLifecycleContext, final_state: AgentState,
        messages: list[Message], total_steps: int
    ) -> AgentState | None:
        # Record total execution time
        elapsed = time.time() - self._start_times.pop(context.run_id, 0)
        self._record_metric("graph.execution_time", elapsed)
        self._record_metric("graph.total_steps", total_steps)
        self._record_metric("graph.message_count", len(messages))
        return None

    async def on_graph_error(
        self, context: GraphLifecycleContext, error: Exception,
        partial_state: AgentState, messages: list[Message],
        step: int, node_name: str
    ) -> tuple[AgentState, str] | None:
        # Record failure metrics
        self._record_metric("graph.errors", 1)
        self._record_metric(f"graph.error.{node_name}", 1)
        return None

from tenxgraph.core.graph import StateGraph
from tenxgraph.utils.callbacks import CallbackManager, GraphLifecycleHook

class MyLifecycleHook(GraphLifecycleHook):
    async def on_graph_start(self, context, state):
        print(f"Starting graph run: {context.run_id}")
        return None

    async def on_graph_end(self, context, final_state, messages, total_steps):
        print(f"Completed in {total_steps} steps with {len(messages)} messages")
        return None

cbm = CallbackManager()
cbm.register_lifecycle_hook(MyLifecycleHook())

builder = StateGraph(MyState)
# ... define nodes ...

app = builder.compile(callback_manager=cbm)

# Hooks fire automatically
result = await app.ainvoke({"messages": [...]})