From 983b6200969a080b1638b59dda61dfac140dbf4d Mon Sep 17 00:00:00 2001
From: Adrian Lyjak <adrianlyjak@gmail.com>
Date: Sat, 17 Jan 2026 21:49:30 -0500
Subject: [PATCH] plugin param

---
 .../src/workflows/context/context.py          | 18 ++++++++++--
 .../src/workflows/workflow.py                 | 29 +++++++++++++++++++
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/packages/llama-index-workflows/src/workflows/context/context.py b/packages/llama-index-workflows/src/workflows/context/context.py
index 6aa7ba3d..436e62aa 100644
--- a/packages/llama-index-workflows/src/workflows/context/context.py
+++ b/packages/llama-index-workflows/src/workflows/context/context.py
@@ -31,7 +31,6 @@ from workflows.events import (
     StopEvent,
 )
 from workflows.handler import WorkflowHandler
-from workflows.plugins.basic import basic_runtime
 from workflows.runtime.broker import WorkflowBroker
 from workflows.runtime.types.internal_state import BrokerState
 from workflows.runtime.types.plugin import Plugin, WorkflowRuntime
@@ -131,13 +130,26 @@ class Context(Generic[MODEL_T]):
         workflow: Workflow,
         previous_context: dict[str, Any] | None = None,
         serializer: BaseSerializer | None = None,
-        plugin: Plugin = basic_runtime,
+        plugin: Plugin | None = None,
     ) -> None:
         self._serializer = serializer or JsonSerializer()
         self._broker_run = None
-        self._plugin = plugin
         self._workflow = workflow
 
+        # Use workflow's plugin, fall back to explicit param for compatibility
+        # TODO: Remove the plugin parameter entirely - it's not a public interface.
+        # Just use workflow.plugin directly.
+        if plugin is not None:
+            warnings.warn(
+                "Passing plugin to Context is deprecated. "
+                "Pass plugin to Workflow instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            self._plugin = plugin
+        else:
+            self._plugin = workflow.plugin
+
         # parse the serialized context
         serializer = serializer or JsonSerializer()
         if previous_context is not None:
diff --git a/packages/llama-index-workflows/src/workflows/workflow.py b/packages/llama-index-workflows/src/workflows/workflow.py
index 3d8f195e..d7505c3c 100644
--- a/packages/llama-index-workflows/src/workflows/workflow.py
+++ b/packages/llama-index-workflows/src/workflows/workflow.py
@@ -16,6 +16,7 @@ from pydantic import ValidationError
 
 if TYPE_CHECKING:  # pragma: no cover
     from .context import Context
+    from .runtime.types.plugin import Plugin
 from .decorators import StepConfig, StepFunction
 from .errors import (
     WorkflowConfigurationError,
@@ -95,6 +96,8 @@ class Workflow(metaclass=WorkflowMeta):
     # Populated by the metaclass; declared here for type checkers.
     _step_functions: dict[str, StepFunction]
 
+    _plugin: Plugin
+
     def __init__(
         self,
         timeout: float | None = 45.0,
@@ -102,6 +105,7 @@ class Workflow(metaclass=WorkflowMeta):
         verbose: bool = False,
         resource_manager: ResourceManager | None = None,
         num_concurrent_runs: int | None = None,
+        plugin: Plugin | None = None,
     ) -> None:
         """
         Initialize a workflow instance.
@@ -115,6 +119,9 @@ class Workflow(metaclass=WorkflowMeta):
             resource_manager (ResourceManager | None): Custom resource manager
                 for dependency injection.
             num_concurrent_runs (int | None): Limit on concurrent `run()` calls.
+            plugin (Plugin | None): Optional plugin to use for this workflow.
+                If not provided, uses the current context-scoped plugin or
+                falls back to basic_runtime.
         """
         # Configuration
         self._timeout = timeout
@@ -133,6 +140,28 @@ class Workflow(metaclass=WorkflowMeta):
         # Instrumentation
         self._dispatcher = dispatcher
 
+        # Plugin registration: explicit > context-scoped > basic_runtime
+        from workflows.plugins._context import get_current_plugin
+        from workflows.runtime.workflow_tracker import WorkflowTracker
+
+        if plugin is not None:
+            self._plugin = plugin
+        else:
+            # get_current_plugin() falls back to basic_runtime
+            self._plugin = get_current_plugin()
+
+        # Register with plugin's tracker (if plugin supports it)
+        # TODO: Replace this getattr kludge with a proper method on Plugin ABC,
+        # e.g., plugin.track_workflow(self) that's a no-op in BasicRuntime.
+        tracker = getattr(self._plugin, "_tracker", None)
+        if isinstance(tracker, WorkflowTracker):
+            tracker.add(self)
+
+    @property
+    def plugin(self) -> Plugin:
+        """The plugin this workflow is registered with."""
+        return self._plugin
+
     def _ensure_start_event_class(self) -> type[StartEvent]:
         """
         Returns the StartEvent type used in this workflow.