agno 2.3.9__tar.gz → 2.3.10__tar.gz

{agno-2.3.9 → agno-2.3.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agno
-Version: 2.3.9
+Version: 2.3.10
 Summary: Agno: a lightweight library for building Multi-Agent Systems
 Author-email: Ashpreet Bedi <ashpreet@agno.com>
 Project-URL: homepage, https://agno.com

{agno-2.3.9 → agno-2.3.10}/agno/agent/agent.py

@@ -1785,9 +1785,6 @@ class Agent:
                 if stream_events is None:
                     stream_events = False if self.stream_events is None else self.stream_events
 
-                self.stream = self.stream or stream
-                self.stream_events = self.stream_events or stream_events
-
                 # Prepare arguments for the model
                 response_format = (
                     self._get_response_format(run_context=run_context) if self.parser_model is None else None
@@ -2711,9 +2708,6 @@ class Agent:
         if stream_events is None:
             stream_events = False if self.stream_events is None else self.stream_events
 
-        self.stream = self.stream or stream
-        self.stream_events = self.stream_events or stream_events
-
         self.model = cast(Model, self.model)
 
         # Get knowledge filters
@@ -3020,9 +3014,6 @@ class Agent:
                 if stream is False:
                     stream_events = False
 
-                self.stream = self.stream or stream
-                self.stream_events = self.stream_events or stream_events
-
                 # Run can be continued from previous run response or from passed run_response context
                 if run_response is not None:
                     # The run is continued from a provided run_response. This contains the updated tools.
@@ -3589,9 +3580,6 @@ class Agent:
         if stream is False:
             stream_events = False
 
-        self.stream = self.stream or stream
-        self.stream_events = self.stream_events or stream_events
-
         # Get knowledge filters
         knowledge_filters = knowledge_filters
         if self.knowledge_filters or knowledge_filters:

{agno-2.3.9 → agno-2.3.10}/agno/db/sqlite/async_sqlite.py

@@ -1825,7 +1825,7 @@ class AsyncSqliteDb(AsyncBaseDb):
             Optional[KnowledgeRow]: The upserted knowledge row, or None if the operation fails.
         """
         try:
-            table = await self._get_table(table_type="knowledge")
+            table = await self._get_table(table_type="knowledge", create_table_if_not_found=True)
             if table is None:
                 return None
 

{agno-2.3.9 → agno-2.3.10}/agno/eval/accuracy.py

@@ -282,7 +282,8 @@ Remember: You must only compare the agent_output to the expected_output. The exp
     ) -> Optional[AccuracyEvaluation]:
         """Orchestrate the evaluation process."""
         try:
-            accuracy_agent_response = evaluator_agent.run(evaluation_input).content
+            response = evaluator_agent.run(evaluation_input, stream=False)
+            accuracy_agent_response = response.content
             if accuracy_agent_response is None or not isinstance(accuracy_agent_response, AccuracyAgentResponse):
                 raise EvalError(f"Evaluator Agent returned an invalid response: {accuracy_agent_response}")
             return AccuracyEvaluation(
@@ -306,7 +307,7 @@ Remember: You must only compare the agent_output to the expected_output. The exp
     ) -> Optional[AccuracyEvaluation]:
         """Orchestrate the evaluation process asynchronously."""
         try:
-            response = await evaluator_agent.arun(evaluation_input)
+            response = await evaluator_agent.arun(evaluation_input, stream=False)
             accuracy_agent_response = response.content
             if accuracy_agent_response is None or not isinstance(accuracy_agent_response, AccuracyAgentResponse):
                 raise EvalError(f"Evaluator Agent returned an invalid response: {accuracy_agent_response}")
@@ -362,9 +363,11 @@ Remember: You must only compare the agent_output to the expected_output. The exp
                 agent_session_id = f"eval_{self.eval_id}_{i + 1}"
 
                 if self.agent is not None:
-                    output = self.agent.run(input=eval_input, session_id=agent_session_id).content
+                    agent_response = self.agent.run(input=eval_input, session_id=agent_session_id, stream=False)
+                    output = agent_response.content
                 elif self.team is not None:
-                    output = self.team.run(input=eval_input, session_id=agent_session_id).content
+                    team_response = self.team.run(input=eval_input, session_id=agent_session_id, stream=False)
+                    output = team_response.content
 
                 if not output:
                     logger.error(f"Failed to generate a valid answer on iteration {i + 1}: {output}")
@@ -505,11 +508,11 @@ Remember: You must only compare the agent_output to the expected_output. The exp
                 agent_session_id = f"eval_{self.eval_id}_{i + 1}"
 
                 if self.agent is not None:
-                    response = await self.agent.arun(input=eval_input, session_id=agent_session_id)
-                    output = response.content
+                    agent_response = await self.agent.arun(input=eval_input, session_id=agent_session_id, stream=False)
+                    output = agent_response.content
                 elif self.team is not None:
-                    response = await self.team.arun(input=eval_input, session_id=agent_session_id)  # type: ignore
-                    output = response.content
+                    team_response = await self.team.arun(input=eval_input, session_id=agent_session_id, stream=False)
+                    output = team_response.content
 
                 if not output:
                     logger.error(f"Failed to generate a valid answer on iteration {i + 1}: {output}")

{agno-2.3.9 → agno-2.3.10}/agno/eval/agent_as_judge.py

@@ -281,24 +281,25 @@ class AgentAsJudgeEval(BaseEval):
                 </output>
             """)
 
-            response = evaluator_agent.run(prompt).content
-            if not isinstance(response, (NumericJudgeResponse, BinaryJudgeResponse)):
-                raise EvalError(f"Invalid response: {response}")
+            response = evaluator_agent.run(prompt, stream=False)
+            judge_response = response.content
+            if not isinstance(judge_response, (NumericJudgeResponse, BinaryJudgeResponse)):
+                raise EvalError(f"Invalid response: {judge_response}")
 
             # Determine pass/fail based on scoring strategy and response type
-            if isinstance(response, NumericJudgeResponse):
-                score = response.score
+            if isinstance(judge_response, NumericJudgeResponse):
+                score = judge_response.score
                 passed = score >= self.threshold
             else:  # BinaryJudgeResponse
                 score = None
-                passed = response.passed
+                passed = judge_response.passed
 
             evaluation = AgentAsJudgeEvaluation(
                 input=input,
                 output=output,
                 criteria=self.criteria,
                 score=score,
-                reason=response.reason,
+                reason=judge_response.reason,
                 passed=passed,
             )
 
@@ -332,7 +333,7 @@ class AgentAsJudgeEval(BaseEval):
                 </output>
             """)
 
-            response = await evaluator_agent.arun(prompt)
+            response = await evaluator_agent.arun(prompt, stream=False)
             judge_response = response.content
             if not isinstance(judge_response, (NumericJudgeResponse, BinaryJudgeResponse)):
                 raise EvalError(f"Invalid response: {judge_response}")

{agno-2.3.9 → agno-2.3.10}/agno/knowledge/embedder/openai.py

@@ -4,7 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple
 from typing_extensions import Literal
 
 from agno.knowledge.embedder.base import Embedder
-from agno.utils.log import log_warning, log_info
+from agno.utils.log import log_info, log_warning
 
 try:
     from openai import AsyncOpenAI

{agno-2.3.9 → agno-2.3.10}/agno/models/base.py

@@ -1801,6 +1801,17 @@ class Model(ABC):
                 log_error(f"Error while iterating function result generator for {function_call.function.name}: {e}")
                 function_call.error = str(e)
                 function_call_success = False
+
+            # For generators, re-capture updated_session_state after consumption
+            # since session_state modifications were made during iteration
+            if function_execution_result.updated_session_state is None:
+                if (
+                    function_call.function._run_context is not None
+                    and function_call.function._run_context.session_state is not None
+                ):
+                    function_execution_result.updated_session_state = function_call.function._run_context.session_state
+                elif function_call.function._session_state is not None:
+                    function_execution_result.updated_session_state = function_call.function._session_state
         else:
             from agno.tools.function import ToolResult
 
@@ -2327,7 +2338,29 @@ class Model(ABC):
                     log_error(f"Error while iterating function result generator for {function_call.function.name}: {e}")
                     function_call.error = str(e)
                     function_call_success = False
-            else:
+
+            # For generators (sync or async), re-capture updated_session_state after consumption
+            # since session_state modifications were made during iteration
+            if async_function_call_output is not None or isinstance(
+                function_call.result,
+                (GeneratorType, collections.abc.Iterator, AsyncGeneratorType, collections.abc.AsyncIterator),
+            ):
+                if updated_session_state is None:
+                    if (
+                        function_call.function._run_context is not None
+                        and function_call.function._run_context.session_state is not None
+                    ):
+                        updated_session_state = function_call.function._run_context.session_state
+                    elif function_call.function._session_state is not None:
+                        updated_session_state = function_call.function._session_state
+
+            if not (
+                async_function_call_output is not None
+                or isinstance(
+                    function_call.result,
+                    (GeneratorType, collections.abc.Iterator, AsyncGeneratorType, collections.abc.AsyncIterator),
+                )
+            ):
                 from agno.tools.function import ToolResult
 
                 if isinstance(function_execution_result.result, ToolResult):

{agno-2.3.9 → agno-2.3.10}/agno/models/google/gemini.py

@@ -36,6 +36,7 @@ try:
         GenerateContentResponseUsageMetadata,
         GoogleSearch,
         GoogleSearchRetrieval,
+        GroundingMetadata,
         Operation,
         Part,
         Retrieval,
@@ -244,8 +245,8 @@ class Gemini(Model):
         builtin_tools = []
 
         if self.grounding:
-            log_info(
-                "Grounding enabled. This is a legacy tool. For Gemini 2.0+ Please use enable `search` flag instead."
+            log_debug(
+                "Gemini Grounding enabled. This is a legacy tool. For Gemini 2.0+ Please use enable `search` flag instead."
             )
             builtin_tools.append(
                 Tool(
@@ -258,15 +259,15 @@ class Gemini(Model):
             )
 
         if self.search:
-            log_info("Google Search enabled.")
+            log_debug("Gemini Google Search enabled.")
             builtin_tools.append(Tool(google_search=GoogleSearch()))
 
         if self.url_context:
-            log_info("URL context enabled.")
+            log_debug("Gemini URL context enabled.")
             builtin_tools.append(Tool(url_context=UrlContext()))
 
         if self.vertexai_search:
-            log_info("Vertex AI Search enabled.")
+            log_debug("Gemini Vertex AI Search enabled.")
             if not self.vertexai_search_datastore:
                 log_error("vertexai_search_datastore must be provided when vertexai_search is enabled.")
                 raise ValueError("vertexai_search_datastore must be provided when vertexai_search is enabled.")
@@ -1008,27 +1009,24 @@ class Gemini(Model):
             citations = Citations()
             citations_raw = {}
             citations_urls = []
+            web_search_queries: List[str] = []
 
             if response.candidates and response.candidates[0].grounding_metadata is not None:
-                grounding_metadata = response.candidates[0].grounding_metadata.model_dump()
-                citations_raw["grounding_metadata"] = grounding_metadata
+                grounding_metadata: GroundingMetadata = response.candidates[0].grounding_metadata
+                citations_raw["grounding_metadata"] = grounding_metadata.model_dump()
 
-                chunks = grounding_metadata.get("grounding_chunks", []) or []
-                citation_pairs = []
+                chunks = grounding_metadata.grounding_chunks or []
+                web_search_queries = grounding_metadata.web_search_queries or []
                 for chunk in chunks:
-                    if not isinstance(chunk, dict):
+                    if not chunk:
                         continue
-                    web = chunk.get("web")
-                    if not isinstance(web, dict):
+                    web = chunk.web
+                    if not web:
                         continue
-                    uri = web.get("uri")
-                    title = web.get("title")
+                    uri = web.uri
+                    title = web.title
                     if uri:
-                        citation_pairs.append((uri, title))
-
-                # Create citation objects from filtered pairs
-                grounding_urls = [UrlCitation(url=url, title=title) for url, title in citation_pairs]
-                citations_urls.extend(grounding_urls)
+                        citations_urls.append(UrlCitation(url=uri, title=title))
 
             # Handle URLs from URL context tool
             if (
@@ -1036,22 +1034,29 @@ class Gemini(Model):
                 and hasattr(response.candidates[0], "url_context_metadata")
                 and response.candidates[0].url_context_metadata is not None
             ):
-                url_context_metadata = response.candidates[0].url_context_metadata.model_dump()
-                citations_raw["url_context_metadata"] = url_context_metadata
+                url_context_metadata = response.candidates[0].url_context_metadata
+                citations_raw["url_context_metadata"] = url_context_metadata.model_dump()
 
-                url_metadata_list = url_context_metadata.get("url_metadata", [])
+                url_metadata_list = url_context_metadata.url_metadata or []
                 for url_meta in url_metadata_list:
-                    retrieved_url = url_meta.get("retrieved_url")
-                    status = url_meta.get("url_retrieval_status", "UNKNOWN")
+                    retrieved_url = url_meta.retrieved_url
+                    status = "UNKNOWN"
+                    if url_meta.url_retrieval_status:
+                        status = url_meta.url_retrieval_status.value
                     if retrieved_url and status == "URL_RETRIEVAL_STATUS_SUCCESS":
                         # Avoid duplicate URLs
                         existing_urls = [citation.url for citation in citations_urls]
                         if retrieved_url not in existing_urls:
                             citations_urls.append(UrlCitation(url=retrieved_url, title=retrieved_url))
 
+            if citations_raw:
+                citations.raw = citations_raw
+            if citations_urls:
+                citations.urls = citations_urls
+            if web_search_queries:
+                citations.search_queries = web_search_queries
+
             if citations_raw or citations_urls:
-                citations.raw = citations_raw if citations_raw else None
-                citations.urls = citations_urls if citations_urls else None
                 model_response.citations = citations
 
         # Extract usage metadata if present
@@ -1150,28 +1155,52 @@ class Gemini(Model):
 
                         model_response.tool_calls.append(tool_call)
 
-            if response_delta.candidates[0].grounding_metadata is not None:
-                citations = Citations()
-                grounding_metadata = response_delta.candidates[0].grounding_metadata.model_dump()
-                citations.raw = grounding_metadata
+            citations = Citations()
+            citations.raw = {}
+            citations.urls = []
 
+            if (
+                hasattr(response_delta.candidates[0], "grounding_metadata")
+                and response_delta.candidates[0].grounding_metadata is not None
+            ):
+                grounding_metadata = response_delta.candidates[0].grounding_metadata
+                citations.raw["grounding_metadata"] = grounding_metadata.model_dump()
+                citations.search_queries = grounding_metadata.web_search_queries or []
                 # Extract url and title
-                chunks = grounding_metadata.pop("grounding_chunks", None) or []
-                citation_pairs = []
+                chunks = grounding_metadata.grounding_chunks or []
                 for chunk in chunks:
-                    if not isinstance(chunk, dict):
+                    if not chunk:
                         continue
-                    web = chunk.get("web")
-                    if not isinstance(web, dict):
+                    web = chunk.web
+                    if not web:
                         continue
-                    uri = web.get("uri")
-                    title = web.get("title")
+                    uri = web.uri
+                    title = web.title
                     if uri:
-                        citation_pairs.append((uri, title))
+                        citations.urls.append(UrlCitation(url=uri, title=title))
+
+            # Handle URLs from URL context tool
+            if (
+                hasattr(response_delta.candidates[0], "url_context_metadata")
+                and response_delta.candidates[0].url_context_metadata is not None
+            ):
+                url_context_metadata = response_delta.candidates[0].url_context_metadata
 
-                # Create citation objects from filtered pairs
-                citations.urls = [UrlCitation(url=url, title=title) for url, title in citation_pairs]
+                citations.raw["url_context_metadata"] = url_context_metadata.model_dump()
+
+                url_metadata_list = url_context_metadata.url_metadata or []
+                for url_meta in url_metadata_list:
+                    retrieved_url = url_meta.retrieved_url
+                    status = "UNKNOWN"
+                    if url_meta.url_retrieval_status:
+                        status = url_meta.url_retrieval_status.value
+                    if retrieved_url and status == "URL_RETRIEVAL_STATUS_SUCCESS":
+                        # Avoid duplicate URLs
+                        existing_urls = [citation.url for citation in citations.urls]
+                        if retrieved_url not in existing_urls:
+                            citations.urls.append(UrlCitation(url=retrieved_url, title=retrieved_url))
 
+            if citations.raw or citations.urls:
                 model_response.citations = citations
 
             # Extract usage metadata if present

{agno-2.3.9 → agno-2.3.10}/agno/models/message.py

@@ -42,6 +42,9 @@ class Citations(BaseModel):
     # Raw citations from the model
     raw: Optional[Any] = None
 
+    # Search queries used to retrieve the citations
+    search_queries: Optional[List[str]] = None
+
     # URLs of the citations.
     urls: Optional[List[UrlCitation]] = None
 

{agno-2.3.9 → agno-2.3.10}/agno/os/routers/evals/utils.py

@@ -66,14 +66,14 @@ async def run_agent_as_judge_eval(
 
     # Run agent/team to get output
     if agent:
-        agent_response = await agent.arun(eval_run_input.input)
+        agent_response = await agent.arun(eval_run_input.input, stream=False)
         output = str(agent_response.content) if agent_response.content else ""
         model_id = agent.model.id if agent and agent.model else None
         model_provider = agent.model.provider if agent and agent.model else None
         agent_id = agent.id
         team_id = None
     elif team:
-        team_response = await team.arun(eval_run_input.input)
+        team_response = await team.arun(eval_run_input.input, stream=False)
         output = str(team_response.content) if team_response.content else ""
         model_id = team.model.id if team and team.model else None
         model_provider = team.model.provider if team and team.model else None
@@ -125,39 +125,21 @@ async def run_performance_eval(
     default_model: Optional[Model] = None,
 ) -> EvalSchema:
     """Run a performance evaluation for the given agent or team"""
-    # Create sync or async function based on DB type
-    if isinstance(db, AsyncBaseDb):
-        if agent:
-
-            async def run_component():  # type: ignore
-                return await agent.arun(eval_run_input.input)
-
-            model_id = agent.model.id if agent and agent.model else None
-            model_provider = agent.model.provider if agent and agent.model else None
-
-        elif team:
-
-            async def run_component():  # type: ignore
-                return await team.arun(eval_run_input.input)
-
-            model_id = team.model.id if team and team.model else None
-            model_provider = team.model.provider if team and team.model else None
-    else:
-        if agent:
+    if agent:
 
-            def run_component():  # type: ignore
-                return agent.run(eval_run_input.input)
+        async def run_component():  # type: ignore
+            return await agent.arun(eval_run_input.input, stream=False)
 
-            model_id = agent.model.id if agent and agent.model else None
-            model_provider = agent.model.provider if agent and agent.model else None
+        model_id = agent.model.id if agent and agent.model else None
+        model_provider = agent.model.provider if agent and agent.model else None
 
-        elif team:
+    elif team:
 
-            def run_component():
-                return team.run(eval_run_input.input)
+        async def run_component():  # type: ignore
+            return await team.arun(eval_run_input.input, stream=False)
 
-            model_id = team.model.id if team and team.model else None
-            model_provider = team.model.provider if team and team.model else None
+        model_id = team.model.id if team and team.model else None
+        model_provider = team.model.provider if team and team.model else None
 
     performance_eval = PerformanceEval(
         db=db,
@@ -171,11 +153,7 @@ async def run_performance_eval(
         model_provider=model_provider,
     )
 
-    # PerformanceEval needs sync/async check because it wraps a function
-    if isinstance(db, AsyncBaseDb):
-        result = await performance_eval.arun(print_results=False, print_summary=False)
-    else:
-        result = performance_eval.run(print_results=False, print_summary=False)
+    result = await performance_eval.arun(print_results=False, print_summary=False)
     if not result:
         raise HTTPException(status_code=500, detail="Failed to run performance evaluation")
 
@@ -210,7 +188,7 @@ async def run_reliability_eval(
         raise HTTPException(status_code=400, detail="expected_tool_calls is required for reliability evaluations")
 
     if agent:
-        agent_response = await agent.arun(eval_run_input.input)
+        agent_response = await agent.arun(eval_run_input.input, stream=False)
         reliability_eval = ReliabilityEval(
             db=db,
             name=eval_run_input.name,
@@ -221,7 +199,7 @@ async def run_reliability_eval(
         model_provider = agent.model.provider if agent and agent.model else None
 
     elif team:
-        team_response = await team.arun(eval_run_input.input)
+        team_response = await team.arun(eval_run_input.input, stream=False)
         reliability_eval = ReliabilityEval(
             db=db,
             name=eval_run_input.name,

{agno-2.3.9 → agno-2.3.10}/agno/team/team.py

@@ -1541,6 +1541,8 @@ class Team:
             add_history_to_context=add_history_to_context,
             add_session_state_to_context=add_session_state_to_context,
             add_dependencies_to_context=add_dependencies_to_context,
+            stream=False,
+            stream_events=False,
         )
 
         # 3. Prepare run messages
@@ -1753,6 +1755,8 @@ class Team:
             add_history_to_context=add_history_to_context,
             add_session_state_to_context=add_session_state_to_context,
             add_dependencies_to_context=add_dependencies_to_context,
+            stream=True,
+            stream_events=stream_events,
         )
 
         # 3. Prepare run messages
@@ -2182,9 +2186,6 @@ class Team:
         if stream_events is None:
             stream_events = False if self.stream_events is None else self.stream_events
 
-        self.stream = self.stream or stream
-        self.stream_events = self.stream_events or stream_events
-
         self.model = cast(Model, self.model)
 
         if self.metadata is not None:
@@ -2394,6 +2395,8 @@ class Team:
             add_history_to_context=add_history_to_context,
             add_dependencies_to_context=add_dependencies_to_context,
             add_session_state_to_context=add_session_state_to_context,
+            stream=False,
+            stream_events=False,
         )
 
         # 5. Prepare run messages
@@ -2638,6 +2641,10 @@ class Team:
             files=run_input.files,
             debug_mode=debug_mode,
             add_history_to_context=add_history_to_context,
+            add_dependencies_to_context=add_dependencies_to_context,
+            add_session_state_to_context=add_session_state_to_context,
+            stream=True,
+            stream_events=stream_events,
         )
 
         # 6. Prepare run messages
@@ -3046,9 +3053,6 @@ class Team:
         if stream_events is None:
             stream_events = False if self.stream_events is None else self.stream_events
 
-        self.stream = self.stream or stream
-        self.stream_events = self.stream_events or stream_events
-
         self.model = cast(Model, self.model)
 
         if self.metadata is not None:
@@ -5397,6 +5401,8 @@ class Team:
         add_history_to_context: Optional[bool] = None,
         add_dependencies_to_context: Optional[bool] = None,
         add_session_state_to_context: Optional[bool] = None,
+        stream: Optional[bool] = None,
+        stream_events: Optional[bool] = None,
         check_mcp_tools: bool = True,
     ) -> List[Union[Function, dict]]:
         # Connect tools that require connection management
@@ -5491,8 +5497,8 @@ class Team:
                 team_run_context=team_run_context,
                 input=user_message_content,
                 user_id=user_id,
-                stream=self.stream or False,
-                stream_events=self.stream_events or False,
+                stream=stream or False,
+                stream_events=stream_events or False,
                 async_mode=async_mode,
                 images=images,  # type: ignore
                 videos=videos,  # type: ignore

{agno-2.3.9 → agno-2.3.10}/agno/tools/function.py

@@ -908,21 +908,16 @@ class FunctionCall(BaseModel):
             else:
                 result = self.function.entrypoint(**entrypoint_args, **self.arguments)  # type: ignore
 
-            updated_session_state = None
-            if entrypoint_args.get("run_context") is not None:
-                run_context = entrypoint_args.get("run_context")
-                updated_session_state = (
-                    run_context.session_state
-                    if run_context is not None and run_context.session_state is not None
-                    else None
-                )
-            else:
-                if self.function._session_state is not None:
-                    updated_session_state = self.function._session_state
-
             # Handle generator case
             if isgenerator(result):
                 self.result = result  # Store generator directly, can't cache
+                # For generators, don't capture updated_session_state yet -
+                # session_state is passed by reference, so mutations made during
+                # generator iteration are already reflected in the original dict.
+                # Returning None prevents stale state from being merged later.
+                execution_result = FunctionExecutionResult(
+                    status="success", result=self.result, updated_session_state=None
+                )
             else:
                 self.result = result
                 # Only cache non-generator results
@@ -931,9 +926,21 @@ class FunctionCall(BaseModel):
                     cache_file = self.function._get_cache_file_path(cache_key)
                     self.function._save_to_cache(cache_file, self.result)
 
-            execution_result = FunctionExecutionResult(
-                status="success", result=self.result, updated_session_state=updated_session_state
-            )
+                updated_session_state = None
+                if entrypoint_args.get("run_context") is not None:
+                    run_context = entrypoint_args.get("run_context")
+                    updated_session_state = (
+                        run_context.session_state
+                        if run_context is not None and run_context.session_state is not None
+                        else None
+                    )
+                else:
+                    if self.function._session_state is not None:
+                        updated_session_state = self.function._session_state
+
+                execution_result = FunctionExecutionResult(
+                    status="success", result=self.result, updated_session_state=updated_session_state
+                )
 
         except AgentRunException as e:
             log_debug(f"{e.__class__.__name__}: {e}")
@@ -1143,14 +1150,21 @@ class FunctionCall(BaseModel):
                 cache_file = self.function._get_cache_file_path(cache_key)
                 self.function._save_to_cache(cache_file, self.result)
 
-            updated_session_state = None
-            if entrypoint_args.get("run_context") is not None:
-                run_context = entrypoint_args.get("run_context")
-                updated_session_state = (
-                    run_context.session_state
-                    if run_context is not None and run_context.session_state is not None
-                    else None
-                )
+            # For generators, don't capture updated_session_state -
+            # session_state is passed by reference, so mutations made during
+            # generator iteration are already reflected in the original dict.
+            # Returning None prevents stale state from being merged later.
+            if isgenerator(self.result) or isasyncgen(self.result):
+                updated_session_state = None
+            else:
+                updated_session_state = None
+                if entrypoint_args.get("run_context") is not None:
+                    run_context = entrypoint_args.get("run_context")
+                    updated_session_state = (
+                        run_context.session_state
+                        if run_context is not None and run_context.session_state is not None
+                        else None
+                    )
 
             execution_result = FunctionExecutionResult(
                 status="success", result=self.result, updated_session_state=updated_session_state