griptape-nodes 0.42.0__py3-none-any.whl → 0.43.0__py3-none-any.whl

griptape_nodes/exe_types/node_types.py

@@ -65,6 +65,9 @@ class BaseNode(ABC):
     stop_flow: bool = False
     root_ui_element: BaseNodeElement
     _tracked_parameters: list[BaseNodeElement]
+    _entry_control_parameter: Parameter | None = (
+        None  # The control input parameter used to enter this node during execution
+    )
 
     @property
     def parameters(self) -> list[Parameter]:
@@ -92,6 +95,7 @@ class BaseNode(ABC):
         self.root_ui_element._node_context = self
         self.process_generator = None
         self._tracked_parameters = []
+        self.set_entry_control_parameter(None)
 
     # This is gross and we need to have a universal pass on resolution state changes and emission of events. That's what this ticket does!
     # https://github.com/griptape-ai/griptape-nodes/issues/994
@@ -106,6 +110,18 @@ class BaseNode(ABC):
                 )
             )
         self.state = NodeResolutionState.UNRESOLVED
+        # NOTE: _entry_control_parameter is NOT cleared here as it represents execution context
+        # that should persist through the resolve/unresolve cycle during a single execution
+
+    def set_entry_control_parameter(self, parameter: Parameter | None) -> None:
+        """Set the control parameter that was used to enter this node.
+
+        This should only be called by the ControlFlowContext during execution.
+
+        Args:
+            parameter: The control input parameter that triggered this node's execution, or None to clear
+        """
+        self._entry_control_parameter = parameter
 
     def emit_parameter_changes(self) -> None:
         if self._tracked_parameters:
@@ -577,7 +593,7 @@ class BaseNode(ABC):
 
         def _flatten(items: Iterable[Any]) -> Generator[Any, None, None]:
             for item in items:
-                if isinstance(item, Iterable) and not isinstance(item, (str, bytes)):
+                if isinstance(item, Iterable) and not isinstance(item, (str, bytes, dict)):
                     yield from _flatten(item)
                 elif item:
                     yield item

griptape_nodes/machines/control_flow.py

@@ -2,10 +2,12 @@
 from __future__ import annotations
 
 import logging
+from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 from griptape.events import EventBus
 
+from griptape_nodes.exe_types.core_types import Parameter
 from griptape_nodes.exe_types.node_types import BaseNode, NodeResolutionState
 from griptape_nodes.exe_types.type_validator import TypeValidator
 from griptape_nodes.machines.fsm import FSM, State
@@ -17,6 +19,15 @@ from griptape_nodes.retained_mode.events.execution_events import (
     SelectedControlOutputEvent,
 )
 
+
+@dataclass
+class NextNodeInfo:
+    """Information about the next node to execute and how to reach it."""
+
+    node: BaseNode
+    entry_parameter: Parameter | None
+
+
 if TYPE_CHECKING:
     from griptape_nodes.exe_types.core_types import Parameter
     from griptape_nodes.exe_types.flow import ControlFlow
@@ -36,18 +47,26 @@ class ControlFlowContext:
         self.resolution_machine = NodeResolutionMachine()
         self.current_node = None
 
-    def get_next_node(self, output_parameter: Parameter) -> BaseNode | None:
+    def get_next_node(self, output_parameter: Parameter) -> NextNodeInfo | None:
+        """Get the next node and the target parameter that will receive the control flow.
+
+        Returns:
+            NextNodeInfo | None: Information about the next node or None if no connection
+        """
         if self.current_node is not None:
             from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
 
-            node = GriptapeNodes.FlowManager().get_connections().get_connected_node(self.current_node, output_parameter)
-            if node is not None:
-                node, _ = node
+            node_connection = (
+                GriptapeNodes.FlowManager().get_connections().get_connected_node(self.current_node, output_parameter)
+            )
+            if node_connection is not None:
+                node, entry_parameter = node_connection
+                return NextNodeInfo(node=node, entry_parameter=entry_parameter)
             # Continue Execution to the next node that needs to be executed using global execution queue
-            else:
-                # Get the next node in the execution queue, or None if queue is empty
-                node = GriptapeNodes.FlowManager().get_next_node_from_execution_queue()
-            return node
+            # Get the next node in the execution queue, or None if queue is empty
+            node = GriptapeNodes.FlowManager().get_next_node_from_execution_queue()
+            if node is not None:
+                return NextNodeInfo(node=node, entry_parameter=None)
         return None
 
     def reset(self) -> None:
@@ -111,10 +130,11 @@ class NextNodeState(State):
             context.current_node.stop_flow = False
             return CompleteState
         next_output = context.current_node.get_next_control_output()
-        next_node = None
+        next_node_info = None
+
         if next_output is not None:
             context.selected_output = next_output
-            next_node = context.get_next_node(context.selected_output)
+            next_node_info = context.get_next_node(context.selected_output)
             EventBus.publish_event(
                 ExecutionGriptapeNodeEvent(
                     wrapped_event=ExecutionEvent(
@@ -130,11 +150,18 @@ class NextNodeState(State):
 
             # Get the next node in the execution queue, or None if queue is empty
             next_node = GriptapeNodes.FlowManager().get_next_node_from_execution_queue()
+            if next_node is not None:
+                next_node_info = NextNodeInfo(node=next_node, entry_parameter=None)
+
         # The parameter that will be evaluated next
-        if next_node is None:
+        if next_node_info is None:
             # If no node attached
             return CompleteState
-        context.current_node = next_node
+
+        # Always set the entry control parameter (None for execution queue nodes)
+        next_node_info.node.set_entry_control_parameter(next_node_info.entry_parameter)
+
+        context.current_node = next_node_info.node
         context.selected_output = None
         if not context.paused:
             return ResolveNodeState
@@ -177,6 +204,8 @@ class ControlFlowMachine(FSM[ControlFlowContext]):
 
     def start_flow(self, start_node: BaseNode, debug_mode: bool = False) -> None:  # noqa: FBT001, FBT002
         self._context.current_node = start_node
+        # Set entry control parameter for initial node (None for workflow start)
+        start_node.set_entry_control_parameter(None)
         # Set up to debug
         self._context.paused = debug_mode
         self.start(ResolveNodeState)  # Begins the flow

griptape_nodes/machines/fsm.py

@@ -1,4 +1,4 @@
-from typing import Any, Generic, TypeVar
+from typing import Any, TypeVar
 
 T = TypeVar("T")
 
@@ -19,8 +19,13 @@ class State:
         """Called when exiting the state."""
         return
 
+    @staticmethod
+    def on_event(context: Any, event: Any) -> type["State"] | None:  # noqa: ARG004
+        """Called on an event, which may trigger a State transition."""
+        return None
+
 
-class FSM(Generic[T]):
+class FSM[T]:
     def __init__(self, context: T) -> None:
         self._context = context
         self._current_state = None
@@ -53,3 +58,12 @@ class FSM(Generic[T]):
 
         if new_state is not None:
             self.transition_state(new_state)
+
+    def handle_event(self, event: Any) -> None:
+        if self._current_state is None:
+            new_state = None
+        else:
+            new_state = self._current_state.on_event(self._context, event)
+
+        if new_state is not None:
+            self.transition_state(new_state)

griptape_nodes/mcp_server/__init__.py

@@ -0,0 +1 @@
+"""Runs the Griptape Nodes MCP server."""

griptape_nodes/mcp_server/server.py

@@ -0,0 +1,126 @@
+import contextlib
+import json
+import logging
+import os
+from collections.abc import AsyncIterator
+
+import uvicorn
+from fastapi import FastAPI
+from mcp.server.lowlevel import Server
+from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
+from mcp.types import (
+    TextContent,
+    Tool,
+)
+from pydantic import TypeAdapter
+from rich.logging import RichHandler
+from starlette.types import Receive, Scope, Send
+
+from griptape_nodes.mcp_server.ws_request_manager import AsyncRequestManager, WebSocketConnectionManager
+from griptape_nodes.retained_mode.events.base_events import RequestPayload
+from griptape_nodes.retained_mode.events.connection_events import (
+    CreateConnectionRequest,
+    DeleteConnectionRequest,
+    ListConnectionsForNodeRequest,
+)
+from griptape_nodes.retained_mode.events.flow_events import ListNodesInFlowRequest
+from griptape_nodes.retained_mode.events.node_events import (
+    CreateNodeRequest,
+    DeleteNodeRequest,
+    GetNodeMetadataRequest,
+    GetNodeResolutionStateRequest,
+    ListParametersOnNodeRequest,
+    SetNodeMetadataRequest,
+)
+from griptape_nodes.retained_mode.events.parameter_events import (
+    GetParameterValueRequest,
+    SetParameterValueRequest,
+)
+from griptape_nodes.retained_mode.managers.config_manager import ConfigManager
+from griptape_nodes.retained_mode.managers.secrets_manager import SecretsManager
+
+SUPPORTED_REQUEST_EVENTS: dict[str, type[RequestPayload]] = {
+    # Nodes
+    "CreateNodeRequest": CreateNodeRequest,
+    "DeleteNodeRequest": DeleteNodeRequest,
+    "ListNodesInFlowRequest": ListNodesInFlowRequest,
+    "GetNodeResolutionStateRequest": GetNodeResolutionStateRequest,
+    "GetNodeMetadataRequest": GetNodeMetadataRequest,
+    "SetNodeMetadataRequest": SetNodeMetadataRequest,
+    # Connections
+    "CreateConnectionRequest": CreateConnectionRequest,
+    "DeleteConnectionRequest": DeleteConnectionRequest,
+    "ListConnectionsForNodeRequest": ListConnectionsForNodeRequest,
+    # Parameters
+    "ListParametersOnNodeRequest": ListParametersOnNodeRequest,
+    "GetParameterValueRequest": GetParameterValueRequest,
+    "SetParameterValueRequest": SetParameterValueRequest,
+}
+
+GTN_MCP_SERVER_PORT = int(os.getenv("GTN_MCP_SERVER_PORT", "9927"))
+
+config_manager = ConfigManager()
+secrets_manager = SecretsManager(config_manager)
+
+
+def main(api_key: str) -> None:
+    """Main entry point for the Griptape Nodes MCP server."""
+    mcp_server_logger = logging.getLogger("griptape_nodes_mcp_server")
+    mcp_server_logger.addHandler(RichHandler(show_time=True, show_path=False, markup=True, rich_tracebacks=True))
+    mcp_server_logger.setLevel(logging.INFO)
+    mcp_server_logger.info("Starting MCP GTN server...")
+
+    # Give these a session ID
+    connection_manager = WebSocketConnectionManager()
+    request_manager = AsyncRequestManager(connection_manager, api_key)
+
+    app = Server("mcp-gtn")
+
+    @app.list_tools()
+    async def list_tools() -> list[Tool]:
+        return [
+            Tool(name=event.__name__, description=event.__doc__, inputSchema=TypeAdapter(event).json_schema())
+            for (name, event) in SUPPORTED_REQUEST_EVENTS.items()
+        ]
+
+    @app.call_tool()
+    async def call_tool(name: str, arguments: dict) -> list[TextContent]:
+        if name not in SUPPORTED_REQUEST_EVENTS:
+            msg = f"Unsupported tool: {name}"
+            raise ValueError(msg)
+
+        request_payload = SUPPORTED_REQUEST_EVENTS[name](**arguments)
+
+        await request_manager.connect()
+        result = await request_manager.create_request_event(
+            request_payload.__class__.__name__, request_payload.__dict__, timeout_ms=5000
+        )
+        mcp_server_logger.debug("Got result: %s", result)
+
+        return [TextContent(type="text", text=json.dumps(result))]
+
+    # Create the session manager with our app and event store
+    session_manager = StreamableHTTPSessionManager(
+        app=app,
+    )
+
+    @contextlib.asynccontextmanager
+    async def lifespan(_: FastAPI) -> AsyncIterator[None]:
+        """Context manager for managing session manager lifecycle."""
+        async with session_manager.run():
+            mcp_server_logger.info("GTN MCP server started with StreamableHTTP session manager!")
+            try:
+                yield
+            finally:
+                mcp_server_logger.info("GTN MCP server shutting down...")
+
+    # Create an ASGI application using the transport
+    mcp_server_app = FastAPI(lifespan=lifespan)
+
+    # ASGI handler for streamable HTTP connections
+    async def handle_streamable_http(scope: Scope, receive: Receive, send: Send) -> None:
+        await session_manager.handle_request(scope, receive, send)
+
+    mcp_server_app.mount("/mcp", app=handle_streamable_http)
+
+    uvicorn.run(mcp_server_app, host="127.0.0.1", port=GTN_MCP_SERVER_PORT)

griptape_nodes/mcp_server/ws_request_manager.py

@@ -0,0 +1,268 @@
+import asyncio
+import contextlib
+import json
+import logging
+import os
+import uuid
+from collections.abc import Callable
+from typing import Any, Generic, TypeVar
+from urllib.parse import urljoin
+
+import websockets
+
+logger = logging.getLogger("griptape_nodes_mcp_server")
+
+T = TypeVar("T")
+
+
+class WebSocketConnectionManager:
+    """Python equivalent of WebSocketConnectionManager in TypeScript."""
+
+    def __init__(
+        self,
+        websocket_url: str = urljoin(
+            os.getenv("GRIPTAPE_NODES_API_BASE_URL", "https://api.nodes.griptape.ai").replace("http", "ws"),
+            "/ws/engines/events?version=v2",
+        ),
+    ):
+        self.websocket_url = websocket_url
+        self.websocket: Any = None
+        self.connected = False
+        self.event_handlers: dict[str, list[Callable]] = {}
+        self.request_handlers: dict[str, tuple[Callable, Callable]] = {}
+        self._process_task: asyncio.Task | None = None
+
+    async def send(self, data: dict[str, Any]) -> None:
+        """Send a message to the WebSocket server."""
+        if not self.websocket:
+            msg = "Not connected to WebSocket server"
+            raise ConnectionError(msg)
+
+        try:
+            message = json.dumps(data)
+            await self.websocket.send(message)
+            logger.debug("📤 Sent message: %s", message)
+        except Exception as e:
+            logger.error("Failed to send message: %s", e)
+            raise
+
+    async def _process_messages(self) -> None:
+        """Process incoming WebSocket messages."""
+        if not self.websocket:
+            logger.warning("WebSocket is not connected, cannot process messages")
+            return
+
+        try:
+            async for message in self.websocket:
+                try:
+                    data = json.loads(message)
+                    logger.debug("📥 Received message: %s", message)
+                    await self._handle_message(data)
+                except json.JSONDecodeError:
+                    logger.error("Failed to parse message: %s", message)
+                except Exception as e:
+                    logger.error("Error processing message: %s", e)
+        except websockets.exceptions.ConnectionClosed:
+            logger.warning("WebSocket connection closed")
+            self.connected = False
+        except asyncio.CancelledError:
+            # Task was cancelled, just exit
+            pass
+        except Exception as e:
+            logger.error("Error in message processing loop: %s", e)
+            self.connected = False
+
+    async def _handle_message(self, data: dict[str, Any]) -> None:
+        request = data.get("payload", {}).get("request", {})
+        request_id = request.get("request_id")
+
+        if request_id and request_id in self.request_handlers:
+            success_handler, failure_handler = self.request_handlers[request_id]
+            try:
+                if data.get("type") == "success_result":
+                    success_handler(data, request)
+                else:
+                    failure_handler(data, request)
+            except Exception as e:
+                logger.error("Error in request handler: %s", e)
+
+    def subscribe_to_request_event(
+        self, success_handler: Callable[[Any, Any], None], failure_handler: Callable[[Any, Any], None]
+    ) -> str:
+        """Subscribe to a request-response event."""
+        request_id = str(uuid.uuid4())
+        self.request_handlers[request_id] = (success_handler, failure_handler)
+        return request_id
+
+    def unsubscribe_from_request_event(self, request_id: str) -> None:
+        """Unsubscribe from a request-response event."""
+        if request_id in self.request_handlers:
+            del self.request_handlers[request_id]
+
+
+class AsyncRequestManager(Generic[T]):  # noqa: UP046
+    def __init__(self, connection_manager: WebSocketConnectionManager, api_key: str):
+        self.connection_manager = connection_manager
+        self.api_key = api_key
+
+    async def _subscribe_to_topic(self, topic: str) -> None:
+        """Subscribe to a specific topic in the message bus."""
+        if self.connection_manager.websocket is None:
+            logger.warning("WebSocket connection not available for subscribing to topic")
+            return
+
+        try:
+            body = {"type": "subscribe", "topic": topic, "payload": {}}
+            await self.connection_manager.websocket.send(json.dumps(body))
+            logger.debug("Subscribed to topic: %s", topic)
+        except websockets.exceptions.WebSocketException as e:
+            logger.error("Error subscribing to topic %s: %s", topic, e)
+        except Exception as e:
+            logger.error("Unexpected error while subscribing to topic %s: %s", topic, e)
+
+    async def _subscribe_to_topics(self) -> None:
+        from griptape_nodes.retained_mode.managers.session_manager import SessionManager
+        from griptape_nodes.retained_mode.utils.engine_identity import EngineIdentity
+
+        # Subscribe to response topic (engine discovery)
+        await self._subscribe_to_topic("response")
+
+        # Get engine ID and subscribe to engine_id/response
+        engine_id = EngineIdentity.get_engine_id()
+        if engine_id:
+            await self._subscribe_to_topic(f"engines/{engine_id}/response")
+        else:
+            logger.warning("Engine ID not available for subscription")
+
+        # Get session ID and subscribe to session_id/response if available
+        session_id = SessionManager.get_saved_session_id()
+        if session_id:
+            topic = f"sessions/{session_id}/response"
+            await self._subscribe_to_topic(topic)
+        else:
+            logger.info("No session ID available for subscription")
+
+    async def connect(self) -> None:
+        """Connect to the WebSocket server."""
+        from griptape_nodes.app.app import _create_websocket_connection
+
+        try:
+            self.connection_manager.websocket = await _create_websocket_connection(self.api_key)
+            logger.debug("🟢 WebSocket connection established: %s", self.connection_manager.websocket)
+
+            await self._subscribe_to_topics()
+
+            # Start processing messages
+            self.connection_manager._process_task = asyncio.create_task(self.connection_manager._process_messages())
+
+        except Exception as e:
+            self.connection_manager.connected = False
+            logger.error("🔴 WebSocket connection failed: %s", str(e))
+            msg = f"Failed to connect to WebSocket: {e!s}"
+            raise ConnectionError(msg) from e
+
+    async def disconnect(self) -> None:
+        """Disconnect from the WebSocket server."""
+        if self.connection_manager.websocket:
+            await self.connection_manager.websocket.close()
+            self.connection_manager.websocket = None
+        self.connection_manager.connected = False
+
+        # Cancel processing task if it's running
+        if self.connection_manager._process_task:
+            self.connection_manager._process_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self.connection_manager._process_task
+            self.connection_manager._process_task = None
+
+        logger.debug("WebSocket disconnected")
+
+    async def send_api_message(self, data: dict[str, Any]) -> None:
+        """Send a message to the API via WebSocket."""
+        try:
+            await self.connection_manager.send(data)
+        except ConnectionError as e:
+            logger.error("Failed to send API message: %s", e)
+            raise
+        except Exception as e:
+            logger.error("Unexpected error sending API message: %s", e)
+            raise
+
+    async def create_event(self, request_type: str, payload: dict[str, Any]) -> None:
+        """Send an event to the API without waiting for a response."""
+        from griptape_nodes.app.app import _determine_request_topic
+
+        logger.debug("📝 Creating Event: %s - %s", request_type, json.dumps(payload))
+
+        data = {"event_type": "EventRequest", "request_type": request_type, "request": payload}
+        topic = _determine_request_topic()
+
+        request_data = {"payload": data, "type": data["event_type"], "topic": topic}
+
+        if not request_data["payload"]["request"].get("request_id"):
+            request_data["payload"]["request"]["request_id"] = ""
+
+        await self.send_api_message(request_data)
+
+    async def create_request_event(
+        self, request_type: str, payload: dict[str, Any], timeout_ms: int | None = None
+    ) -> T:
+        """Send a request and wait for its response.
+
+        Args:
+            request_type: Type of request to send
+            payload: Data to send with the request
+            timeout_ms: Optional timeout in milliseconds
+
+        Returns:
+            The response data
+
+        Raises:
+            asyncio.TimeoutError: If the request times out
+            Exception: If the request fails
+        """
+        # Create a future that will be resolved when the response arrives
+        response_future = asyncio.Future()
+
+        # Convert timeout from milliseconds to seconds for asyncio
+        timeout_sec = timeout_ms / 1000 if timeout_ms else None
+
+        # Define handlers that will resolve/reject the future
+        def success_handler(response: Any, _: Any) -> None:
+            if not response_future.done():
+                result = response.get("payload", {}).get("result", "Success")
+                logger.debug("✅ Request succeeded: %s", result)
+                response_future.set_result(result)
+
+        def failure_handler(response: Any, _: Any) -> None:
+            if not response_future.done():
+                error = (
+                    response.get("payload", {}).get("result", {}).get("exception", "Unknown error") or "Unknown error"
+                )
+                logger.error("❌ Request failed: %s", error)
+                response_future.set_exception(Exception(error))
+
+        # Generate request ID and subscribe
+        request_id = self.connection_manager.subscribe_to_request_event(success_handler, failure_handler)
+        payload["request_id"] = request_id
+
+        logger.debug("🚀 Request (%s): %s %s", request_id, request_type, json.dumps(payload))
+
+        try:
+            # Send the event
+            await self.create_event(request_type, payload)
+
+            # Wait for the response with optional timeout
+            if timeout_sec:
+                return await asyncio.wait_for(response_future, timeout=timeout_sec)
+            return await response_future
+
+        except TimeoutError:
+            logger.error("Request timed out after %s ms: %s", timeout_ms, request_id)
+            self.connection_manager.unsubscribe_from_request_event(request_id)
+            raise
+
+        except Exception as e:
+            logger.error("Request failed: %s - %s", request_id, e)
+            self.connection_manager.unsubscribe_from_request_event(request_id)
+            raise

griptape_nodes/node_library/workflow_registry.py

@@ -13,7 +13,7 @@ from griptape_nodes.utils.metaclasses import SingletonMeta
 
 
 class WorkflowMetadata(BaseModel):
-    LATEST_SCHEMA_VERSION: ClassVar[str] = "0.6.0"
+    LATEST_SCHEMA_VERSION: ClassVar[str] = "0.6.1"
 
     name: str
     schema_version: str

griptape_nodes/retained_mode/events/app_events.py

@@ -82,6 +82,12 @@ class AppInitializationComplete(AppPayload):
     """Application initialization completed successfully. All subsystems ready."""
 
 
+@dataclass
+@PayloadRegistry.register
+class AppConnectionEstablished(AppPayload):
+    """Notification that a connection to the API has been established."""
+
+
 @dataclass
 @PayloadRegistry.register
 class GetEngineVersionRequest(RequestPayload):
@@ -192,7 +198,6 @@ class EngineHeartbeatResultSuccess(ResultPayloadSuccess):
         instance_region: Cloud instance region (None if not applicable)
         instance_provider: Cloud provider name (None if not applicable)
         deployment_type: Type of deployment (None if not applicable)
-        public_ip: Public IP address (None if not available)
         current_workflow: Name of active workflow (None if none)
         workflow_file_path: Path to workflow file (None if none)
         has_active_flow: Whether there's an active flow running
@@ -208,7 +213,6 @@ class EngineHeartbeatResultSuccess(ResultPayloadSuccess):
     instance_region: str | None
     instance_provider: str | None
     deployment_type: str | None
-    public_ip: str | None
     current_workflow: str | None
     workflow_file_path: str | None
     has_active_flow: bool