feat: workflow step interceptor

evan-ak · evan-ak · commit 29125f72c302 · 2026-01-14T22:22:20.000+09:00
diff --git a/src/memu/app/service.py b/src/memu/app/service.py
@@ -30,6 +30,7 @@
     LLMInterceptorHandle,
     LLMInterceptorRegistry,
 )
+from memu.workflow.interceptor import WorkflowInterceptorHandle, WorkflowInterceptorRegistry
 from memu.workflow.pipeline import PipelineManager
 from memu.workflow.runner import WorkflowRunner, resolve_workflow_runner
 from memu.workflow.step import WorkflowState, WorkflowStep
@@ -83,6 +84,7 @@ def __init__(
         # Initialize client caches (lazy creation on first use)
         self._llm_clients: dict[str, Any] = {}
         self._llm_interceptors = LLMInterceptorRegistry()
+        self._workflow_interceptors = WorkflowInterceptorRegistry()
 
         self._workflow_runner = resolve_workflow_runner(workflow_runner)
 
@@ -240,6 +242,45 @@ def intercept_on_error_llm_call(
     ) -> LLMInterceptorHandle:
         return self._llm_interceptors.register_on_error(fn, name=name, priority=priority, where=where)
 
+    def intercept_before_workflow_step(
+        self,
+        fn: Callable[..., Any],
+        *,
+        name: str | None = None,
+    ) -> WorkflowInterceptorHandle:
+        """
+        Register an interceptor to be called before each workflow step.
+
+        The interceptor receives (step_context: WorkflowStepContext, state: WorkflowState).
+        """
+        return self._workflow_interceptors.register_before(fn, name=name)
+
+    def intercept_after_workflow_step(
+        self,
+        fn: Callable[..., Any],
+        *,
+        name: str | None = None,
+    ) -> WorkflowInterceptorHandle:
+        """
+        Register an interceptor to be called after each workflow step.
+
+        The interceptor receives (step_context: WorkflowStepContext, state: WorkflowState).
+        """
+        return self._workflow_interceptors.register_after(fn, name=name)
+
+    def intercept_on_error_workflow_step(
+        self,
+        fn: Callable[..., Any],
+        *,
+        name: str | None = None,
+    ) -> WorkflowInterceptorHandle:
+        """
+        Register an interceptor to be called when a workflow step raises an exception.
+
+        The interceptor receives (step_context: WorkflowStepContext, state: WorkflowState, error: Exception).
+        """
+        return self._workflow_interceptors.register_on_error(fn, name=name)
+
     def _get_context(self) -> Context:
         return self._context
 
@@ -292,7 +333,13 @@ async def _run_workflow(self, workflow_name: str, initial_state: WorkflowState)
         """Execute a workflow through the configured runner backend."""
         steps = self._pipelines.build(workflow_name)
         runner_context = {"workflow_name": workflow_name}
-        return await self._workflow_runner.run(workflow_name, steps, initial_state, runner_context)
+        return await self._workflow_runner.run(
+            workflow_name,
+            steps,
+            initial_state,
+            runner_context,
+            interceptor_registry=self._workflow_interceptors,
+        )
 
     @staticmethod
     def _extract_json_blob(raw: str) -> str:
diff --git a/src/memu/workflow/__init__.py b/src/memu/workflow/__init__.py
@@ -1,3 +1,8 @@
+from memu.workflow.interceptor import (
+    WorkflowInterceptorHandle,
+    WorkflowInterceptorRegistry,
+    WorkflowStepContext,
+)
 from memu.workflow.pipeline import PipelineManager, PipelineRevision
 from memu.workflow.runner import (
     LocalWorkflowRunner,
@@ -12,9 +17,12 @@
     "PipelineManager",
     "PipelineRevision",
     "WorkflowContext",
+    "WorkflowInterceptorHandle",
+    "WorkflowInterceptorRegistry",
     "WorkflowRunner",
     "WorkflowState",
     "WorkflowStep",
+    "WorkflowStepContext",
     "register_workflow_runner",
     "resolve_workflow_runner",
     "run_steps",
diff --git a/src/memu/workflow/interceptor.py b/src/memu/workflow/interceptor.py
@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import inspect
+import logging
+import threading
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from memu.workflow.step import WorkflowState
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class WorkflowStepContext:
+    """Context information for a workflow step execution."""
+
+    workflow_name: str
+    step_id: str
+    step_role: str
+    step_context: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class _WorkflowInterceptor:
+    interceptor_id: int
+    fn: Callable[..., Any]
+    name: str | None
+
+
+@dataclass(frozen=True)
+class _WorkflowInterceptorSnapshot:
+    before: tuple[_WorkflowInterceptor, ...]
+    after: tuple[_WorkflowInterceptor, ...]
+    on_error: tuple[_WorkflowInterceptor, ...]
+
+
+class WorkflowInterceptorHandle:
+    """Handle for disposing a registered workflow interceptor."""
+
+    def __init__(self, registry: WorkflowInterceptorRegistry, interceptor_id: int) -> None:
+        self._registry = registry
+        self._interceptor_id = interceptor_id
+        self._disposed = False
+
+    def dispose(self) -> bool:
+        """Remove the interceptor from the registry. Returns True if removed."""
+        if self._disposed:
+            return False
+        self._disposed = True
+        return self._registry.remove(self._interceptor_id)
+
+
+class WorkflowInterceptorRegistry:
+    """
+    Registry for workflow step interceptors.
+
+    Interceptors are called before and after each workflow step execution.
+    Unlike LLM interceptors, workflow interceptors do not support filtering,
+    priority, or ordering - they are called in registration order.
+    """
+
+    def __init__(self, *, strict: bool = False) -> None:
+        self._before: tuple[_WorkflowInterceptor, ...] = ()
+        self._after: tuple[_WorkflowInterceptor, ...] = ()
+        self._on_error: tuple[_WorkflowInterceptor, ...] = ()
+        self._lock = threading.Lock()
+        self._seq = 0
+        self._strict = strict
+
+    @property
+    def strict(self) -> bool:
+        """If True, interceptor exceptions will propagate instead of being logged."""
+        return self._strict
+
+    def register_before(
+        self,
+        fn: Callable[..., Any],
+        *,
+        name: str | None = None,
+    ) -> WorkflowInterceptorHandle:
+        """
+        Register an interceptor to be called before each step.
+
+        The interceptor receives (step_context: WorkflowStepContext, state: WorkflowState).
+        """
+        return self._register("before", fn, name=name)
+
+    def register_after(
+        self,
+        fn: Callable[..., Any],
+        *,
+        name: str | None = None,
+    ) -> WorkflowInterceptorHandle:
+        """
+        Register an interceptor to be called after each step.
+
+        The interceptor receives (step_context: WorkflowStepContext, state: WorkflowState).
+        """
+        return self._register("after", fn, name=name)
+
+    def register_on_error(
+        self,
+        fn: Callable[..., Any],
+        *,
+        name: str | None = None,
+    ) -> WorkflowInterceptorHandle:
+        """
+        Register an interceptor to be called when a step raises an exception.
+
+        The interceptor receives (step_context: WorkflowStepContext, state: WorkflowState, error: Exception).
+        """
+        return self._register("on_error", fn, name=name)
+
+    def _register(
+        self,
+        kind: str,
+        fn: Callable[..., Any],
+        *,
+        name: str | None,
+    ) -> WorkflowInterceptorHandle:
+        if not callable(fn):
+            msg = "Interceptor must be callable"
+            raise TypeError(msg)
+        with self._lock:
+            self._seq += 1
+            interceptor = _WorkflowInterceptor(
+                interceptor_id=self._seq,
+                fn=fn,
+                name=name,
+            )
+            if kind == "before":
+                self._before = (*self._before, interceptor)
+            elif kind == "after":
+                self._after = (*self._after, interceptor)
+            elif kind == "on_error":
+                self._on_error = (*self._on_error, interceptor)
+            else:
+                msg = f"Unknown interceptor kind '{kind}'"
+                raise ValueError(msg)
+        return WorkflowInterceptorHandle(self, interceptor.interceptor_id)
+
+    def remove(self, interceptor_id: int) -> bool:
+        """Remove an interceptor by ID. Returns True if found and removed."""
+        with self._lock:
+            removed = False
+            before = tuple(i for i in self._before if i.interceptor_id != interceptor_id)
+            after = tuple(i for i in self._after if i.interceptor_id != interceptor_id)
+            on_error = tuple(i for i in self._on_error if i.interceptor_id != interceptor_id)
+            if len(before) != len(self._before):
+                removed = True
+                self._before = before
+            if len(after) != len(self._after):
+                removed = True
+                self._after = after
+            if len(on_error) != len(self._on_error):
+                removed = True
+                self._on_error = on_error
+        return removed
+
+    def snapshot(self) -> _WorkflowInterceptorSnapshot:
+        """Get a point-in-time snapshot of registered interceptors."""
+        return _WorkflowInterceptorSnapshot(self._before, self._after, self._on_error)
+
+
+async def run_before_interceptors(
+    interceptors: tuple[_WorkflowInterceptor, ...],
+    step_context: WorkflowStepContext,
+    state: WorkflowState,
+    *,
+    strict: bool = False,
+) -> None:
+    """Run all before-step interceptors."""
+    for interceptor in interceptors:
+        await _safe_invoke_interceptor(interceptor, strict, step_context, state)
+
+
+async def run_after_interceptors(
+    interceptors: tuple[_WorkflowInterceptor, ...],
+    step_context: WorkflowStepContext,
+    state: WorkflowState,
+    *,
+    strict: bool = False,
+) -> None:
+    """Run all after-step interceptors in reverse order."""
+    for interceptor in reversed(interceptors):
+        await _safe_invoke_interceptor(interceptor, strict, step_context, state)
+
+
+async def run_on_error_interceptors(
+    interceptors: tuple[_WorkflowInterceptor, ...],
+    step_context: WorkflowStepContext,
+    state: WorkflowState,
+    error: Exception,
+    *,
+    strict: bool = False,
+) -> None:
+    """Run all on-error interceptors in reverse order."""
+    for interceptor in reversed(interceptors):
+        await _safe_invoke_interceptor(interceptor, strict, step_context, state, error)
+
+
+async def _safe_invoke_interceptor(
+    interceptor: _WorkflowInterceptor,
+    strict: bool,
+    *args: Any,
+) -> None:
+    """Safely invoke an interceptor, handling exceptions based on strict mode."""
+    try:
+        result = interceptor.fn(*args)
+        if inspect.isawaitable(result):
+            await result
+    except Exception:
+        if strict:
+            raise
+        logger.exception("Workflow interceptor failed: %s", interceptor.name or interceptor.interceptor_id)
diff --git a/src/memu/workflow/runner.py b/src/memu/workflow/runner.py
@@ -1,10 +1,13 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 from memu.workflow.step import WorkflowContext, WorkflowState, WorkflowStep, run_steps
 
+if TYPE_CHECKING:
+    from memu.workflow.interceptor import WorkflowInterceptorRegistry
+
 
 @runtime_checkable
 class WorkflowRunner(Protocol):
@@ -18,6 +21,7 @@ async def run(
         steps: list[WorkflowStep],
         initial_state: WorkflowState,
         context: WorkflowContext = None,
+        interceptor_registry: WorkflowInterceptorRegistry | None = None,
     ) -> WorkflowState: ...
 
 
@@ -30,8 +34,9 @@ async def run(
         steps: list[WorkflowStep],
         initial_state: WorkflowState,
         context: WorkflowContext = None,
+        interceptor_registry: WorkflowInterceptorRegistry | None = None,
     ) -> WorkflowState:
-        return await run_steps(workflow_name, steps, initial_state, context)
+        return await run_steps(workflow_name, steps, initial_state, context, interceptor_registry)
 
 
 RunnerFactory = Callable[[], WorkflowRunner]
diff --git a/src/memu/workflow/step.py b/src/memu/workflow/step.py
