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

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

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