google-adk 0.5.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

google/adk/cli/fast_api.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 import asyncio
 from contextlib import asynccontextmanager
 import importlib
@@ -20,8 +21,9 @@ import json
 import logging
 import os
 from pathlib import Path
-import re
+import signal
 import sys
+import time
 import traceback
 import typing
 from typing import Any
@@ -30,7 +32,6 @@ from typing import Literal
 from typing import Optional
 
 import click
-from click import Tuple
 from fastapi import FastAPI
 from fastapi import HTTPException
 from fastapi import Query
@@ -48,16 +49,22 @@ from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter
 from opentelemetry.sdk.trace import export
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace import TracerProvider
-from pydantic import BaseModel
+from pydantic import Field
 from pydantic import ValidationError
 from starlette.types import Lifespan
+from typing_extensions import override
 
 from ..agents import RunConfig
+from ..agents.base_agent import BaseAgent
 from ..agents.live_request_queue import LiveRequest
 from ..agents.live_request_queue import LiveRequestQueue
 from ..agents.llm_agent import Agent
+from ..agents.llm_agent import LlmAgent
 from ..agents.run_config import StreamingMode
 from ..artifacts import InMemoryArtifactService
+from ..evaluation.eval_case import EvalCase
+from ..evaluation.eval_case import SessionInput
+from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
 from ..events.event import Event
 from ..memory.in_memory_memory_service import InMemoryMemoryService
 from ..runners import Runner
@@ -65,17 +72,23 @@ from ..sessions.database_session_service import DatabaseSessionService
 from ..sessions.in_memory_session_service import InMemorySessionService
 from ..sessions.session import Session
 from ..sessions.vertex_ai_session_service import VertexAiSessionService
+from ..tools.base_toolset import BaseToolset
 from .cli_eval import EVAL_SESSION_ID_PREFIX
+from .cli_eval import EvalCaseResult
 from .cli_eval import EvalMetric
 from .cli_eval import EvalMetricResult
+from .cli_eval import EvalMetricResultPerInvocation
+from .cli_eval import EvalSetResult
 from .cli_eval import EvalStatus
+from .utils import common
 from .utils import create_empty_state
 from .utils import envs
 from .utils import evals
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 _EVAL_SET_FILE_EXTENSION = ".evalset.json"
+_EVAL_SET_RESULT_FILE_EXTENSION = ".evalset_result.json"
 
 
 class ApiServerSpanExporter(export.SpanExporter):
@@ -103,7 +116,45 @@ class ApiServerSpanExporter(export.SpanExporter):
     return True
 
 
-class AgentRunRequest(BaseModel):
+class InMemoryExporter(export.SpanExporter):
+
+  def __init__(self, trace_dict):
+    super().__init__()
+    self._spans = []
+    self.trace_dict = trace_dict
+
+  @override
+  def export(
+      self, spans: typing.Sequence[ReadableSpan]
+  ) -> export.SpanExportResult:
+    for span in spans:
+      trace_id = span.context.trace_id
+      if span.name == "call_llm":
+        attributes = dict(span.attributes)
+        session_id = attributes.get("gcp.vertex.agent.session_id", None)
+        if session_id:
+          if session_id not in self.trace_dict:
+            self.trace_dict[session_id] = [trace_id]
+          else:
+            self.trace_dict[session_id] += [trace_id]
+    self._spans.extend(spans)
+    return export.SpanExportResult.SUCCESS
+
+  @override
+  def force_flush(self, timeout_millis: int = 30000) -> bool:
+    return True
+
+  def get_finished_spans(self, session_id: str):
+    trace_ids = self.trace_dict.get(session_id, None)
+    if trace_ids is None or not trace_ids:
+      return []
+    return [x for x in self._spans if x.context.trace_id in trace_ids]
+
+  def clear(self):
+    self._spans.clear()
+
+
+class AgentRunRequest(common.BaseModel):
   app_name: str
   user_id: str
   session_id: str
@@ -111,25 +162,38 @@ class AgentRunRequest(BaseModel):
   streaming: bool = False
 
 
-class AddSessionToEvalSetRequest(BaseModel):
+class AddSessionToEvalSetRequest(common.BaseModel):
   eval_id: str
   session_id: str
   user_id: str
 
 
-class RunEvalRequest(BaseModel):
+class RunEvalRequest(common.BaseModel):
   eval_ids: list[str]  # if empty, then all evals in the eval set are run.
   eval_metrics: list[EvalMetric]
 
 
-class RunEvalResult(BaseModel):
+class RunEvalResult(common.BaseModel):
+  eval_set_file: str
   eval_set_id: str
   eval_id: str
   final_eval_status: EvalStatus
-  eval_metric_results: list[tuple[EvalMetric, EvalMetricResult]]
+  eval_metric_results: list[tuple[EvalMetric, EvalMetricResult]] = Field(
+      deprecated=True,
+      description=(
+          "This field is deprecated, use overall_eval_metric_results instead."
+      ),
+  )
+  overall_eval_metric_results: list[EvalMetricResult]
+  eval_metric_result_per_invocation: list[EvalMetricResultPerInvocation]
+  user_id: str
   session_id: str
 
 
+class GetEventGraphResult(common.BaseModel):
+  dot_src: str
+
+
 def get_fast_api_app(
     *,
     agent_dir: str,
@@ -141,12 +205,15 @@ def get_fast_api_app(
 ) -> FastAPI:
   # InMemory tracing dict.
   trace_dict: dict[str, Any] = {}
+  session_trace_dict: dict[str, Any] = {}
 
   # Set up tracing in the FastAPI server.
   provider = TracerProvider()
   provider.add_span_processor(
       export.SimpleSpanProcessor(ApiServerSpanExporter(trace_dict))
   )
+  memory_exporter = InMemoryExporter(session_trace_dict)
+  provider.add_span_processor(export.SimpleSpanProcessor(memory_exporter))
   if trace_to_cloud:
     envs.load_dotenv_for_agent("", agent_dir)
     if project_id := os.environ.get("GOOGLE_CLOUD_PROJECT", None):
@@ -155,26 +222,82 @@ def get_fast_api_app(
       )
       provider.add_span_processor(processor)
     else:
-      logging.warning(
+      logger.warning(
           "GOOGLE_CLOUD_PROJECT environment variable is not set. Tracing will"
           " not be enabled."
       )
 
   trace.set_tracer_provider(provider)
 
-  exit_stacks = []
+  toolsets_to_close: set[BaseToolset] = set()
 
   @asynccontextmanager
   async def internal_lifespan(app: FastAPI):
-    if lifespan:
-      async with lifespan(app) as lifespan_context:
+    # Set up signal handlers for graceful shutdown
+    original_sigterm = signal.getsignal(signal.SIGTERM)
+    original_sigint = signal.getsignal(signal.SIGINT)
+
+    def cleanup_handler(sig, frame):
+      # Log the signal
+      logger.info("Received signal %s, performing pre-shutdown cleanup", sig)
+      # Do synchronous cleanup if needed
+      # Then call original handler if it exists
+      if sig == signal.SIGTERM and callable(original_sigterm):
+        original_sigterm(sig, frame)
+      elif sig == signal.SIGINT and callable(original_sigint):
+        original_sigint(sig, frame)
+
+    # Install cleanup handlers
+    signal.signal(signal.SIGTERM, cleanup_handler)
+    signal.signal(signal.SIGINT, cleanup_handler)
+
+    try:
+      if lifespan:
+        async with lifespan(app) as lifespan_context:
+          yield lifespan_context
+      else:
         yield
+    finally:
+      # During shutdown, properly clean up all toolsets
+      logger.info(
+          "Server shutdown initiated, cleaning up %s toolsets",
+          len(toolsets_to_close),
+      )
 
-        if exit_stacks:
-          for stack in exit_stacks:
-            await stack.aclose()
-    else:
-      yield
+      # Create tasks for all toolset closures to run concurrently
+      cleanup_tasks = []
+      for toolset in toolsets_to_close:
+        task = asyncio.create_task(close_toolset_safely(toolset))
+        cleanup_tasks.append(task)
+
+      if cleanup_tasks:
+        # Wait for all cleanup tasks with timeout
+        done, pending = await asyncio.wait(
+            cleanup_tasks,
+            timeout=10.0,  # 10 second timeout for cleanup
+            return_when=asyncio.ALL_COMPLETED,
+        )
+
+        # If any tasks are still pending, log it
+        if pending:
+          logger.warning(
+              f"{len(pending)} toolset cleanup tasks didn't complete in time"
+          )
+          for task in pending:
+            task.cancel()
+
+      # Restore original signal handlers
+      signal.signal(signal.SIGTERM, original_sigterm)
+      signal.signal(signal.SIGINT, original_sigint)
+
+  async def close_toolset_safely(toolset):
+    """Safely close a toolset with error handling."""
+    try:
+      logger.info(f"Closing toolset: {type(toolset).__name__}")
+      await toolset.close()
+      logger.info(f"Successfully closed toolset: {type(toolset).__name__}")
+    except Exception as e:
+      logger.error(f"Error closing toolset {type(toolset).__name__}: {e}")
 
   # Run the FastAPI server.
   app = FastAPI(lifespan=internal_lifespan)
@@ -198,6 +321,8 @@ def get_fast_api_app(
   artifact_service = InMemoryArtifactService()
   memory_service = InMemoryMemoryService()
 
+  eval_sets_manager = LocalEvalSetsManager(agent_dir=agent_dir)
+
   # Build the Session service
   agent_engine_id = ""
   if session_db_url:
@@ -240,14 +365,34 @@ def get_fast_api_app(
       raise HTTPException(status_code=404, detail="Trace not found")
     return event_dict
 
+  @app.get("/debug/trace/session/{session_id}")
+  def get_session_trace(session_id: str) -> Any:
+    spans = memory_exporter.get_finished_spans(session_id)
+    if not spans:
+      return []
+    return [
+        {
+            "name": s.name,
+            "span_id": s.context.span_id,
+            "trace_id": s.context.trace_id,
+            "start_time": s.start_time,
+            "end_time": s.end_time,
+            "attributes": dict(s.attributes),
+            "parent_span_id": s.parent.span_id if s.parent else None,
+        }
+        for s in spans
+    ]
+
   @app.get(
       "/apps/{app_name}/users/{user_id}/sessions/{session_id}",
       response_model_exclude_none=True,
   )
-  def get_session(app_name: str, user_id: str, session_id: str) -> Session:
+  async def get_session(
+      app_name: str, user_id: str, session_id: str
+  ) -> Session:
     # Connect to managed session if agent_engine_id is set.
     app_name = agent_engine_id if agent_engine_id else app_name
-    session = session_service.get_session(
+    session = await session_service.get_session(
         app_name=app_name, user_id=user_id, session_id=session_id
     )
     if not session:
@@ -258,14 +403,15 @@ def get_fast_api_app(
       "/apps/{app_name}/users/{user_id}/sessions",
       response_model_exclude_none=True,
   )
-  def list_sessions(app_name: str, user_id: str) -> list[Session]:
+  async def list_sessions(app_name: str, user_id: str) -> list[Session]:
     # Connect to managed session if agent_engine_id is set.
     app_name = agent_engine_id if agent_engine_id else app_name
+    list_sessions_response = await session_service.list_sessions(
+        app_name=app_name, user_id=user_id
+    )
     return [
         session
-        for session in session_service.list_sessions(
-            app_name=app_name, user_id=user_id
-        ).sessions
+        for session in list_sessions_response.sessions
         # Remove sessions that were generated as a part of Eval.
         if not session.id.startswith(EVAL_SESSION_ID_PREFIX)
     ]
@@ -274,7 +420,7 @@ def get_fast_api_app(
       "/apps/{app_name}/users/{user_id}/sessions/{session_id}",
       response_model_exclude_none=True,
   )
-  def create_session_with_id(
+  async def create_session_with_id(
       app_name: str,
       user_id: str,
       session_id: str,
@@ -283,7 +429,7 @@ def get_fast_api_app(
     # Connect to managed session if agent_engine_id is set.
     app_name = agent_engine_id if agent_engine_id else app_name
     if (
-        session_service.get_session(
+        await session_service.get_session(
             app_name=app_name, user_id=user_id, session_id=session_id
         )
         is not None
@@ -292,9 +438,8 @@ def get_fast_api_app(
       raise HTTPException(
           status_code=400, detail=f"Session already exists: {session_id}"
       )
-
     logger.info("New session created: %s", session_id)
-    return session_service.create_session(
+    return await session_service.create_session(
         app_name=app_name, user_id=user_id, state=state, session_id=session_id
     )
 
@@ -302,16 +447,15 @@ def get_fast_api_app(
       "/apps/{app_name}/users/{user_id}/sessions",
       response_model_exclude_none=True,
   )
-  def create_session(
+  async def create_session(
       app_name: str,
       user_id: str,
       state: Optional[dict[str, Any]] = None,
   ) -> Session:
     # Connect to managed session if agent_engine_id is set.
     app_name = agent_engine_id if agent_engine_id else app_name
-
     logger.info("New session created")
-    return session_service.create_session(
+    return await session_service.create_session(
         app_name=app_name, user_id=user_id, state=state
     )
 
@@ -331,28 +475,13 @@ def get_fast_api_app(
       eval_set_id: str,
   ):
     """Creates an eval set, given the id."""
-    pattern = r"^[a-zA-Z0-9_]+$"
-    if not bool(re.fullmatch(pattern, eval_set_id)):
+    try:
+      eval_sets_manager.create_eval_set(app_name, eval_set_id)
+    except ValueError as ve:
       raise HTTPException(
           status_code=400,
-          detail=(
-              f"Invalid eval set id. Eval set id should have the `{pattern}`"
-              " format"
-          ),
-      )
-    # Define the file path
-    new_eval_set_path = _get_eval_set_file_path(
-        app_name, agent_dir, eval_set_id
-    )
-
-    logger.info("Creating eval set file `%s`", new_eval_set_path)
-
-    if not os.path.exists(new_eval_set_path):
-      # Write the JSON string to the file
-      logger.info("Eval set file doesn't exist, we will create a new one.")
-      with open(new_eval_set_path, "w") as f:
-        empty_content = json.dumps([], indent=2)
-        f.write(empty_content)
+          detail=str(ve),
+      ) from ve
 
   @app.get(
       "/apps/{app_name}/eval_sets",
@@ -360,15 +489,7 @@ def get_fast_api_app(
   )
   def list_eval_sets(app_name: str) -> list[str]:
     """Lists all eval sets for the given app."""
-    eval_set_file_path = os.path.join(agent_dir, app_name)
-    eval_sets = []
-    for file in os.listdir(eval_set_file_path):
-      if file.endswith(_EVAL_SET_FILE_EXTENSION):
-        eval_sets.append(
-            os.path.basename(file).removesuffix(_EVAL_SET_FILE_EXTENSION)
-        )
-
-    return sorted(eval_sets)
+    return eval_sets_manager.list_eval_sets(app_name)
 
   @app.post(
       "/apps/{app_name}/eval_sets/{eval_set_id}/add_session",
@@ -377,54 +498,33 @@ def get_fast_api_app(
   async def add_session_to_eval_set(
       app_name: str, eval_set_id: str, req: AddSessionToEvalSetRequest
   ):
-    pattern = r"^[a-zA-Z0-9_]+$"
-    if not bool(re.fullmatch(pattern, req.eval_id)):
-      raise HTTPException(
-          status_code=400,
-          detail=f"Invalid eval id. Eval id should have the `{pattern}` format",
-      )
-
     # Get the session
-    session = session_service.get_session(
+    session = await session_service.get_session(
         app_name=app_name, user_id=req.user_id, session_id=req.session_id
     )
     assert session, "Session not found."
-    # Load the eval set file data
-    eval_set_file_path = _get_eval_set_file_path(
-        app_name, agent_dir, eval_set_id
-    )
-    with open(eval_set_file_path, "r") as file:
-      eval_set_data = json.load(file)  # Load JSON into a list
 
-    if [x for x in eval_set_data if x["name"] == req.eval_id]:
-      raise HTTPException(
-          status_code=400,
-          detail=(
-              f"Eval id `{req.eval_id}` already exists in `{eval_set_id}`"
-              " eval set."
-          ),
-      )
-
-    # Convert the session data to evaluation format
-    test_data = evals.convert_session_to_eval_format(session)
+    # Convert the session data to eval invocations
+    invocations = evals.convert_session_to_eval_invocations(session)
 
     # Populate the session with initial session state.
     initial_session_state = create_empty_state(
         await _get_root_agent_async(app_name)
     )
 
-    eval_set_data.append({
-        "name": req.eval_id,
-        "data": test_data,
-        "initial_session": {
-            "state": initial_session_state,
-            "app_name": app_name,
-            "user_id": req.user_id,
-        },
-    })
-    # Serialize the test data to JSON and write to the eval set file.
-    with open(eval_set_file_path, "w") as f:
-      f.write(json.dumps(eval_set_data, indent=2))
+    new_eval_case = EvalCase(
+        eval_id=req.eval_id,
+        conversation=invocations,
+        session_input=SessionInput(
+            app_name=app_name, user_id=req.user_id, state=initial_session_state
+        ),
+        creation_timestamp=time.time(),
+    )
+
+    try:
+      eval_sets_manager.add_eval_case(app_name, eval_set_id, new_eval_case)
+    except ValueError as ve:
+      raise HTTPException(status_code=400, detail=str(ve)) from ve
 
   @app.get(
       "/apps/{app_name}/eval_sets/{eval_set_id}/evals",
@@ -435,14 +535,9 @@ def get_fast_api_app(
       eval_set_id: str,
   ) -> list[str]:
     """Lists all evals in an eval set."""
-    # Load the eval set file data
-    eval_set_file_path = _get_eval_set_file_path(
-        app_name, agent_dir, eval_set_id
-    )
-    with open(eval_set_file_path, "r") as file:
-      eval_set_data = json.load(file)  # Load JSON into a list
+    eval_set_data = eval_sets_manager.get_eval_set(app_name, eval_set_id)
 
-    return sorted([x["name"] for x in eval_set_data])
+    return sorted([x.eval_id for x in eval_set_data.eval_cases])
 
   @app.post(
       "/apps/{app_name}/eval_sets/{eval_set_id}/run_eval",
@@ -451,51 +546,136 @@ def get_fast_api_app(
   async def run_eval(
       app_name: str, eval_set_id: str, req: RunEvalRequest
   ) -> list[RunEvalResult]:
+    """Runs an eval given the details in the eval request."""
     from .cli_eval import run_evals
 
-    """Runs an eval given the details in the eval request."""
     # Create a mapping from eval set file to all the evals that needed to be
     # run.
-    eval_set_file_path = _get_eval_set_file_path(
-        app_name, agent_dir, eval_set_id
-    )
-    eval_set_to_evals = {eval_set_file_path: req.eval_ids}
+    envs.load_dotenv_for_agent(os.path.basename(app_name), agent_dir)
 
-    if not req.eval_ids:
-      logger.info(
-          "Eval ids to run list is empty. We will all evals in the eval set."
-      )
-    root_agent = await _get_root_agent_async(app_name)
-    eval_results = list(
-        await run_evals(
-            eval_set_to_evals,
-            root_agent,
-            getattr(root_agent, "reset_data", None),
-            req.eval_metrics,
-            session_service=session_service,
-            artifact_service=artifact_service,
-        )
-    )
+    eval_set = eval_sets_manager.get_eval_set(app_name, eval_set_id)
+
+    if req.eval_ids:
+      eval_cases = [e for e in eval_set.eval_cases if e.eval_id in req.eval_ids]
+      eval_set_to_evals = {eval_set_id: eval_cases}
+    else:
+      logger.info("Eval ids to run list is empty. We will run all eval cases.")
+      eval_set_to_evals = {eval_set_id: eval_set.eval_cases}
 
+    root_agent = await _get_root_agent_async(app_name)
     run_eval_results = []
-    for eval_result in eval_results:
+    eval_case_results = []
+    async for eval_case_result in run_evals(
+        eval_set_to_evals,
+        root_agent,
+        getattr(root_agent, "reset_data", None),
+        req.eval_metrics,
+        session_service=session_service,
+        artifact_service=artifact_service,
+    ):
       run_eval_results.append(
           RunEvalResult(
               app_name=app_name,
+              eval_set_file=eval_case_result.eval_set_file,
               eval_set_id=eval_set_id,
-              eval_id=eval_result.eval_id,
-              final_eval_status=eval_result.final_eval_status,
-              eval_metric_results=eval_result.eval_metric_results,
-              session_id=eval_result.session_id,
+              eval_id=eval_case_result.eval_id,
+              final_eval_status=eval_case_result.final_eval_status,
+              eval_metric_results=eval_case_result.eval_metric_results,
+              overall_eval_metric_results=eval_case_result.overall_eval_metric_results,
+              eval_metric_result_per_invocation=eval_case_result.eval_metric_result_per_invocation,
+              user_id=eval_case_result.user_id,
+              session_id=eval_case_result.session_id,
           )
       )
+      eval_case_result.session_details = await session_service.get_session(
+          app_name=app_name,
+          user_id=eval_case_result.user_id,
+          session_id=eval_case_result.session_id,
+      )
+      eval_case_results.append(eval_case_result)
+
+    timestamp = time.time()
+    eval_set_result_name = app_name + "_" + eval_set_id + "_" + str(timestamp)
+    eval_set_result = EvalSetResult(
+        eval_set_result_id=eval_set_result_name,
+        eval_set_result_name=eval_set_result_name,
+        eval_set_id=eval_set_id,
+        eval_case_results=eval_case_results,
+        creation_timestamp=timestamp,
+    )
+
+    # Write eval result file, with eval_set_result_name.
+    app_eval_history_dir = os.path.join(
+        agent_dir, app_name, ".adk", "eval_history"
+    )
+    if not os.path.exists(app_eval_history_dir):
+      os.makedirs(app_eval_history_dir)
+    # Convert to json and write to file.
+    eval_set_result_json = eval_set_result.model_dump_json()
+    eval_set_result_file_path = os.path.join(
+        app_eval_history_dir,
+        eval_set_result_name + _EVAL_SET_RESULT_FILE_EXTENSION,
+    )
+    logger.info("Writing eval result to file: %s", eval_set_result_file_path)
+    with open(eval_set_result_file_path, "w") as f:
+      f.write(json.dumps(eval_set_result_json, indent=2))
+
     return run_eval_results
 
+  @app.get(
+      "/apps/{app_name}/eval_results/{eval_result_id}",
+      response_model_exclude_none=True,
+  )
+  def get_eval_result(
+      app_name: str,
+      eval_result_id: str,
+  ) -> EvalSetResult:
+    """Gets the eval result for the given eval id."""
+    # Load the eval set file data
+    maybe_eval_result_file_path = (
+        os.path.join(
+            agent_dir, app_name, ".adk", "eval_history", eval_result_id
+        )
+        + _EVAL_SET_RESULT_FILE_EXTENSION
+    )
+    if not os.path.exists(maybe_eval_result_file_path):
+      raise HTTPException(
+          status_code=404,
+          detail=f"Eval result `{eval_result_id}` not found.",
+      )
+    with open(maybe_eval_result_file_path, "r") as file:
+      eval_result_data = json.load(file)  # Load JSON into a list
+    try:
+      eval_result = EvalSetResult.model_validate_json(eval_result_data)
+      return eval_result
+    except ValidationError as e:
+      logger.exception("get_eval_result validation error: %s", e)
+
+  @app.get(
+      "/apps/{app_name}/eval_results",
+      response_model_exclude_none=True,
+  )
+  def list_eval_results(app_name: str) -> list[str]:
+    """Lists all eval results for the given app."""
+    app_eval_history_directory = os.path.join(
+        agent_dir, app_name, ".adk", "eval_history"
+    )
+
+    if not os.path.exists(app_eval_history_directory):
+      return []
+
+    eval_result_files = [
+        file.removesuffix(_EVAL_SET_RESULT_FILE_EXTENSION)
+        for file in os.listdir(app_eval_history_directory)
+        if file.endswith(_EVAL_SET_RESULT_FILE_EXTENSION)
+    ]
+    return eval_result_files
+
   @app.delete("/apps/{app_name}/users/{user_id}/sessions/{session_id}")
-  def delete_session(app_name: str, user_id: str, session_id: str):
+  async def delete_session(app_name: str, user_id: str, session_id: str):
     # Connect to managed session if agent_engine_id is set.
     app_name = agent_engine_id if agent_engine_id else app_name
-    session_service.delete_session(
+    await session_service.delete_session(
         app_name=app_name, user_id=user_id, session_id=session_id
     )
 
@@ -590,7 +770,7 @@ def get_fast_api_app(
   async def agent_run(req: AgentRunRequest) -> list[Event]:
     # Connect to managed session if agent_engine_id is set.
     app_id = agent_engine_id if agent_engine_id else req.app_name
-    session = session_service.get_session(
+    session = await session_service.get_session(
         app_name=app_id, user_id=req.user_id, session_id=req.session_id
     )
     if not session:
@@ -612,7 +792,7 @@ def get_fast_api_app(
     # Connect to managed session if agent_engine_id is set.
     app_id = agent_engine_id if agent_engine_id else req.app_name
     # SSE endpoint
-    session = session_service.get_session(
+    session = await session_service.get_session(
         app_name=app_id, user_id=req.user_id, session_id=req.session_id
     )
     if not session:
@@ -653,7 +833,7 @@ def get_fast_api_app(
   ):
     # Connect to managed session if agent_engine_id is set.
     app_id = agent_engine_id if agent_engine_id else app_name
-    session = session_service.get_session(
+    session = await session_service.get_session(
         app_name=app_id, user_id=user_id, session_id=session_id
     )
     session_events = session.events if session else []
@@ -673,7 +853,7 @@ def get_fast_api_app(
         from_name = event.author
         to_name = function_call.name
         function_call_highlights.append((from_name, to_name))
-        dot_graph = agent_graph.get_agent_graph(
+        dot_graph = await agent_graph.get_agent_graph(
             root_agent, function_call_highlights
         )
     elif function_responses:
@@ -682,17 +862,17 @@ def get_fast_api_app(
         from_name = function_response.name
         to_name = event.author
         function_responses_highlights.append((from_name, to_name))
-        dot_graph = agent_graph.get_agent_graph(
+        dot_graph = await agent_graph.get_agent_graph(
             root_agent, function_responses_highlights
         )
     else:
       from_name = event.author
       to_name = ""
-      dot_graph = agent_graph.get_agent_graph(
+      dot_graph = await agent_graph.get_agent_graph(
           root_agent, [(from_name, to_name)]
       )
     if dot_graph and isinstance(dot_graph, graphviz.Digraph):
-      return {"dot_src": dot_graph.source}
+      return GetEventGraphResult(dot_src=dot_graph.source)
     else:
       return {}
 
@@ -710,7 +890,7 @@ def get_fast_api_app(
 
     # Connect to managed session if agent_engine_id is set.
     app_id = agent_engine_id if agent_engine_id else app_name
-    session = session_service.get_session(
+    session = await session_service.get_session(
         app_name=app_id, user_id=user_id, session_id=session_id
     )
     if not session:
@@ -766,6 +946,16 @@ def get_fast_api_app(
       for task in pending:
         task.cancel()
 
+  def _get_all_toolsets(agent: BaseAgent) -> set[BaseToolset]:
+    toolsets = set()
+    if isinstance(agent, LlmAgent):
+      for tool_union in agent.tools:
+        if isinstance(tool_union, BaseToolset):
+          toolsets.add(tool_union)
+    for sub_agent in agent.sub_agents:
+      toolsets.update(_get_all_toolsets(sub_agent))
+    return toolsets
+
   async def _get_root_agent_async(app_name: str) -> Agent:
     """Returns the root agent for the given app."""
     if app_name in root_agent_dict:
@@ -776,16 +966,8 @@ def get_fast_api_app(
     else:
       raise ValueError(f'Unable to find "root_agent" from {app_name}.')
 
-    # Handle an awaitable root agent and await for the actual agent.
-    if inspect.isawaitable(root_agent):
-      try:
-        agent, exit_stack = await root_agent
-        exit_stacks.append(exit_stack)
-        root_agent = agent
-      except Exception as e:
-        raise RuntimeError(f"error getting root agent, {e}") from e
-
     root_agent_dict[app_name] = root_agent
+    toolsets_to_close.update(_get_all_toolsets(root_agent))
     return root_agent
 
   async def _get_runner_async(app_name: str) -> Runner: