async def arun(
        self,
        system_instruction: str | None = None,
        messages: list[Message] | None = None,
        *args,
        **kwargs,
    ) -> AsyncGenerator[str, None]:
        """
        Main streaming entrypoint: returns an async generator of SSE lines.
        """
        self._reset()
        await self._setup(system_instruction)

        if messages:
            for m in messages:
                await self.conversation.add_message(m)

        # Initialize citation tracker for this run
        citation_tracker = CitationTracker()

        # Dictionary to store citation payloads by ID
        citation_payloads = {}

        # Track all citations emitted during streaming for final persistence
        self.streaming_citations: list[dict] = []

        async def sse_generator() -> AsyncGenerator[str, None]:
            iterations_count = 0

            try:
                # Keep streaming until we complete
                while (
                    not self._completed
                    and iterations_count < self.config.max_iterations
                ):
                    iterations_count += 1
                    # 1) Get current messages
                    msg_list = await self.conversation.get_messages()
                    gen_cfg = self.get_generation_config(
                        msg_list[-1], stream=True
                    )

                    # 2) Start streaming from LLM
                    llm_stream = self.llm_provider.aget_completion_stream(
                        msg_list, gen_cfg
                    )

                    # Create state variables for each iteration
                    iteration_buffer = ""
                    yielded_first_event = False
                    in_action_block = False
                    is_thinking = False
                    accumulated_thinking = ""
                    thinking_signatures = {}

                    async for chunk in llm_stream:
                        delta = chunk.choices[0].delta
                        finish_reason = chunk.choices[0].finish_reason

                        # Handle thinking if present
                        if hasattr(delta, "thinking") and delta.thinking:
                            # Accumulate thinking for later use in messages
                            accumulated_thinking += delta.thinking

                            # Emit SSE "thinking" event
                            async for (
                                line
                            ) in SSEFormatter.yield_thinking_event(
                                delta.thinking
                            ):
                                yield line

                        # Add this new handler for thinking signatures
                        if hasattr(delta, "thinking_signature"):
                            thinking_signatures[accumulated_thinking] = (
                                delta.thinking_signature
                            )
                            accumulated_thinking = ""

                        # 3) If new text, accumulate it
                        if delta.content:
                            iteration_buffer += delta.content

                            # Check if we have accumulated enough text for a `<Thought>` block
                            if len(iteration_buffer) < len("<Thought>"):
                                continue

                            # Check if we have yielded the first event
                            if not yielded_first_event:
                                # Emit the first chunk
                                if self.THOUGHT_OPEN.findall(iteration_buffer):
                                    is_thinking = True
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_thinking_event(
                                        iteration_buffer
                                    ):
                                        yield line
                                else:
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_message_event(
                                        iteration_buffer
                                    ):
                                        yield line

                                # Mark as yielded
                                yielded_first_event = True
                                continue

                            # Check if we are in a thinking block
                            if is_thinking:
                                # Still thinking, so keep yielding thinking events
                                if not self.THOUGHT_CLOSE.findall(
                                    iteration_buffer
                                ):
                                    # Emit SSE "thinking" event
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_thinking_event(
                                        delta.content
                                    ):
                                        yield line

                                    continue
                                # Done thinking, so emit the last thinking event
                                else:
                                    is_thinking = False
                                    thought_text = delta.content.split(
                                        "</Thought>"
                                    )[0].split("</think>")[0]
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_thinking_event(
                                        thought_text
                                    ):
                                        yield line
                                    post_thought_text = delta.content.split(
                                        "</Thought>"
                                    )[-1].split("</think>")[-1]
                                    delta.content = post_thought_text

                            # (b) Find new citation spans in the accumulated text
                            new_citation_spans = find_new_citation_spans(
                                iteration_buffer, citation_tracker
                            )

                            # Process each new citation span
                            for cid, spans in new_citation_spans.items():
                                for span in spans:
                                    # Check if this is the first time we've seen this citation ID
                                    is_new_citation = (
                                        citation_tracker.is_new_citation(cid)
                                    )

                                    # Get payload if it's a new citation
                                    payload = None
                                    if is_new_citation:
                                        source_obj = self.search_results_collector.find_by_short_id(
                                            cid
                                        )
                                        if source_obj:
                                            # Store payload for reuse
                                            payload = dump_obj(source_obj)
                                            citation_payloads[cid] = payload

                                    # Create citation event payload
                                    citation_data = {
                                        "id": cid,
                                        "object": "citation",
                                        "is_new": is_new_citation,
                                        "span": {
                                            "start": span[0],
                                            "end": span[1],
                                        },
                                    }

                                    # Only include full payload for new citations
                                    if is_new_citation and payload:
                                        citation_data["payload"] = payload

                                    # Add to streaming citations for final answer
                                    self.streaming_citations.append(
                                        citation_data
                                    )

                                    # Emit the citation event
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_citation_event(
                                        citation_data
                                    ):
                                        yield line

                            # Now prepare to emit the newly streamed text as a "message" event
                            if (
                                iteration_buffer.count("<")
                                and not in_action_block
                            ):
                                in_action_block = True

                            if (
                                in_action_block
                                and len(
                                    self.ACTION_PATTERN.findall(
                                        iteration_buffer
                                    )
                                )
                                < 2
                            ):
                                continue

                            elif in_action_block:
                                in_action_block = False
                                # Emit the post action block text, if it is there
                                post_action_text = iteration_buffer.split(
                                    "</Action>"
                                )[-1]
                                if post_action_text:
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_message_event(
                                        post_action_text
                                    ):
                                        yield line

                            else:
                                async for (
                                    line
                                ) in SSEFormatter.yield_message_event(
                                    delta.content
                                ):
                                    yield line

                        elif finish_reason == "stop":
                            break

                    # Process any accumulated thinking
                    await self._handle_thinking(
                        thinking_signatures, accumulated_thinking
                    )

                    # 6) The LLM is done. If we have any leftover partial text,
                    #    finalize it in the conversation
                    if iteration_buffer:
                        # Create the final message with metadata including citations
                        final_message = Message(
                            role="assistant",
                            content=iteration_buffer,
                            metadata={"citations": self.streaming_citations},
                        )

                        # Add it to the conversation
                        await self.conversation.add_message(final_message)

                    # --- 4) Process any <Action>/<ToolCalls> blocks, or mark completed
                    action_matches = self.ACTION_PATTERN.findall(
                        iteration_buffer
                    )

                    if len(action_matches) > 0:
                        # Process each ToolCall
                        xml_toolcalls = "<ToolCalls>"

                        for action_block in action_matches:
                            tool_calls_text = []
                            # Look for ToolCalls wrapper, or use the raw action block
                            calls_wrapper = self.TOOLCALLS_PATTERN.findall(
                                action_block
                            )
                            if calls_wrapper:
                                for tw in calls_wrapper:
                                    tool_calls_text.append(tw)
                            else:
                                tool_calls_text.append(action_block)

                            for calls_region in tool_calls_text:
                                calls_found = self.TOOLCALL_PATTERN.findall(
                                    calls_region
                                )
                                for tc_block in calls_found:
                                    tool_name, tool_params = (
                                        self._parse_single_tool_call(tc_block)
                                    )
                                    if tool_name:
                                        # Emit SSE event for tool call
                                        tool_call_id = (
                                            f"call_{abs(hash(tc_block))}"
                                        )
                                        call_evt_data = {
                                            "tool_call_id": tool_call_id,
                                            "name": tool_name,
                                            "arguments": json.dumps(
                                                tool_params
                                            ),
                                        }
                                        async for line in (
                                            SSEFormatter.yield_tool_call_event(
                                                call_evt_data
                                            )
                                        ):
                                            yield line

                                        try:
                                            tool_result = await self.handle_function_or_tool_call(
                                                tool_name,
                                                json.dumps(tool_params),
                                                tool_id=tool_call_id,
                                                save_messages=False,
                                            )
                                            result_content = tool_result.llm_formatted_result
                                        except Exception as e:
                                            result_content = f"Error in tool '{tool_name}': {str(e)}"

                                        xml_toolcalls += (
                                            f"<ToolCall>"
                                            f"<Name>{tool_name}</Name>"
                                            f"<Parameters>{json.dumps(tool_params)}</Parameters>"
                                            f"<Result>{result_content}</Result>"
                                            f"</ToolCall>"
                                        )

                                        # Emit SSE tool result for non-result tools
                                        result_data = {
                                            "tool_call_id": tool_call_id,
                                            "role": "tool",
                                            "content": json.dumps(
                                                convert_nonserializable_objects(
                                                    result_content
                                                )
                                            ),
                                        }
                                        async for line in SSEFormatter.yield_tool_result_event(
                                            result_data
                                        ):
                                            yield line

                        xml_toolcalls += "</ToolCalls>"
                        pre_action_text = iteration_buffer[
                            : iteration_buffer.find(action_block)
                        ]
                        post_action_text = iteration_buffer[
                            iteration_buffer.find(action_block)
                            + len(action_block) :
                        ]
                        iteration_text = (
                            pre_action_text + xml_toolcalls + post_action_text
                        )

                        # Update the conversation with tool results
                        await self.conversation.add_message(
                            Message(
                                role="assistant",
                                content=iteration_text,
                                metadata={
                                    "citations": self.streaming_citations
                                },
                            )
                        )
                    else:
                        # (a) Prepare final answer with optimized citations
                        consolidated_citations = []
                        # Group citations by ID with all their spans
                        for (
                            cid,
                            spans,
                        ) in citation_tracker.get_all_spans().items():
                            if cid in citation_payloads:
                                consolidated_citations.append(
                                    {
                                        "id": cid,
                                        "object": "citation",
                                        "spans": [
                                            {"start": s[0], "end": s[1]}
                                            for s in spans
                                        ],
                                        "payload": citation_payloads[cid],
                                    }
                                )

                        # Create final answer payload
                        final_evt_payload = {
                            "id": "msg_final",
                            "object": "agent.final_answer",
                            "generated_answer": iteration_buffer,
                            "citations": consolidated_citations,
                        }

                        # Emit final answer event
                        async for (
                            line
                        ) in SSEFormatter.yield_final_answer_event(
                            final_evt_payload
                        ):
                            yield line

                        # (b) Signal the end of the SSE stream
                        yield SSEFormatter.yield_done_event()
                        self._completed = True

                # If we exit the while loop due to hitting max iterations
                if not self._completed:
                    # Generate a summary using the LLM
                    summary = await self._generate_llm_summary(
                        iterations_count
                    )

                    # Send the summary as a message event
                    async for line in SSEFormatter.yield_message_event(
                        summary
                    ):
                        yield line

                    # Add summary to conversation with citations metadata
                    await self.conversation.add_message(
                        Message(
                            role="assistant",
                            content=summary,
                            metadata={"citations": self.streaming_citations},
                        )
                    )

                    # Create and emit a final answer payload with the summary
                    final_evt_payload = {
                        "id": "msg_final",
                        "object": "agent.final_answer",
                        "generated_answer": summary,
                        "citations": consolidated_citations,
                    }

                    async for line in SSEFormatter.yield_final_answer_event(
                        final_evt_payload
                    ):
                        yield line

                    # Signal the end of the SSE stream
                    yield SSEFormatter.yield_done_event()
                    self._completed = True

            except Exception as e:
                logger.error(f"Error in streaming agent: {str(e)}")
                # Emit error event for client
                async for line in SSEFormatter.yield_error_event(
                    f"Agent error: {str(e)}"
                ):
                    yield line
                # Send done event to close the stream
                yield SSEFormatter.yield_done_event()

        # Finally, we return the async generator
        async for line in sse_generator():
            yield line

        Main streaming entrypoint: returns an async generator of SSE lines.
        
py/core/agent/base.py