Merge branch 'main' into live-error-toast

src/google/adk/auth/auth_credential.py

@@ -66,7 +66,8 @@ class OAuth2Auth(BaseModelWithConfig):
   redirect_uri: Optional[str] = None
   auth_response_uri: Optional[str] = None
   auth_code: Optional[str] = None
-  token: Optional[Dict[str, Any]] = None
+  access_token: Optional[str] = None
+  refresh_token: Optional[str] = None
 
 
 class ServiceAccountCredential(BaseModelWithConfig):

src/google/adk/auth/auth_handler.py

@@ -82,7 +82,8 @@ class AuthHandler:
         or not auth_credential.oauth2
         or not auth_credential.oauth2.client_id
         or not auth_credential.oauth2.client_secret
-        or auth_credential.oauth2.token
+        or auth_credential.oauth2.access_token
+        or auth_credential.oauth2.refresh_token
     ):
       return self.auth_config.exchanged_auth_credential
 
@@ -93,7 +94,7 @@ class AuthHandler:
         redirect_uri=auth_credential.oauth2.redirect_uri,
         state=auth_credential.oauth2.state,
     )
-    token = client.fetch_token(
+    tokens = client.fetch_token(
         token_endpoint,
         authorization_response=auth_credential.oauth2.auth_response_uri,
         code=auth_credential.oauth2.auth_code,
@@ -102,7 +103,10 @@ class AuthHandler:
 
     updated_credential = AuthCredential(
         auth_type=AuthCredentialTypes.OAUTH2,
-        oauth2=OAuth2Auth(token=dict(token)),
+        oauth2=OAuth2Auth(
+            access_token=tokens.get("access_token"),
+            refresh_token=tokens.get("refresh_token"),
+        ),
     )
     return updated_credential
 

src/google/adk/cli/cli_deploy.py

@@ -54,7 +54,7 @@ COPY "agents/{app_name}/" "/app/agents/{app_name}/"
 
 EXPOSE {port}
 
-CMD adk {command} --port={port} {trace_to_cloud_option} "/app/agents"
+CMD adk {command} --port={port} {session_db_option} {trace_to_cloud_option} "/app/agents"
 """
 
 
@@ -85,6 +85,7 @@ def to_cloud_run(
     trace_to_cloud: bool,
     with_ui: bool,
     verbosity: str,
+    session_db_url: str,
 ):
   """Deploys an agent to Google Cloud Run.
 
@@ -112,6 +113,7 @@ def to_cloud_run(
     trace_to_cloud: Whether to enable Cloud Trace.
     with_ui: Whether to deploy with UI.
     verbosity: The verbosity level of the CLI.
+    session_db_url: The database URL to connect the session.
   """
   app_name = app_name or os.path.basename(agent_folder)
 
@@ -144,6 +146,9 @@ def to_cloud_run(
         port=port,
         command='web' if with_ui else 'api_server',
         install_agent_deps=install_agent_deps,
+        session_db_option=f'--session_db_url={session_db_url}'
+        if session_db_url
+        else '',
         trace_to_cloud_option='--trace_to_cloud' if trace_to_cloud else '',
     )
     dockerfile_path = os.path.join(temp_folder, 'Dockerfile')

src/google/adk/cli/cli_tools_click.py

@@ -245,12 +245,13 @@ def cli_eval(
 @click.option(
     "--session_db_url",
     help=(
-        "Optional. The database URL to store the session.\n\n  - Use"
-        " 'agentengine://<agent_engine_resource_id>' to connect to Vertex"
-        " managed session service.\n\n  - Use 'sqlite://<path_to_sqlite_file>'"
-        " to connect to a SQLite DB.\n\n  - See"
-        " https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls"
-        " for more details on supported DB URLs."
+        """Optional. The database URL to store the session.
+
+  - Use 'agentengine://<agent_engine_resource_id>' to connect to Agent Engine sessions.
+
+  - Use 'sqlite://<path_to_sqlite_file>' to connect to a SQLite DB.
+
+  - See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls for more details on supported DB URLs."""
     ),
 )
 @click.option(
@@ -366,12 +367,13 @@ def cli_web(
 @click.option(
     "--session_db_url",
     help=(
-        "Optional. The database URL to store the session.\n\n  - Use"
-        " 'agentengine://<agent_engine_resource_id>' to connect to Vertex"
-        " managed session service.\n\n  - Use 'sqlite://<path_to_sqlite_file>'"
-        " to connect to a SQLite DB.\n\n  - See"
-        " https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls"
-        " for more details on supported DB URLs."
+        """Optional. The database URL to store the session.
+
+  - Use 'agentengine://<agent_engine_resource_id>' to connect to Agent Engine sessions.
+
+  - Use 'sqlite://<path_to_sqlite_file>' to connect to a SQLite DB.
+
+  - See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls for more details on supported DB URLs."""
     ),
 )
 @click.option(
@@ -541,6 +543,18 @@ def cli_api_server(
     default="WARNING",
     help="Optional. Override the default verbosity level.",
 )
+@click.option(
+    "--session_db_url",
+    help=(
+        """Optional. The database URL to store the session.
+
+  - Use 'agentengine://<agent_engine_resource_id>' to connect to Agent Engine sessions.
+
+  - Use 'sqlite://<path_to_sqlite_file>' to connect to a SQLite DB.
+
+  - See https://docs.sqlalchemy.org/en/20/core/engines.html#backend-specific-urls for more details on supported DB URLs."""
+    ),
+)
 @click.argument(
     "agent",
     type=click.Path(
@@ -558,6 +572,7 @@ def cli_deploy_cloud_run(
     trace_to_cloud: bool,
     with_ui: bool,
     verbosity: str,
+    session_db_url: str,
 ):
   """Deploys an agent to Cloud Run.
 
@@ -579,6 +594,7 @@ def cli_deploy_cloud_run(
         trace_to_cloud=trace_to_cloud,
         with_ui=with_ui,
         verbosity=verbosity,
+        session_db_url=session_db_url,
     )
   except Exception as e:
     click.secho(f"Deploy failed: {e}", fg="red", err=True)

src/google/adk/cli/fast_api.py

@@ -756,6 +756,12 @@ def get_fast_api_app(
     except Exception as e:
       logger.exception("Error during live websocket communication: %s", e)
       traceback.print_exc()
+      WEBSOCKET_INTERNAL_ERROR_CODE = 1011
+      WEBSOCKET_MAX_BYTES_FOR_REASON = 123
+      await websocket.close(
+          code=WEBSOCKET_INTERNAL_ERROR_CODE,
+          reason=str(e)[:WEBSOCKET_MAX_BYTES_FOR_REASON],
+      )
     finally:
       for task in pending:
         task.cancel()

src/google/adk/tools/openapi_tool/auth/credential_exchangers/oauth2_exchanger.py

@@ -69,7 +69,7 @@ class OAuth2CredentialExchanger(BaseAuthCredentialExchanger):
         HTTP bearer token cannot be generated, return the original credential.
     """
 
-    if "access_token" not in auth_credential.oauth2.token:
+    if not auth_credential.oauth2.access_token:
       return auth_credential
 
     # Return the access token as a bearer token.
@@ -78,7 +78,7 @@ class OAuth2CredentialExchanger(BaseAuthCredentialExchanger):
         http=HttpAuth(
             scheme="bearer",
             credentials=HttpCredentials(
-                token=auth_credential.oauth2.token["access_token"]
+                token=auth_credential.oauth2.access_token
             ),
         ),
     )
@@ -111,7 +111,7 @@ class OAuth2CredentialExchanger(BaseAuthCredentialExchanger):
       return auth_credential
 
     # If access token is exchanged, exchange a HTTPBearer token.
-    if auth_credential.oauth2.token:
+    if auth_credential.oauth2.access_token:
       return self.generate_auth_token(auth_credential)
 
     return None