38class Session:
 39    """
 40    The main class for managing an execution session.
 41
 42    This class is responsible for setting up all the necessary components for running a Railtracks execution, including the coordinator, publisher, and state management.
 43
 44    For the configuration parameters of the setting. It will follow this precedence:
 45    1. The parameters in the `Session` constructor.
 46    2. The parameters in global context variables.
 47    3. The default values.
 48
 49    Default Values:
 50    - `name`: None
 51    - `timeout`: 150.0 seconds
 52    - `end_on_error`: False
 53    - `broadcast_callback`: None (no callback for broadcast messages)
 54    - `prompt_injection`: True (the prompt will be automatically injected from context variables)
 55    - `save_state`: True (the state of the execution will be saved to a file at the end of the run in the `.railtracks/data/sessions/` directory)
 56
 57
 58    Args:
 59        name (str | None, optional): Optional name for the session. This name will be included in the saved state file if `save_state` is True.
 60        context (Dict[str, Any], optional): A dictionary of global context variables to be used during the execution.
 61        flow_name (str | None, optional): The name of the flow this session is associated with.
 62        flow_id (str | None, optional): The unique identifier of the flow this session is associated with.
 63        timeout (float, optional): The maximum number of seconds to wait for a response to your top-level request.
 64        end_on_error (bool, optional): If True, the execution will stop when an exception is encountered.
 65        broadcast_callback (Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None, optional): A callback function that will be called with the broadcast messages.
 66        prompt_injection (bool, optional): If True, the prompt will be automatically injected from context variables.
 67        save_state (bool, optional): If True, the state of the execution will be saved to a file at the end of the run in the `.railtracks/data/sessions/` directory.
 68    """
 69
 70    def __init__(
 71        self,
 72        context: Dict[str, Any] | None = None,
 73        *,
 74        flow_name: str | None = None,
 75        flow_id: str | None = None,
 76        name: str | None = None,
 77        timeout: float | None = None,
 78        end_on_error: bool | None = None,
 79        broadcast_callback: (
 80            Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
 81        ) = None,
 82        prompt_injection: bool | None = None,
 83        save_state: bool | None = None,
 84        payload_callback: Callable[[dict[str, Any]], None] | None = None,
 85    ):
 86        # first lets read from defaults if nessecary for the provided input config
 87
 88        if flow_name is None:
 89            warnings.warn(
 90                "Sessions should be tied to a flow for better observability and state management. Please use the Flow object to create and manage your sessions (see __ for more details). This warning will become an error in future versions.",
 91                DeprecationWarning,
 92            )
 93
 94        self.executor_config = self.global_config_precedence(
 95            timeout=timeout,
 96            end_on_error=end_on_error,
 97            broadcast_callback=broadcast_callback,
 98            prompt_injection=prompt_injection,
 99            save_state=save_state,
100            payload_callback=payload_callback,
101        )
102
103        if context is None:
104            context = {}
105
106        self.name = name
107        self.flow_name = flow_name
108        self.flow_id = flow_id
109
110        self.publisher: RTPublisher = RTPublisher()
111
112        self._identifier = str(uuid.uuid4())
113
114        executor_info = ExecutionInfo.create_new()
115        self.coordinator = Coordinator(
116            execution_modes={"async": AsyncioExecutionStrategy()}
117        )
118        self.rt_state = RTState(
119            executor_info, self.executor_config, self.coordinator, self.publisher
120        )
121
122        self.coordinator.start(self.publisher)
123        self._setup_subscriber()
124        register_globals(
125            session_id=self._identifier,
126            rt_publisher=self.publisher,
127            parent_id=None,
128            executor_config=self.executor_config,
129            global_context_vars=context,
130        )
131
132        self._start_time = time.time()
133
134        logger.debug("Session %s is initialized" % self._identifier)
135
136    @classmethod
137    def global_config_precedence(
138        cls,
139        timeout: float | None,
140        end_on_error: bool | None,
141        broadcast_callback: (
142            Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
143        ),
144        prompt_injection: bool | None,
145        save_state: bool | None,
146        payload_callback: Callable[[dict[str, Any]], None] | None,
147    ) -> ExecutorConfig:
148        """
149        Uses the following precedence order to determine the configuration parameters:
150        1. The parameters in the method parameters.
151        2. The parameters in global context variables.
152        3. The default values.
153        """
154        global_executor_config = get_global_config()
155
156        return global_executor_config.precedence_overwritten(
157            timeout=timeout,
158            end_on_error=end_on_error,
159            subscriber=broadcast_callback,
160            prompt_injection=prompt_injection,
161            save_state=save_state,
162            payload_callback=payload_callback,
163        )
164
165    def __enter__(self):
166        return self
167
168    def __exit__(self, exc_type, exc_val, exc_tb):
169        if self.executor_config.save_state:
170            try:
171                railtracks_home = os.environ.get("RAILTRACKS_HOME", ".railtracks")
172                railtracks_dir = Path(railtracks_home)
173                sessions_dir = railtracks_dir / "data" / "sessions"
174                sessions_dir.mkdir(
175                    parents=True, exist_ok=True
176                )  # Creates directory structure if doesn't exist, skips otherwise.
177
178                # Try to create file path with name, fallback to identifier only if there's an issue
179                if self.flow_name is not None:
180                    name = self.flow_name
181                elif self.name is not None:
182                    name = self.name
183                else:
184                    name = ""
185
186                try:
187                    file_path = sessions_dir / f"{name}_{self._identifier}.json"
188                    file_path.touch()
189                except FileNotFoundError:
190                    logger.warning(
191                        get_message(
192                            ExceptionMessageKey.INVALID_SESSION_FILE_NAME_WARN
193                        ).format(name=name, identifier=self._identifier)
194                    )
195                    file_path = sessions_dir / f"{self._identifier}.json"
196
197                logger.info("Saving execution info to %s" % file_path)
198
199                file_path.write_text(json.dumps(self.payload()))
200
201            except Exception as e:
202                logger.error(
203                    "Error while saving to execution info to file",
204                    exc_info=e,
205                )
206        try:
207            if self.executor_config.payload_callback is not None:
208                self.executor_config.payload_callback(self.payload())
209        except Exception:
210            # TODO: add logging here.
211            pass
212
213        self._close()
214
215    def _setup_subscriber(self):
216        """
217        Prepares and attaches the saved broadcast_callback to the publisher attached to this runner.
218        """
219
220        if self.executor_config.subscriber is not None:
221            self.publisher.subscribe(
222                stream_subscriber(self.executor_config.subscriber),
223                name="Streaming Subscriber",
224            )
225
226    def _close(self):
227        """
228        Closes the runner and cleans up all resources.
229
230        - Shuts down the state object
231        - Deletes all the global variables that were registered in the context
232        """
233        # FIX: Resource leak - publisher background task wasn't being shut down on Session exit
234        # VISION: Session owns publisher lifecycle and must clean up all resources when exiting
235        if self.publisher.is_running():
236            try:
237                # Signal shutdown by setting the flag - the loop will check this and exit
238                self.publisher._running = False
239
240                # Try to cancel the background task if it exists and isn't done
241                if (
242                    self.publisher.pub_loop is not None
243                    and not self.publisher.pub_loop.done()
244                ):
245                    try:
246                        # Cancel the task - it will check _running and exit naturally
247                        self.publisher.pub_loop.cancel()
248                    except Exception:
249                        # Task might be done or in a different loop, that's okay
250                        pass
251            except Exception:
252                # If shutdown fails for any reason, log it but don't crash
253                logger.warning(
254                    "Failed to shutdown publisher during Session cleanup. "
255                    "This may indicate a resource leak.",
256                    exc_info=True,
257                )
258
259        self.rt_state.shutdown()
260
261        delete_globals()
262        # by deleting all of the state variables we are ensuring that the next time we create a runner it is fresh
263
264    @property
265    def info(self) -> ExecutionInfo:
266        """
267        Returns the current state of the runner.
268
269        This is useful for debugging and viewing the current state of the run.
270        """
271        return self.rt_state.info
272
273    def payload(self) -> Dict[str, Any]:
274        """
275        Gets the complete json payload tied to this session.
276
277        The outputted json schema is maintained in (link here)
278        """
279        info = self.info
280
281        run_list = info.graph_serialization()
282
283        full_dict = {
284            "flow_name": self.flow_name,
285            "flow_id": self.flow_id,
286            "session_id": self._identifier,
287            "session_name": self.name,
288            "start_time": self._start_time,
289            "end_time": time.time(),
290            "runs": run_list,
291        }
292
293        return json.loads(json.dumps(full_dict))

350def session(
351    func: Callable[_P, Coroutine[Any, Any, _TOutput]] | None = None,
352    *,
353    name: str | None = None,
354    context: Dict[str, Any] | None = None,
355    timeout: float | None = None,
356    end_on_error: bool | None = None,
357    broadcast_callback: (
358        Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
359    ) = None,
360    prompt_injection: bool | None = None,
361    save_state: bool | None = None,
362) -> (
363    Callable[_P, Coroutine[Any, Any, Tuple[_TOutput, Session]]]
364    | Callable[
365        [Callable[_P, Coroutine[Any, Any, _TOutput]]],
366        Callable[_P, Coroutine[Any, Any, Tuple[_TOutput, Session]]],
367    ]
368):
369    """
370    This decorator automatically creates and manages a Session context for the decorated function,
371    allowing async functions to use Railtracks operations without manually managing the session lifecycle.
372
373    Can be used as:
374    - @session (without parentheses) - uses default settings
375    - @session() (with empty parentheses) - uses default settings
376    - @session(name="my_task", timeout=30) (with configuration parameters)
377
378    When using this decorator, the function returns a tuple containing:
379    1. The original function's return value
380    2. The Session object used during execution
381
382    This allows access to session information (like execution state, logs, etc.) after the function completes,
383    while maintaining the simplicity of decorator usage.
384
385    Args:
386        name (str | None, optional): Optional name for the session. This name will be included in the saved state file if `save_state` is True.
387        context (Dict[str, Any], optional): A dictionary of global context variables to be used during the execution.
388        timeout (float, optional): The maximum number of seconds to wait for a response to your top-level request.
389        end_on_error (bool, optional): If True, the execution will stop when an exception is encountered.
390        broadcast_callback (Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None, optional): A callback function that will be called with the broadcast messages.
391        prompt_injection (bool, optional): If True, the prompt will be automatically injected from context variables.
392        save_state (bool, optional): If True, the state of the execution will be saved to a file at the end of the run in the `.railtracks/data/sessions/` directory.
393
394    Returns:
395        When used as @session (without parentheses): Returns the decorated function that returns (result, session).
396        When used as @session(...) (with parameters): Returns a decorator function that takes an async function
397        and returns a new async function that returns (result, session).
398    """
399
400    def decorator(
401        target_func: Callable[_P, Coroutine[Any, Any, _TOutput]],
402    ) -> Callable[_P, Coroutine[Any, Any, Tuple[_TOutput, Session]]]:
403        # Validate that the decorated function is async
404        if not inspect.iscoroutinefunction(target_func):
405            raise TypeError(
406                f"@session decorator can only be applied to async functions. "
407                f"Function '{target_func.__name__}' is not async. "
408                f"Add 'async' keyword to your function definition."
409            )
410
411        @wraps(target_func)
412        async def wrapper(
413            *args: _P.args, **kwargs: _P.kwargs
414        ) -> Tuple[_TOutput, Session]:
415            session_obj = Session(
416                context=context,
417                timeout=timeout,
418                end_on_error=end_on_error,
419                broadcast_callback=broadcast_callback,
420                name=name,
421                prompt_injection=prompt_injection,
422                save_state=save_state,
423            )
424
425            with session_obj:
426                result = await target_func(*args, **kwargs)
427                return result, session_obj
428
429        return wrapper
430
431    # If used as @session without parentheses
432    if func is not None:
433        return decorator(func)
434
435    # If used as @session(...)
436    return decorator

 59async def call(
 60    node_: Callable[_P, Node[_TOutput]] | RTFunction[_P, _TOutput],
 61    *args: _P.args,
 62    **kwargs: _P.kwargs,
 63) -> _TOutput:
 64    """
 65    Call a node from within a node inside the framework. This will return a coroutine that you can interact with
 66    in whatever way using async/await logic.
 67
 68    Usage:
 69    ```python
 70    # for sequential operation
 71    result = await call(NodeA, "hello world", 42)
 72
 73    # for parallel operation
 74    tasks = [call(NodeA, "hello world", i) for i in range(10)]
 75    results = await asyncio.gather(*tasks)
 76    ```
 77
 78    Args:
 79        node: The node type you would like to create. This could be a function decorated with `@function_node`, a function, or a Node instance.
 80        *args: The arguments to pass to the node
 81        **kwargs: The keyword arguments to pass to the node
 82    """
 83    node: Callable[_P, Node[_TOutput]]
 84    # this entire section is a bit of a typing nightmare becuase all overloads we provide.
 85    if isinstance(node_, FunctionType):
 86        node = extract_node_from_function(node_)
 87    else:
 88        node = node_
 89    # if the context is none then we will need to create a wrapper for the state object to work with.
 90    if not is_context_present():
 91        # we have to use lazy import here to prevent a circular import issue. This is a must have unfortunately.
 92        from railtracks import Session
 93
 94        with Session():
 95            result = await _start(node, args=args, kwargs=kwargs)
 96            return result
 97
 98    # if the context is not active then we know this is the top level request
 99    if not is_context_active():
100        result = await _start(node, args=args, kwargs=kwargs)
101        return result
102
103    # if the context is active then we can just run the node
104    result = await _run(node, args=args, kwargs=kwargs)
105    return result

 6async def broadcast(item: str):
 7    """
 8    Streams the given message
 9
10    This will trigger the broadcast_callback callback you have already provided.
11
12    Args:
13        item (str): The item you want to stream.
14    """
15    publisher = get_publisher()
16
17    await publisher.publish(Streaming(node_id=get_parent_id(), streamed_object=item))

27async def call_batch(
28    node: Callable[..., Node[_TOutput]]
29    | Callable[..., _TOutput]
30    | _AsyncNodeAttachedFunc[_P, _TOutput]
31    | _SyncNodeAttachedFunc[_P, _TOutput],
32    *iterables: Iterable[Any],
33    return_exceptions: bool = True,
34):
35    """
36    Complete a node over multiple iterables, allowing for parallel execution.
37
38    Note the results will be returned in the order of the iterables, not the order of completion.
39
40    If one of the nodes returns an exception, the thrown exception will be included as a response.
41
42    Args:
43        node: The node type to create.
44        *iterables: The iterables to map the node over.
45        return_exceptions: If True, exceptions will be returned as part of the results.
46            If False, exceptions will be raised immediately, and you will lose access to the results.
47            Defaults to true.
48
49    Returns:
50        An iterable of results from the node.
51
52    Usage:
53        ```python
54        results = await batch(NodeA, ["hello world"] * 10)
55        for result in results:
56            handle(result)
57        ```
58    """
59    # this is big typing disaster but there is no way around it. Try if if you want to.
60    contracts = [call(node, *args) for args in zip(*iterables)]
61
62    results = await asyncio.gather(*contracts, return_exceptions=return_exceptions)
63    return results

 19class ExecutionInfo:
 20    """
 21    A class that contains the full details of the state of a run at any given point in time.
 22
 23    The class is designed to be used as a snapshot of state that can be used to display the state of the run, or to
 24    create a graphical representation of the system.
 25    """
 26
 27    def __init__(
 28        self,
 29        request_forest: RequestForest,
 30        node_forest: NodeForest,
 31        stamper: StampManager,
 32    ):
 33        self.request_forest = request_forest
 34        self.node_forest = node_forest
 35        self.stamper = stamper
 36
 37    @classmethod
 38    def default(cls) -> ExecutionInfo:
 39        """Creates a new "empty" instance of the ExecutionInfo class with the default values."""
 40        return cls.create_new()
 41
 42    @classmethod
 43    def create_new(
 44        cls,
 45    ) -> ExecutionInfo:
 46        """
 47        Creates a new empty instance of state variables with the provided executor configuration.
 48
 49        """
 50        request_heap = RequestForest()
 51        node_heap = NodeForest()
 52        stamper = StampManager()
 53
 54        return ExecutionInfo(
 55            request_forest=request_heap,
 56            node_forest=node_heap,
 57            stamper=stamper,
 58        )
 59
 60    @property
 61    def answer(self):
 62        """Convenience method to access the answer of the run."""
 63        return self.request_forest.answer
 64
 65    @property
 66    def all_stamps(self) -> List[Stamp]:
 67        """Convenience method to access all the stamps of the run."""
 68        return self.stamper.all_stamps
 69
 70    @property
 71    def name(self):
 72        """
 73        Gets the name of the graph by pulling the name of the insertion request. It will raise a ValueError if the insertion
 74        request is not present or there are multiple insertion requests.
 75        """
 76        insertion_requests = self.insertion_requests
 77
 78        # The name is only defined for the length of 1.
 79        # NOTE: Maybe we should send a warning once to user in other cases.
 80        if len(insertion_requests) != 1:
 81            return None
 82
 83        i_r = insertion_requests[0]
 84
 85        return self.node_forest.get_node_type(i_r.sink_id).name()
 86
 87    @property
 88    def insertion_requests(self):
 89        """A convenience method to access all the insertion requests of the run."""
 90        return self.request_forest.insertion_request
 91
 92    def _get_info(self, ids: List[str] | str | None = None) -> ExecutionInfo:
 93        """
 94        Gets a subset of the current state based on the provided node ids. It will contain all the children of the provided node ids
 95
 96        Note: If no ids are provided, the full state is returned.
 97
 98        Args:
 99            ids (List[str] | str | None): A list of node ids to filter the state by. If None, the full state is returned.
100
101        Returns:
102            ExecutionInfo: A new instance of ExecutionInfo containing only the children of the provided ids.
103
104        """
105        if ids is None:
106            return self
107        else:
108            # firstly lets
109            if isinstance(ids, str):
110                ids = [ids]
111
112            # we need to quickly check to make sure these ids are valid
113            for identifier in ids:
114                if identifier not in self.request_forest:
115                    raise ValueError(
116                        f"Identifier '{identifier}' not found in the current state."
117                    )
118
119            new_node_forest, new_request_forest = create_sub_state_info(
120                self.node_forest.heap(),
121                self.request_forest.heap(),
122                ids,
123            )
124            return ExecutionInfo(
125                node_forest=new_node_forest,
126                request_forest=new_request_forest,
127                stamper=self.stamper,
128            )
129
130    def _to_graph(self) -> Tuple[List[Vertex], List[Edge]]:
131        """
132        Converts the current state into its graph representation.
133
134        Returns:
135            List[Node]: An iterable of nodes in the graph.
136            List[Edge]: An iterable of edges in the graph.
137        """
138        return self.node_forest.to_vertices(), self.request_forest.to_edges()
139
140    def graph_serialization(self) -> dict[str, Any]:
141        """
142                Creates a string (JSON) representation of this info object designed to be used to construct a graph for this
143                info object.
144
145                Some important notes about its structure are outlined below:
146                - The `nodes` key contains a list of all the nodes in the graph, represented as `Vertex` objects.
147                - The `edges` key contains a list of all the edges in the graph, represented as `Edge` objects.
148                - The `stamps` key contains an ease of use list of all the stamps associated with the run, represented as `Stamp` objects.
149
150                - The "nodes" and "requests" key will be outlined with normal graph details like connections and identifiers in addition to a loose details object.
151                - However, both will carry an addition param called "stamp" which is a timestamp style object.
152                - They also will carry a "parent" param which is a recursive structure that allows you to traverse the graph in time.
153
154
155        ```
156        """
157        parent_nodes = [x.identifier for x in self.insertion_requests]
158
159        infos = [self._get_info(parent_node) for parent_node in parent_nodes]
160
161        runs = []
162
163        for info, parent_node_id in zip(infos, parent_nodes):
164            insertion_requests = info.request_forest.insertion_request
165
166            assert len(insertion_requests) == 1
167            parent_request = insertion_requests[0]
168
169            all_parents = parent_request.get_all_parents()
170
171            start_time = all_parents[-1].stamp.time
172
173            assert len([x for x in all_parents if x.status == "Completed"]) <= 1
174            end_time = None
175            for req in all_parents:
176                if req.status in ["Completed", "Failed"]:
177                    end_time = req.stamp.time
178                    break
179
180            entry = {
181                "name": info.name,
182                "run_id": parent_node_id,
183                "nodes": info.node_forest.to_vertices(),
184                "status": parent_request.status,
185                "edges": info.request_forest.to_edges(),
186                "steps": _get_stamps_from_forests(
187                    info.node_forest, info.request_forest
188                ),
189                "start_time": start_time,
190                "end_time": end_time,
191            }
192            runs.append(entry)
193
194        return json.loads(
195            json.dumps(
196                runs,
197                cls=RTJSONEncoder,
198            )
199        )

 8class ExecutorConfig:
 9    def __init__(
10        self,
11        *,
12        timeout: float | None = None,
13        end_on_error: bool = False,
14        broadcast_callback: (
15            Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
16        ) = None,
17        prompt_injection: bool = True,
18        save_state: bool = True,
19        payload_callback: Callable[[dict[str, Any]], None] | None = None,
20    ):
21        """
22        ExecutorConfig is special configuration object designed to allow customization of the executor in the RT system.
23
24        Args:
25            timeout (float | None): The maximum number of seconds to wait for a response to your top level request. Pass None (or omit) to disable the timeout entirely.
26            end_on_error (bool): If true, the executor will stop execution when an exception is encountered.
27            broadcast_callback (Callable or Coroutine): A function or coroutine that will handle streaming messages.
28            prompt_injection (bool): If true, prompts can be injected with global context
29            save_state (bool): If true, the state of the executor will be saved to disk.
30        """
31        self.timeout = timeout
32        self.end_on_error = end_on_error
33        self.subscriber = broadcast_callback
34        self.prompt_injection = prompt_injection
35        # During test runs, disable save_state by default unless RAILTRACKS_ALLOW_PERSISTENCE is set
36        self._user_save_state = save_state
37
38        self.payload_callback = payload_callback
39
40    # this is done because if we try to lock the save_state in init
41    # later when we want to allow a few tests to actually run persistance, they wont be able to do so
42    @property
43    def save_state(self) -> bool:
44        if os.getenv("RAILTRACKS_TEST_MODE") and not os.getenv(
45            "RAILTRACKS_ALLOW_PERSISTENCE"
46        ):
47            return False
48        return self._user_save_state
49
50    def precedence_overwritten(
51        self,
52        *,
53        timeout: float | None = None,
54        end_on_error: bool | None = None,
55        subscriber: (
56            Callable[[str], None] | Callable[[str], Coroutine[None, None, None]] | None
57        ) = None,
58        prompt_injection: bool | None = None,
59        save_state: bool | None = None,
60        payload_callback: Callable[[dict[str, Any]], None] | None = None,
61    ):
62        """
63        If any of the parameters are provided (not None), it will create a new update the current instance with the new values and return a deep copied reference to it.
64        """
65        return ExecutorConfig(
66            timeout=timeout,
67            end_on_error=end_on_error
68            if end_on_error is not None
69            else self.end_on_error,
70            broadcast_callback=subscriber
71            if subscriber is not None
72            else self.subscriber,
73            prompt_injection=prompt_injection
74            if prompt_injection is not None
75            else self.prompt_injection,
76            save_state=save_state if save_state is not None else self.save_state,
77            payload_callback=payload_callback
78            if payload_callback is not None
79            else self.payload_callback,
80        )
81
82    def __repr__(self):
83        return (
84            f"ExecutorConfig(timeout={self.timeout}, end_on_error={self.end_on_error}, "
85            f"prompt_injection={self.prompt_injection}, "
86            f"save_state={self.save_state}, payload_callback={self.payload_callback})"
87        )