agent.py

import logging
import re
import uuid
import json
from typing import Optional, List, Any, Union, AsyncGenerator, Generator, Literal, Dict
import inspect

from langgraph.graph.state import CompiledStateGraph

try:
    from langchain.schema import BaseMessage, SystemMessage, ToolMessage
except ImportError:
    # Langchain >= 1.0.0
    from langchain_core.messages import BaseMessage, SystemMessage, ToolMessage
    
from langchain_core.runnables import RunnableConfig, ensure_config
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.types import Command

from .types import (
    State,
    LangGraphPlatformMessage,
    MessagesInProgressRecord,
    SchemaKeys,
    MessageInProgress,
    RunMetadata,
    LangGraphEventTypes,
    CustomEventNames,
    LangGraphReasoning
)
from .utils import (
    agui_messages_to_langchain,
    DEFAULT_SCHEMA_KEYS,
    filter_object_by_schema_keys,
    get_stream_payload_input,
    langchain_messages_to_agui,
    resolve_reasoning_content,
    resolve_encrypted_reasoning_content,
    resolve_message_content,
    camel_to_snake,
    json_safe_stringify,
    make_json_safe,
    normalize_tool_content
)

from ag_ui.core import (
    EventType,
    CustomEvent,
    MessagesSnapshotEvent,
    RawEvent,
    RunAgentInput,
    RunErrorEvent,
    RunFinishedEvent,
    RunStartedEvent,
    StateDeltaEvent,
    StateSnapshotEvent,
    StepFinishedEvent,
    StepStartedEvent,
    TextMessageContentEvent,
    TextMessageEndEvent,
    TextMessageStartEvent,
    ToolCallArgsEvent,
    ToolCallEndEvent,
    ToolCallStartEvent,
    ToolCallResultEvent,
    ReasoningStartEvent,
    ReasoningMessageStartEvent,
    ReasoningMessageContentEvent,
    ReasoningMessageEndEvent,
    ReasoningEndEvent,
    ReasoningEncryptedValueEvent,
)
from ag_ui.encoder import EventEncoder

ProcessedEvents = Union[
    TextMessageStartEvent,
    TextMessageContentEvent,
    TextMessageEndEvent,
    ReasoningStartEvent,
    ReasoningMessageStartEvent,
    ReasoningMessageContentEvent,
    ReasoningMessageEndEvent,
    ReasoningEndEvent,
    ReasoningEncryptedValueEvent,
    ToolCallStartEvent,
    ToolCallArgsEvent,
    ToolCallEndEvent,
    StateSnapshotEvent,
    StateDeltaEvent,
    MessagesSnapshotEvent,
    RawEvent,
    CustomEvent,
    RunStartedEvent,
    RunFinishedEvent,
    RunErrorEvent,
    StepStartedEvent,
    StepFinishedEvent,
]

logger = logging.getLogger(__name__)

class LangGraphAgent:
    def __init__(self, *, name: str, graph: CompiledStateGraph, description: Optional[str] = None, config:  Union[Optional[RunnableConfig], dict] = None):
        self.name = name
        self.description = description
        self.graph = graph
        self.config = config or {}
        self.messages_in_process: MessagesInProgressRecord = {}
        self.active_run: Optional[RunMetadata] = None
        self.constant_schema_keys = ['messages', 'tools']

    def clone(self) -> "LangGraphAgent":
        return LangGraphAgent(
            name=self.name,
            graph=self.graph,
            description=self.description,
            config=self.config,
        )

    def _dispatch_event(self, event: ProcessedEvents) -> str:
        if event.type == EventType.RAW:
            event.event = make_json_safe(event.event)
        elif event.raw_event:
            event.raw_event = make_json_safe(event.raw_event)

        return event

    async def run(self, input: RunAgentInput) -> AsyncGenerator[str, None]:
        forwarded_props = {}
        if hasattr(input, "forwarded_props") and input.forwarded_props:
            forwarded_props = {
                camel_to_snake(k): v for k, v in input.forwarded_props.items()
            }
        async for event_str in self._handle_stream_events(input.copy(update={"forwarded_props": forwarded_props})):
            yield event_str

    async def _handle_stream_events(self, input: RunAgentInput) -> AsyncGenerator[str, None]:
        thread_id = input.thread_id or str(uuid.uuid4())
        INITIAL_ACTIVE_RUN = {
            "id": input.run_id,
            "thread_id": thread_id,
            "reasoning_process": None,
            "node_name": None,
            "has_function_streaming": False,
            "model_made_tool_call": False,
            "state_reliable": True,
        }
        self.active_run = INITIAL_ACTIVE_RUN

        forwarded_props = input.forwarded_props
        node_name_input = forwarded_props.get('node_name', None) if forwarded_props else None

        self.active_run["manually_emitted_state"] = None

        config = ensure_config(self.config.copy() if self.config else {})
        config["configurable"] = {**(config.get('configurable', {})), "thread_id": thread_id}

        agent_state = await self.graph.aget_state(config)
        resume_input = forwarded_props.get('command', {}).get('resume', None)

        if resume_input is None and thread_id and self.active_run.get("node_name") != "__end__" and self.active_run.get("node_name"):
            self.active_run["mode"] = "continue"
        else:
            self.active_run["mode"] = "start"

        prepared_stream_response = await self.prepare_stream(input=input, agent_state=agent_state, config=config)

        yield self._dispatch_event(
            RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=self.active_run["id"])
        )
        self.handle_node_change(node_name_input)

        # In case of resume (interrupt), re-start resumed step
        if resume_input and self.active_run.get("node_name"):
            for ev in self.handle_node_change(self.active_run.get("node_name")):
                yield ev

        state = prepared_stream_response["state"]
        stream = prepared_stream_response["stream"]
        config = prepared_stream_response["config"]
        events_to_dispatch = prepared_stream_response.get('events_to_dispatch', None)

        if events_to_dispatch is not None and len(events_to_dispatch) > 0:
            for event in events_to_dispatch:
                yield self._dispatch_event(event)
            return

        should_exit = False
        current_graph_state = state

        try:
            async for event in stream:
                subgraphs_stream_enabled = input.forwarded_props.get('stream_subgraphs') if input.forwarded_props else False
                is_subgraph_stream = (subgraphs_stream_enabled and (
                    event.get("event", "").startswith("events") or
                    event.get("event", "").startswith("values")
                ))
                if event["event"] == "error":
                    yield self._dispatch_event(
                        RunErrorEvent(type=EventType.RUN_ERROR, message=event["data"]["message"], raw_event=event)
                    )
                    break

                current_node_name = event.get("metadata", {}).get("langgraph_node")
                event_type = event.get("event")
                self.active_run["id"] = event.get("run_id")
                exiting_node = False

                if event_type == "on_chain_end" and isinstance(
                        event.get("data", {}).get("output"), dict
                ):
                    output = event["data"]["output"]
                    current_graph_state.update(output)
                    exiting_node = self.active_run["node_name"] == current_node_name
                    # If output contains any key outside the protocol-internal set
                    # ("messages", "tools", "ag-ui"), the local current_graph_state
                    # is reliably up-to-date again.
                    if any(k not in ("messages", "tools", "ag-ui") for k in output):
                        self.active_run["state_reliable"] = True

                should_exit = should_exit or (
                        event_type == "on_custom_event" and
                        event["name"] == "exit"
                    )

                if current_node_name and current_node_name != self.active_run.get("node_name"):
                    for ev in self.handle_node_change(current_node_name):
                        yield ev

                # Track whether the current model turn is making a predict_state tool
                # call so we can suppress the model-node exit snapshot.  The model-node
                # exit fires *before* the tool runs, so current_graph_state still
                # carries the previous value — emitting it would wipe predict_state
                # progress on the client.  This applies to every iteration, not just
                # the first.  Note: _handle_single_event uses the same predict_state
                # metadata check to emit the PredictState custom event — keep both
                # sites in sync if the check logic changes.
                if event_type == LangGraphEventTypes.OnChatModelStream.value:
                    chunk = event.get("data", {}).get("chunk") or {}
                    tool_call_chunks = (
                        chunk.get("tool_call_chunks") or []
                        if isinstance(chunk, dict)
                        else getattr(chunk, "tool_call_chunks", None) or []
                    )
                    if tool_call_chunks:
                        first = tool_call_chunks[0]
                        first_name = (
                            first.get("name") if isinstance(first, dict)
                            else getattr(first, "name", None)
                        )
                        if first_name:
                            predict_state_meta = event.get("metadata", {}).get("predict_state", [])
                            tool_used_to_predict_state = any(
                                (p.get("tool") if isinstance(p, dict) else getattr(p, "tool", None)) == first_name
                                for p in predict_state_meta
                            )
                            if tool_used_to_predict_state:
                                self.active_run["model_made_tool_call"] = True

                updated_state = self.active_run.get("manually_emitted_state") or current_graph_state
                has_state_diff = updated_state != state
                if exiting_node or (has_state_diff and not self.get_message_in_progress(self.active_run["id"])):
                    state = updated_state
                    self.active_run["prev_node_name"] = self.active_run["node_name"]
                    current_graph_state.update(updated_state)
                    mmtc = self.active_run.get("model_made_tool_call")
                    state_reliable = self.active_run.get("state_reliable", True)
                    suppressed = exiting_node and (mmtc or not state_reliable)
                    if suppressed:
                        logger.debug(
                            "Suppressing STATE_SNAPSHOT on node exit (node=%s, model_made_tool_call=%s, state_reliable=%s)",
                            self.active_run.get("node_name"), mmtc, state_reliable,
                        )
                        self.active_run["model_made_tool_call"] = False
                        if mmtc:
                            # A predict_state tool call was detected — the tool has
                            # not yet run, so current_graph_state does not yet reflect
                            # the forthcoming state update.
                            self.active_run["state_reliable"] = False
                    else:
                        yield self._dispatch_event(
                            StateSnapshotEvent(
                                type=EventType.STATE_SNAPSHOT,
                                snapshot=self.get_state_snapshot(state),
                                raw_event=event,
                            )
                        )

                yield self._dispatch_event(
                    RawEvent(type=EventType.RAW, event=event)
                )

                async for single_event in self._handle_single_event(event, state):
                    yield single_event

            state = await self.graph.aget_state(config)

            tasks = state.tasks if len(state.tasks) > 0 else None
            interrupts = tasks[0].interrupts if tasks else []

            writes = state.metadata.get("writes", {}) or {}
            node_name = self.active_run["node_name"] if interrupts else next(iter(writes), None)
            next_nodes = state.next or ()
            is_end_node = len(next_nodes) == 0 and not interrupts

            node_name = "__end__" if is_end_node else node_name

            for interrupt in interrupts:
                yield self._dispatch_event(
                    CustomEvent(
                        type=EventType.CUSTOM,
                        name=LangGraphEventTypes.OnInterrupt.value,
                        value=dump_json_safe(interrupt.value),
                        raw_event=interrupt,
                    )
                )

            if self.active_run.get("node_name") != node_name:
                for ev in self.handle_node_change(node_name):
                    yield ev

            state_values = state.values if state.values else state
            yield self._dispatch_event(
                StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=self.get_state_snapshot(state_values))
            )

            snapshot_messages = self._filter_orphan_tool_messages(state_values.get("messages", []))
            yield self._dispatch_event(
                MessagesSnapshotEvent(
                    type=EventType.MESSAGES_SNAPSHOT,
                    messages=langchain_messages_to_agui(snapshot_messages),
                )
            )

            for ev in self.handle_node_change(None):
                yield ev

            yield self._dispatch_event(
                RunFinishedEvent(type=EventType.RUN_FINISHED, thread_id=thread_id, run_id=self.active_run["id"])
            )
            # Reset active run to how it was before the stream started
            self.active_run = INITIAL_ACTIVE_RUN
        except Exception:
            raise


    async def prepare_stream(self, input: RunAgentInput, agent_state: State, config: RunnableConfig):
        state_input = input.state or {}
        messages = input.messages or []
        forwarded_props = input.forwarded_props or {}
        thread_id = input.thread_id

        state_input["messages"] = agent_state.values.get("messages", [])
        self.active_run["current_graph_state"] = agent_state.values.copy()
        langchain_messages = agui_messages_to_langchain(messages)
        state = self.langgraph_default_merge_state(state_input, langchain_messages, input)
        self.active_run["current_graph_state"].update(state)
        config["configurable"]["thread_id"] = thread_id
        interrupts = agent_state.tasks[0].interrupts if agent_state.tasks and len(agent_state.tasks) > 0 else []
        has_active_interrupts = len(interrupts) > 0
        resume_input = forwarded_props.get('command', {}).get('resume', None)

        self.active_run["schema_keys"] = self.get_schema_keys(config)

        non_system_messages = [msg for msg in langchain_messages if not isinstance(msg, SystemMessage)]
        if len(agent_state.values.get("messages", [])) > len(non_system_messages):
            # Only trigger time-travel regeneration if the incoming messages are NOT already
            # in the checkpoint. If they are, this is a continuation (e.g. after CopilotKit
            # intercepted a tool call), not a time-travel edit — regenerating would loop.
            #
            # We exclude ToolMessages from the ID comparison because CopilotKit assigns new
            # IDs to tool results that won't match the placeholder IDs AgentCoreMemorySaver
            # wrote to the checkpoint. Human and AI message IDs are stable across requests
            # and are sufficient to distinguish continuation from time-travel.
            incoming_non_tool_ids = {
                getattr(m, "id", None)
                for m in langchain_messages
                if getattr(m, "id", None) and not isinstance(m, ToolMessage)
            }
            checkpoint_ids = {getattr(m, "id", None) for m in agent_state.values.get("messages", []) if getattr(m, "id", None)}
            is_continuation = bool(incoming_non_tool_ids) and incoming_non_tool_ids.issubset(checkpoint_ids)

            if not is_continuation:
                last_user_message = None
                for i in range(len(langchain_messages) - 1, -1, -1):
                    if isinstance(langchain_messages[i], HumanMessage):
                        last_user_message = langchain_messages[i]
                        break

                if last_user_message:
                    last_user_id = getattr(last_user_message, "id", None)
                    if last_user_id and last_user_id in checkpoint_ids:
                        return await self.prepare_regenerate_stream(
                            input=input,
                            message_checkpoint=last_user_message,
                            config=config
                        )

        events_to_dispatch = []
        if has_active_interrupts and not resume_input:
            events_to_dispatch.append(
                RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=self.active_run["id"])
            )

            for interrupt in interrupts:
                events_to_dispatch.append(
                    CustomEvent(
                        type=EventType.CUSTOM,
                        name=LangGraphEventTypes.OnInterrupt.value,
                        value=dump_json_safe(interrupt.value),
                        raw_event=interrupt,
                    )
                )

            events_to_dispatch.append(
                RunFinishedEvent(type=EventType.RUN_FINISHED, thread_id=thread_id, run_id=self.active_run["id"])
            )
            return {
                "stream": None,
                "state": None,
                "config": None,
                "events_to_dispatch": events_to_dispatch,
            }

        if self.active_run["mode"] == "continue":
            await self.graph.aupdate_state(config, state, as_node=self.active_run.get("node_name"))

        if resume_input:
            if isinstance(resume_input, str):
                try:
                    resume_input = json.loads(resume_input)
                except json.JSONDecodeError:
                    pass  # Keep as string if not valid JSON
            stream_input = Command(resume=resume_input)
        else:
            payload_input = get_stream_payload_input(
                mode=self.active_run["mode"],
                state=state,
                schema_keys=self.active_run["schema_keys"],
            )
            stream_input = {**forwarded_props, **payload_input} if payload_input else None


        subgraphs_stream_enabled = input.forwarded_props.get('stream_subgraphs') if input.forwarded_props else False

        kwargs = self.get_stream_kwargs(
            input=stream_input,
            config=config,
            subgraphs=bool(subgraphs_stream_enabled),
            version="v2",
        )

        stream = self.graph.astream_events(**kwargs)

        return {
            "stream": stream,
            "state": state,
            "config": config
        }

    async def prepare_regenerate_stream( # pylint: disable=too-many-arguments
            self,
            input: RunAgentInput,
            message_checkpoint: HumanMessage,
            config: RunnableConfig
    ):
        tools = input.tools or []
        thread_id = input.thread_id

        time_travel_checkpoint = await self.get_checkpoint_before_message(message_checkpoint.id, thread_id, config)
        if time_travel_checkpoint is None:
            return None

        fork = await self.graph.aupdate_state(
            time_travel_checkpoint.config,
            time_travel_checkpoint.values,
            as_node=time_travel_checkpoint.next[0] if time_travel_checkpoint.next else "__start__"
        )

        stream_input = self.langgraph_default_merge_state(time_travel_checkpoint.values, [message_checkpoint], input)
        subgraphs_stream_enabled = input.forwarded_props.get('stream_subgraphs') if input.forwarded_props else False

        kwargs = self.get_stream_kwargs(
            input=stream_input,
            config=fork,
            subgraphs=bool(subgraphs_stream_enabled),
            version="v2",
        )
        stream = self.graph.astream_events(**kwargs)

        return {
            "stream": stream,
            "state": time_travel_checkpoint.values,
            "config": config
        }

    def get_message_in_progress(self, run_id: str) -> Optional[MessageInProgress]:
        return self.messages_in_process.get(run_id)

    def set_message_in_progress(self, run_id: str, data: MessageInProgress):
        current_message_in_progress = self.messages_in_process.get(run_id) or {}
        self.messages_in_process[run_id] = {
            **current_message_in_progress,
            **data,
        }

    def get_schema_keys(self, config) -> SchemaKeys:
        try:
            input_schema = self.graph.get_input_jsonschema(config)
            output_schema = self.graph.get_output_jsonschema(config)
            config_schema = self.graph.config_schema().schema()

            input_schema_keys = list(input_schema["properties"].keys()) if "properties" in input_schema else []
            output_schema_keys = list(output_schema["properties"].keys()) if "properties" in output_schema else []
            config_schema_keys = list(config_schema["properties"].keys()) if "properties" in config_schema else []
            context_schema_keys = []

            if hasattr(self.graph, "context_schema") and self.graph.context_schema is not None:
                context_schema = self.graph.context_schema().schema()
                context_schema_keys = list(context_schema["properties"].keys()) if "properties" in context_schema else []


            return {
                "input": [*input_schema_keys, *self.constant_schema_keys],
                "output": [*output_schema_keys, *self.constant_schema_keys],
                "config": config_schema_keys,
                "context": context_schema_keys,
            }
        except Exception:
            return {
                "input": self.constant_schema_keys,
                "output": self.constant_schema_keys,
                "config": [],
                "context": [],
            }

    def langgraph_default_merge_state(self, state: State, messages: List[BaseMessage], input: RunAgentInput) -> State:
        if messages and isinstance(messages[0], SystemMessage):
            messages = messages[1:]

        existing_messages: List[LangGraphPlatformMessage] = state.get("messages", [])

        # Fix tool_call args that are strings instead of dicts.
        # This happens when CopilotKit's after_agent restores frontend tool_calls
        # and the checkpoint saves them with string args. Bedrock Converse API
        # requires toolUse.input to be a JSON object (dict).
        for msg in existing_messages:
            if isinstance(msg, AIMessage) and getattr(msg, 'tool_calls', None):
                for tc in msg.tool_calls:
                    if isinstance(tc.get('args'), str):
                        try:
                            tc['args'] = json.loads(tc['args'])
                        except (json.JSONDecodeError, TypeError):
                            tc['args'] = {}

        # Fix orphan ToolMessages injected by patch_orphan_tool_calls:
        # Find the real content from AG-UI messages and replace the fake content.
        # Only scan from the last HumanMessage to the end of existing_messages.
        # Track replaced tool_call_ids so we don't also add the AG-UI duplicate.
        agui_tool_content = {
            m.tool_call_id: m.content
            for m in messages
            if isinstance(m, ToolMessage) and hasattr(m, 'tool_call_id')
        }
        replaced_tool_call_ids = set()
        if agui_tool_content:
            last_human_idx = -1
            for i in range(len(existing_messages) - 1, -1, -1):
                if isinstance(existing_messages[i], HumanMessage):
                    last_human_idx = i
                    break
            if last_human_idx >= 0:
                for i in range(last_human_idx + 1, len(existing_messages)):
                    msg = existing_messages[i]
                    if (
                            isinstance(msg, ToolMessage)
                            and isinstance(msg.content, str)
                            and self._ORPHAN_TOOL_MSG_RE.match(msg.content)
                            and hasattr(msg, 'tool_call_id')
                            and msg.tool_call_id in agui_tool_content
                    ):
                        msg.content = agui_tool_content[msg.tool_call_id]
                        replaced_tool_call_ids.add(msg.tool_call_id)

        existing_message_ids = {msg.id for msg in existing_messages}

        new_messages = [
            msg for msg in messages
            if msg.id not in existing_message_ids
        ]

        tools = input.tools or []
        tools_as_dicts = []
        if tools:
            for tool in tools:
                if hasattr(tool, "model_dump"):
                    tools_as_dicts.append(tool.model_dump())
                elif hasattr(tool, "dict"):
                    tools_as_dicts.append(tool.dict())
                else:
                    tools_as_dicts.append(tool)

        all_tools = [*state.get("tools", []), *tools_as_dicts]

        # Remove duplicates based on tool name
        seen_names = set()
        unique_tools = []
        for tool in all_tools:
            tool_name = tool.get("name") if isinstance(tool, dict) else getattr(tool, "name", None)
            if tool_name and tool_name not in seen_names:
                seen_names.add(tool_name)
                unique_tools.append(tool)
            elif not tool_name:
                # Keep tools without names (shouldn't happen, but just in case)
                unique_tools.append(tool)

        return {
            **state,
            "messages": new_messages,
            "tools": unique_tools,
            "ag-ui": {
                "tools": unique_tools,
                "context": input.context or []
            },
            "copilotkit": {
                **state.get("copilotkit", {}),
                "actions": unique_tools,
            },
        }

    _ORPHAN_TOOL_MSG_RE = re.compile(
        r"^Tool call '.+' with id '.+' was interrupted before completion\.$"
    )

    def _filter_orphan_tool_messages(self, messages: list) -> list:
        """Remove fake ToolMessages injected by patch_orphan_tool_calls,
        but only between the last user message and the end of the list."""
        # Find the index of the last HumanMessage
        last_human_idx = -1
        for i in range(len(messages) - 1, -1, -1):
            if isinstance(messages[i], HumanMessage):
                last_human_idx = i
                break

        if last_human_idx == -1:
            return messages

        # Keep everything before the last user message as-is,
        # filter the tail
        head = messages[:last_human_idx + 1]
        tail = [
            m for m in messages[last_human_idx + 1:]
            if not (
                    isinstance(m, ToolMessage)
                    and isinstance(m.content, str)
                    and self._ORPHAN_TOOL_MSG_RE.match(m.content)
            )
        ]
        return head + tail

    def get_state_snapshot(self, state: State) -> State:
        schema_keys = self.active_run["schema_keys"]
        if schema_keys and schema_keys.get("output"):
            state = filter_object_by_schema_keys(state, [*DEFAULT_SCHEMA_KEYS, *schema_keys["output"]])
        return state

    async def _handle_single_event(self, event: Any, state: State) -> AsyncGenerator[str, None]:
        event_type = event.get("event")
        if event_type == LangGraphEventTypes.OnChatModelStream:
            should_emit_messages = event.get("metadata", {}).get("emit-messages", True)
            should_emit_tool_calls = event.get("metadata", {}).get("emit-tool-calls", True)

            if event["data"]["chunk"].response_metadata.get('finish_reason', None):
                return

            current_stream = self.get_message_in_progress(self.active_run["id"])
            has_current_stream = bool(current_stream and current_stream.get("id"))
            tool_call_data = event["data"]["chunk"].tool_call_chunks[0] if event["data"]["chunk"].tool_call_chunks else None
            predict_state_metadata = event.get("metadata", {}).get("predict_state", [])
            tool_call_used_to_predict_state = False
            if tool_call_data and tool_call_data.get("name") and predict_state_metadata:
                tool_call_used_to_predict_state = any(
                    (predict_tool.get("tool") if isinstance(predict_tool, dict) else getattr(predict_tool, "tool", None)) == tool_call_data["name"]
                    for predict_tool in predict_state_metadata
                )

            is_tool_call_start_event = not has_current_stream and tool_call_data and tool_call_data.get("name")
            is_tool_call_args_event = has_current_stream and current_stream.get("tool_call_id") and tool_call_data and tool_call_data.get("args")
            is_tool_call_end_event = has_current_stream and current_stream.get("tool_call_id") and not tool_call_data

            if is_tool_call_start_event or is_tool_call_end_event or is_tool_call_args_event:
                self.active_run["has_function_streaming"] = True

            reasoning_data = resolve_reasoning_content(event["data"]["chunk"]) if event["data"]["chunk"] else None
            encrypted_reasoning_data = resolve_encrypted_reasoning_content(event["data"]["chunk"]) if event["data"]["chunk"] else None
            message_content = resolve_message_content(event["data"]["chunk"].content) if event["data"]["chunk"] and event["data"]["chunk"].content else None
            is_message_content_event = tool_call_data is None and message_content
            is_message_end_event = has_current_stream and not current_stream.get("tool_call_id") and not is_message_content_event

            if reasoning_data:
                for event_str in self.handle_reasoning_event(reasoning_data):
                    yield event_str
                return

            # Handle redacted_thinking blocks (encrypted reasoning content)
            if encrypted_reasoning_data and self.active_run.get('reasoning_process', None) is not None:
                reasoning_message_id = self.active_run["reasoning_process"]["message_id"]
                yield self._dispatch_event(
                    ReasoningEncryptedValueEvent(
                        type=EventType.REASONING_ENCRYPTED_VALUE,
                        subtype="message",
                        entity_id=reasoning_message_id,
                        encrypted_value=encrypted_reasoning_data,
                    )
                )
                return

            if reasoning_data is None and self.active_run.get('reasoning_process', None) is not None:
                reasoning_message_id = self.active_run["reasoning_process"]["message_id"]
                # Emit signature as encrypted value if accumulated during reasoning
                if self.active_run["reasoning_process"].get("signature"):
                    yield self._dispatch_event(
                        ReasoningEncryptedValueEvent(
                            type=EventType.REASONING_ENCRYPTED_VALUE,
                            subtype="message",
                            entity_id=reasoning_message_id,
                            encrypted_value=self.active_run["reasoning_process"]["signature"],
                        )
                    )
                yield self._dispatch_event(
                    ReasoningMessageEndEvent(
                        type=EventType.REASONING_MESSAGE_END,
                        message_id=reasoning_message_id,
                    )
                )
                yield self._dispatch_event(
                    ReasoningEndEvent(
                        type=EventType.REASONING_END,
                        message_id=reasoning_message_id,
                    )
                )
                self.active_run["reasoning_process"] = None

            if tool_call_used_to_predict_state:
                yield self._dispatch_event(
                    CustomEvent(
                        type=EventType.CUSTOM,
                        name="PredictState",
                        value=predict_state_metadata,
                        raw_event=event
                    )
                )

            if is_tool_call_end_event:
                yield self._dispatch_event(
                    ToolCallEndEvent(type=EventType.TOOL_CALL_END, tool_call_id=current_stream["tool_call_id"], raw_event=event)
                )
                self.messages_in_process[self.active_run["id"]] = None
                return


            if is_message_end_event:
                yield self._dispatch_event(
                    TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=current_stream["id"], raw_event=event)
                )
                self.messages_in_process[self.active_run["id"]] = None
                return

            if is_tool_call_start_event and should_emit_tool_calls:
                yield self._dispatch_event(
                    ToolCallStartEvent(
                        type=EventType.TOOL_CALL_START,
                        tool_call_id=tool_call_data["id"],
                        tool_call_name=tool_call_data["name"],
                        parent_message_id=event["data"]["chunk"].id,
                        raw_event=event,
                    )
                )
                self.set_message_in_progress(
                    self.active_run["id"],
                    MessageInProgress(id=event["data"]["chunk"].id, tool_call_id=tool_call_data["id"], tool_call_name=tool_call_data["name"])
                )
                return

            if is_tool_call_args_event and should_emit_tool_calls:
                yield self._dispatch_event(
                    ToolCallArgsEvent(
                        type=EventType.TOOL_CALL_ARGS,
                        tool_call_id=current_stream["tool_call_id"],
                        delta=tool_call_data["args"],
                        raw_event=event
                    )
                )
                return

            if is_message_content_event and should_emit_messages:
                if bool(current_stream and current_stream.get("id")) == False:
                    yield self._dispatch_event(
                        TextMessageStartEvent(
                            type=EventType.TEXT_MESSAGE_START,
                            role="assistant",
                            message_id=event["data"]["chunk"].id,
                            raw_event=event,
                        )
                    )
                    self.set_message_in_progress(
                        self.active_run["id"],
                        MessageInProgress(
                            id=event["data"]["chunk"].id,
                            tool_call_id=None,
                            tool_call_name=None
                        )
                    )
                    current_stream = self.get_message_in_progress(self.active_run["id"])

                yield self._dispatch_event(
                    TextMessageContentEvent(
                        type=EventType.TEXT_MESSAGE_CONTENT,
                        message_id=current_stream["id"],
                        delta=message_content,
                        raw_event=event,
                    )
                )
                return

        elif event_type == LangGraphEventTypes.OnChatModelEnd:
            if self.get_message_in_progress(self.active_run["id"]) and self.get_message_in_progress(self.active_run["id"]).get("tool_call_id"):
                resolved = self._dispatch_event(
                    ToolCallEndEvent(type=EventType.TOOL_CALL_END, tool_call_id=self.get_message_in_progress(self.active_run["id"])["tool_call_id"], raw_event=event)
                )
                if resolved:
                    self.messages_in_process[self.active_run["id"]] = None
                yield resolved
            elif self.get_message_in_progress(self.active_run["id"]) and self.get_message_in_progress(self.active_run["id"]).get("id"):
                resolved = self._dispatch_event(
                    TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=self.get_message_in_progress(self.active_run["id"])["id"], raw_event=event)
                )
                if resolved:
                    self.messages_in_process[self.active_run["id"]] = None
                yield resolved

        elif event_type == LangGraphEventTypes.OnCustomEvent:
            if event["name"] == CustomEventNames.ManuallyEmitMessage:
                yield self._dispatch_event(
                    TextMessageStartEvent(type=EventType.TEXT_MESSAGE_START, role="assistant", message_id=event["data"]["message_id"], raw_event=event)
                )
                yield self._dispatch_event(
                    TextMessageContentEvent(
                        type=EventType.TEXT_MESSAGE_CONTENT,
                        message_id=event["data"]["message_id"],
                        delta=event["data"]["message"],
                        raw_event=event,
                    )
                )
                yield self._dispatch_event(
                    TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=event["data"]["message_id"], raw_event=event)
                )

            elif event["name"] == CustomEventNames.ManuallyEmitToolCall:
                yield self._dispatch_event(
                    ToolCallStartEvent(
                        type=EventType.TOOL_CALL_START,
                        tool_call_id=event["data"]["id"],
                        tool_call_name=event["data"]["name"],
                        parent_message_id=event["data"]["id"],
                        raw_event=event,
                    )
                )
                yield self._dispatch_event(
                    ToolCallArgsEvent(
                        type=EventType.TOOL_CALL_ARGS,
                        tool_call_id=event["data"]["id"],
                        delta=event["data"]["args"] if isinstance(event["data"]["args"], str) else json.dumps(
                            event["data"]["args"]),
                        raw_event=event
                    )
                )
                yield self._dispatch_event(
                    ToolCallEndEvent(type=EventType.TOOL_CALL_END, tool_call_id=event["data"]["id"], raw_event=event)
                )

            elif event["name"] == CustomEventNames.ManuallyEmitState:
                self.active_run["manually_emitted_state"] = event["data"]
                yield self._dispatch_event(
                    StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=self.get_state_snapshot(self.active_run["manually_emitted_state"]), raw_event=event)
                )
            
            yield self._dispatch_event(
                CustomEvent(type=EventType.CUSTOM, name=event["name"], value=event["data"], raw_event=event)
            )

        elif event_type == LangGraphEventTypes.OnToolEnd:
            tool_call_output = event["data"]["output"]

            if isinstance(tool_call_output, Command):
                # Extract ToolMessages from Command.update
                messages = tool_call_output.update.get('messages', [])
                tool_messages = [m for m in messages if isinstance(m, ToolMessage)]

                # Process each tool message
                for tool_msg in tool_messages:
                    if not self.active_run["has_function_streaming"]:
                        yield self._dispatch_event(
                            ToolCallStartEvent(
                                type=EventType.TOOL_CALL_START,
                                tool_call_id=tool_msg.tool_call_id,
                                tool_call_name=tool_msg.name or event.get("name", ""),
                                parent_message_id=tool_msg.id,
                                raw_event=event,
                            )
                        )
                        yield self._dispatch_event(
                            ToolCallArgsEvent(
                                type=EventType.TOOL_CALL_ARGS,
                                tool_call_id=tool_msg.tool_call_id,
                                delta=json.dumps(event["data"].get("input", {})),
                                raw_event=event
                            )
                        )
                        yield self._dispatch_event(
                            ToolCallEndEvent(
                                type=EventType.TOOL_CALL_END,
                                tool_call_id=tool_msg.tool_call_id,
                                raw_event=event
                            )
                        )

                    yield self._dispatch_event(
                        ToolCallResultEvent(
                            type=EventType.TOOL_CALL_RESULT,
                            tool_call_id=tool_msg.tool_call_id,
                            message_id=str(uuid.uuid4()),
                            content=normalize_tool_content(tool_msg.content),
                            role="tool"
                        )
                    )
                self.active_run["has_function_streaming"] = False
                return

            if not self.active_run["has_function_streaming"]:
                yield self._dispatch_event(
                    ToolCallStartEvent(
                        type=EventType.TOOL_CALL_START,
                        tool_call_id=tool_call_output.tool_call_id,
                        tool_call_name=tool_call_output.name or event.get("name", ""),
                        parent_message_id=tool_call_output.id,
                        raw_event=event,
                    )
                )
                yield self._dispatch_event(
                    ToolCallArgsEvent(
                        type=EventType.TOOL_CALL_ARGS,
                        tool_call_id=tool_call_output.tool_call_id,
                        delta=dump_json_safe(event["data"]["input"]),
                        raw_event=event
                    )
                )
                yield self._dispatch_event(
                    ToolCallEndEvent(
                        type=EventType.TOOL_CALL_END,
                        tool_call_id=tool_call_output.tool_call_id,
                        raw_event=event
                    )
                )

            yield self._dispatch_event(
                ToolCallResultEvent(
                    type=EventType.TOOL_CALL_RESULT,
                    tool_call_id=tool_call_output.tool_call_id,
                    message_id=str(uuid.uuid4()),
                    content=normalize_tool_content(tool_call_output.content),
                    role="tool"
                )
            )

            self.active_run["model_made_tool_call"] = False
            self.active_run["state_reliable"] = True
            self.active_run["has_function_streaming"] = False

    def handle_reasoning_event(self, reasoning_data: LangGraphReasoning) -> Generator[str, Any, str | None]:
        if not reasoning_data or "type" not in reasoning_data or "text" not in reasoning_data:
            return ""

        reasoning_step_index = reasoning_data.get("index")

        if (self.active_run.get("reasoning_process") and
                self.active_run["reasoning_process"].get("index") and
                self.active_run["reasoning_process"]["index"] != reasoning_step_index):

            reasoning_message_id = self.active_run["reasoning_process"]["message_id"]
            if self.active_run["reasoning_process"].get("type"):
                yield self._dispatch_event(
                    ReasoningMessageEndEvent(
                        type=EventType.REASONING_MESSAGE_END,
                        message_id=reasoning_message_id,
                    )
                )
            yield self._dispatch_event(
                ReasoningEndEvent(
                    type=EventType.REASONING_END,
                    message_id=reasoning_message_id,
                )
            )
            self.active_run["reasoning_process"] = None

        if not self.active_run.get("reasoning_process"):
            message_id = str(uuid.uuid4())
            yield self._dispatch_event(
                ReasoningStartEvent(
                    type=EventType.REASONING_START,
                    message_id=message_id,
                )
            )
            self.active_run["reasoning_process"] = {
                "index": reasoning_step_index,
                "message_id": message_id,
            }

        if self.active_run["reasoning_process"].get("type") != reasoning_data["type"]:
            yield self._dispatch_event(
                ReasoningMessageStartEvent(
                    type=EventType.REASONING_MESSAGE_START,
                    message_id=self.active_run["reasoning_process"]["message_id"],
                    role="assistant",
                )
            )
            self.active_run["reasoning_process"]["type"] = reasoning_data["type"]

        # Accumulate signature if present (Anthropic extended thinking)
        if reasoning_data.get("signature"):
            self.active_run["reasoning_process"]["signature"] = reasoning_data["signature"]

        if self.active_run["reasoning_process"].get("type"):
            yield self._dispatch_event(
                ReasoningMessageContentEvent(
                    type=EventType.REASONING_MESSAGE_CONTENT,
                    message_id=self.active_run["reasoning_process"]["message_id"],
                    delta=reasoning_data["text"]
                )
            )

    async def get_checkpoint_before_message(self, message_id: str, thread_id: str, config: Optional[RunnableConfig] = None):
        if not thread_id:
            raise ValueError("Missing thread_id in config")

        history_list = []
        async for snapshot in self.graph.aget_state_history(history_config):
            history_list.append(snapshot)

        history_list.reverse()
        for idx, snapshot in enumerate(history_list):
            messages = snapshot.values.get("messages", [])
            if any(getattr(m, "id", None) == message_id for m in messages):
                if idx == 0:
                    # No snapshot before this
                    # Return synthetic "empty before" version
                    empty_snapshot = snapshot
                    empty_snapshot.values["messages"] = []
                    return empty_snapshot

                snapshot_values_without_messages = snapshot.values.copy()
                del snapshot_values_without_messages["messages"]
                checkpoint = history_list[idx - 1]

                merged_values = {**checkpoint.values, **snapshot_values_without_messages}
                checkpoint = checkpoint._replace(values=merged_values)

                return checkpoint

        raise ValueError("Message ID not found in history")

    def handle_node_change(self, node_name: Optional[str]):
        """
        Centralized method to handle node name changes and step transitions.
        Automatically manages step start/end events based on node name changes.
        """
        if node_name == "__end__":
            node_name = None

        if node_name != self.active_run.get("node_name"):
            # End current step if we have one
            if self.active_run.get("node_name"):
                yield self.end_step()

            # Start new step if we have a node name
            if node_name:
                for event in self.start_step(node_name):
                    yield event

        self.active_run["node_name"] = node_name

    def start_step(self, step_name: str):
        """Simple step start event dispatcher - node_name management handled by handle_node_change"""
        yield self._dispatch_event(
            StepStartedEvent(
                type=EventType.STEP_STARTED,
                step_name=step_name
            )
        )

    def end_step(self):
        """Simple step end event dispatcher - node_name management handled by handle_node_change"""
        if not self.active_run.get("node_name"):
            raise ValueError("No active step to end")

        return self._dispatch_event(
            StepFinishedEvent(
                type=EventType.STEP_FINISHED,
                step_name=self.active_run["node_name"]
            )
        )

    # Check if some kwargs are enabled per LG version, to "catch all versions" and backwards compatibility
    def get_stream_kwargs(
            self,
            input: Any,
            subgraphs: bool = False,
            version: Literal["v1", "v2"] = "v2",
            config: Optional[RunnableConfig] = None,
            context: Optional[Dict[str, Any]] = None,
            fork: Optional[Any] = None,
    ):
        kwargs = dict(
            input=input,
            subgraphs=subgraphs,
            version=version,
        )

        # Only add context if supported
        sig = inspect.signature(self.graph.astream_events)
        if 'context' in sig.parameters:
            base_context = {}
            if isinstance(config, dict) and 'configurable' in config and isinstance(config['configurable'], dict):
                base_context.update(config['configurable'])
            if context:  # context might be None or {}
                base_context.update(context)
            if base_context:  # only add if there's something to pass
                kwargs['context'] = base_context

        if config:
            kwargs['config'] = config

        if fork:
            kwargs.update(fork)

        return kwargs


def dump_json_safe(value):
    return json.dumps(value, default=json_safe_stringify) if not isinstance(value, str) else value