Skip to content

Commit d93cdf3

Browse files
DeanChensjDeanChensj
committed
feat(workflow): add FunctionNode auth_config
Add auth_config field to FunctionNode that pauses for user authentication before running. The auth gate lives inside FunctionNode.run(), keeping _node_runner generic. Utility functions in _workflow_hitl_utils: - create_auth_request_event: builds adk_request_credential FC event with user-facing message - process_auth_resume: stores credentials with plain value fallback - has_auth_credential: checks if credentials exist in state Unwrap {"result":...} for auth responses in _process_resumptions, consistent with request_input handling. Includes sample agent and tests for auth flow and credential reuse. Change-Id: Ib81bf9cf63b8282efd7898def17f6e47a95b2d05
1 parent 891e8b7 commit d93cdf3

11 files changed

Lines changed: 670 additions & 3 deletions

File tree

‎contributing/workflow_samples/auth_config/README.md‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
# ADK Workflow Auth Config Sample
2+
3+
## Overview
4+
5+
This sample demonstrates how to use `auth_config` on a `FunctionNode` to require user authentication before the node runs.
6+
7+
When a node has `auth_config`, the workflow automatically:
8+
1. Pauses the node and emits an `adk_request_credential` FunctionCall event
9+
2. The invocation ends — the node is marked as waiting
10+
3. The client sends a new request with the credential as a FunctionResponse
11+
4. The workflow stores the credential in session state and re-runs the node
12+
13+
The **ADK web UI** (`adk web`) handles step 3 automatically — it recognizes auth
14+
requests and presents an auth dialog. If you use a custom client, you need to
15+
handle the `adk_request_credential` FunctionCall and respond with the credential
16+
yourself.
17+
18+
This sample uses **API key** authentication (the simplest credential type).
19+
20+
## No External Setup Required
21+
22+
This sample uses a mock weather lookup. No external API key or server is needed. When the auth UI prompts for a key, you can enter any value (e.g., `my-test-key-123`).
23+
24+
## Sample Inputs
25+
26+
Send any message (e.g., `go`) to start the workflow.
27+
28+
## Graph
29+
30+
```text
31+
[ START ]
32+
|
33+
v
34+
[fetch_weather] <-- pauses for auth on first run
35+
|
36+
v
37+
[summarize]
38+
```
39+
40+
## How To
41+
42+
1. Define an `AuthConfig` with the auth scheme and credential type:
43+
44+
```python
45+
from google.adk.auth.auth_tool import AuthConfig
46+
from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes
47+
48+
auth_config = AuthConfig(
49+
auth_scheme=APIKey(**{'in': APIKeyIn.header, 'name': 'X-Api-Key'}),
50+
raw_auth_credential=AuthCredential(
51+
auth_type=AuthCredentialTypes.API_KEY,
52+
api_key='placeholder',
53+
),
54+
credential_key='weather_api_key',
55+
)
56+
```
57+
58+
2. Use the `@node` decorator with `auth_config` and `rerun_on_resume=True`:
59+
60+
```python
61+
@node(auth_config=auth_config, rerun_on_resume=True)
62+
def fetch_weather(ctx: Context):
63+
...
64+
```
65+
66+
3. Inside the function, retrieve the credential from `ctx`:
67+
68+
```python
69+
def fetch_weather(ctx: Context):
70+
cred = ctx.get_auth_response(auth_config)
71+
api_key = cred.api_key
72+
# Use api_key to call your API...
73+
```
74+
75+
## OAuth2
76+
77+
The same `auth_config` pattern works with OAuth2 and OpenID Connect. The key
78+
differences:
79+
80+
- **Auth scheme**: Use `OAuth2` (from `fastapi.openapi.models`) instead of
81+
`APIKey`. Configure the authorization and token URLs in the OAuth2 flows.
82+
- **Raw credential**: Set `auth_type=AuthCredentialTypes.OAUTH2` and provide
83+
`client_id`, `client_secret`, and `redirect_uri` in the `oauth2` field.
84+
- **Web UI flow**: The ADK web UI recognizes OAuth2 auth requests and opens
85+
an authorization popup automatically. The user authenticates with the
86+
provider, and the UI sends the full `AuthConfig` response back. No special
87+
handling is needed in the node.
88+
- **Token exchange**: The framework automatically exchanges the authorization
89+
code for an access token via `AuthHandler.exchange_auth_token()`.
90+
91+
```python
92+
from fastapi.openapi.models import OAuth2, OAuthFlowAuthorizationCode, OAuthFlows
93+
94+
auth_config = AuthConfig(
95+
auth_scheme=OAuth2(
96+
flows=OAuthFlows(
97+
authorizationCode=OAuthFlowAuthorizationCode(
98+
authorizationUrl='https://provider.com/authorize',
99+
tokenUrl='https://provider.com/token',
100+
scopes={'read': 'Read access'},
101+
)
102+
)
103+
),
104+
raw_auth_credential=AuthCredential(
105+
auth_type=AuthCredentialTypes.OAUTH2,
106+
oauth2=OAuth2Auth(
107+
client_id='YOUR_CLIENT_ID',
108+
client_secret='YOUR_CLIENT_SECRET',
109+
redirect_uri='http://localhost:8000/callback',
110+
),
111+
),
112+
credential_key='my_oauth_credential',
113+
)
114+
```

‎contributing/workflow_samples/auth_config/agent.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Auth Config sample: FunctionNode with API key authentication.
16+
17+
Demonstrates how to use `auth_config` on a FunctionNode to pause
18+
the workflow and request user credentials before running the node.
19+
20+
Flow:
21+
1. User sends any message to start the workflow.
22+
2. The `fetch_weather` node pauses and requests an API key.
23+
3. The user provides the API key through the auth UI.
24+
4. The node runs with the credential available in session state.
25+
5. The `summarize` node displays the result.
26+
"""
27+
28+
from fastapi.openapi.models import APIKey
29+
from fastapi.openapi.models import APIKeyIn
30+
from google.adk import Event
31+
from google.adk import Workflow
32+
from google.adk.agents.context import Context
33+
from google.adk.auth.auth_credential import AuthCredential
34+
from google.adk.auth.auth_credential import AuthCredentialTypes
35+
from google.adk.auth.auth_tool import AuthConfig
36+
from google.adk.workflow import node
37+
38+
# --- Auth configuration ---
39+
# Uses API key auth: the simplest credential type.
40+
# The user will be prompted to provide an API key via the auth UI.
41+
auth_config = AuthConfig(
42+
auth_scheme=APIKey(**{'in': APIKeyIn.header, 'name': 'X-Api-Key'}),
43+
raw_auth_credential=AuthCredential(
44+
auth_type=AuthCredentialTypes.API_KEY,
45+
api_key='placeholder',
46+
),
47+
credential_key='weather_api_key',
48+
)
49+
50+
51+
@node(auth_config=auth_config, rerun_on_resume=True)
52+
def fetch_weather(ctx: Context):
53+
"""Fetches weather data using the authenticated API key."""
54+
# After auth completes, the credential is available via ctx.
55+
cred = ctx.get_auth_response(auth_config)
56+
api_key = cred.api_key if cred else 'unknown'
57+
58+
# In a real agent, you would use the api_key to call an external API.
59+
# For this sample, we just echo it back (masked).
60+
masked = api_key[:4] + '****' if len(api_key) > 4 else '****'
61+
return {
62+
'city': 'San Francisco',
63+
'temperature': '18C',
64+
'condition': 'Sunny',
65+
'api_key_used': masked,
66+
}
67+
68+
69+
def summarize(node_input: dict):
70+
"""Displays the weather result."""
71+
yield Event(
72+
message=(
73+
f"Weather for {node_input['city']}:"
74+
f" {node_input['temperature']}, {node_input['condition']}."
75+
f" (Authenticated with key: {node_input['api_key_used']})"
76+
)
77+
)
78+
79+
80+
root_agent = Workflow(
81+
name='auth_config',
82+
edges=[('START', fetch_weather, summarize)],
83+
)

‎src/google/adk/workflow/_function_node.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,14 @@
2929
from pydantic import TypeAdapter
3030
from typing_extensions import override
3131

32+
from ..auth.auth_tool import AuthConfig
3233
from ..events.event import Event
3334
from ..events.request_input import RequestInput
3435
from ._base_node import BaseNode
3536
from ._retry_config import RetryConfig
37+
from .utils._workflow_hitl_utils import create_auth_request_event
38+
from .utils._workflow_hitl_utils import has_auth_credential
39+
from .utils._workflow_hitl_utils import process_auth_resume
3640

3741
logger = logging.getLogger('google_adk.' + __name__)
3842

@@ -96,6 +100,16 @@ class FunctionNode(BaseNode):
96100
- All other values are validated/coerced by Pydantic's ``TypeAdapter``.
97101
"""
98102

103+
auth_config: AuthConfig | None = None
104+
"""If set, the framework requests user authentication before running.
105+
106+
When the node runs for the first time and no credential is found in
107+
session state, it yields an ``adk_request_credential`` event and
108+
interrupts. On resume, the credential is stored and the node
109+
re-runs with the credential available via
110+
``AuthHandler(auth_config).get_auth_response(ctx.state)``.
111+
"""
112+
99113
# Private attributes (won't be serialized)
100114
_func: Callable[..., Any] = PrivateAttr()
101115
_sig: inspect.Signature = PrivateAttr()
@@ -109,6 +123,7 @@ def __init__(
109123
rerun_on_resume: bool = False,
110124
retry_config: RetryConfig | None = None,
111125
timeout: float | None = None,
126+
auth_config: AuthConfig | None = None,
112127
):
113128
"""Initializes FunctionNode.
114129
@@ -125,11 +140,20 @@ def __init__(
125140
retry_config: If provided, the node will be retried on failure based on
126141
this configuration.
127142
timeout: Maximum time in seconds for this node to complete.
143+
auth_config: If provided, the framework requests user authentication
144+
before running the node. Requires rerun_on_resume=True (the node
145+
must rerun after credentials are provided).
128146
"""
129147

130148
if not callable(func):
131149
raise TypeError('Function must be callable.')
132150

151+
if auth_config and not rerun_on_resume:
152+
raise ValueError(
153+
'FunctionNode with auth_config requires rerun_on_resume=True.'
154+
' The node must rerun after credentials are provided.'
155+
)
156+
133157
inferred_name = name or getattr(func, '__name__', None)
134158
if not inferred_name:
135159
raise ValueError(
@@ -143,6 +167,7 @@ def __init__(
143167
rerun_on_resume=rerun_on_resume,
144168
retry_config=retry_config,
145169
timeout=timeout,
170+
auth_config=auth_config,
146171
)
147172

148173
sig = inspect.signature(func)
@@ -298,6 +323,16 @@ async def _run_impl(
298323
ctx: Context,
299324
node_input: Any,
300325
) -> AsyncGenerator[Any, None]:
326+
# --- Auth gate ---
327+
if self.auth_config:
328+
interrupt_id = f'wf_auth:{ctx.node_path}'
329+
auth_response = ctx.resume_inputs.get(interrupt_id)
330+
if auth_response is not None:
331+
await process_auth_resume(auth_response, self.auth_config, ctx.state)
332+
elif not has_auth_credential(self.auth_config, ctx.state):
333+
yield create_auth_request_event(self.auth_config, interrupt_id)
334+
return
335+
301336
kwargs: dict[str, Any] = {}
302337
for param_name, param in self._sig.parameters.items():
303338
if param_name == 'ctx':

‎src/google/adk/workflow/_node.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from typing import Any
2020
from typing import Callable
2121
from typing import overload
22+
from typing import TYPE_CHECKING
2223
from typing import TypeVar
2324

2425
from pydantic import Field
@@ -32,6 +33,9 @@
3233
from ._retry_config import RetryConfig
3334
from .utils import _workflow_graph_utils as workflow_graph_utils
3435

36+
if TYPE_CHECKING:
37+
from ..auth.auth_tool import AuthConfig
38+
3539
T = TypeVar("T", bound=Callable[..., Any])
3640

3741

@@ -43,6 +47,7 @@ def node(
4347
retry_config: RetryConfig | None = None,
4448
timeout: float | None = None,
4549
parallel_worker: bool = False,
50+
auth_config: AuthConfig | None = None,
4651
) -> Callable[
4752
[T], function_node.FunctionNode | parallel_worker_lib._ParallelWorker
4853
]:
@@ -58,6 +63,7 @@ def node(
5863
retry_config: RetryConfig | None = None,
5964
timeout: float | None = None,
6065
parallel_worker: bool = False,
66+
auth_config: AuthConfig | None = None,
6167
) -> base_node.BaseNode:
6268
...
6369

@@ -70,6 +76,7 @@ def node(
7076
retry_config: RetryConfig | None = None,
7177
timeout: float | None = None,
7278
parallel_worker: bool = False,
79+
auth_config: AuthConfig | None = None,
7380
) -> Any:
7481
"""Decorator or function to wrap a NodeLike in a node or override its properties.
7582
@@ -96,6 +103,8 @@ async def my_func3(): ...
96103
wrapped node.
97104
timeout: If provided, overrides the timeout property of the wrapped node.
98105
parallel_worker: If True, wraps the node in a _ParallelWorker.
106+
auth_config: If provided, the framework requests user authentication
107+
before running the node. Requires rerun_on_resume=True.
99108
100109
Returns:
101110
If used as a decorator factory (@node() or @node(...)), returns a decorator.
@@ -114,6 +123,7 @@ def wrapper(
114123
else False,
115124
retry_config=retry_config,
116125
timeout=timeout,
126+
auth_config=auth_config,
117127
)
118128
if parallel_worker:
119129
return parallel_worker_lib._ParallelWorker(built_node)
@@ -129,6 +139,7 @@ def wrapper(
129139
rerun_on_resume=rerun_on_resume,
130140
retry_config=retry_config,
131141
timeout=timeout,
142+
auth_config=auth_config,
132143
)
133144
if parallel_worker:
134145
return parallel_worker_lib._ParallelWorker(built_node)

‎src/google/adk/workflow/_workflow.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
from .utils._node_path_utils import join_paths
6262
from .utils._retry_utils import _get_retry_delay
6363
from .utils._retry_utils import _should_retry_node
64+
from .utils._workflow_hitl_utils import REQUEST_CREDENTIAL_FUNCTION_CALL_NAME
6465
from .utils._workflow_hitl_utils import REQUEST_INPUT_FUNCTION_CALL_NAME
6566
from .utils._workflow_hitl_utils import unwrap_response
6667

@@ -417,9 +418,9 @@ def _process_resumptions(
417418
]
418419

419420
response_data = response_part.function_response.response
420-
if (
421-
response_part.function_response.name
422-
== REQUEST_INPUT_FUNCTION_CALL_NAME
421+
if response_part.function_response.name in (
422+
REQUEST_INPUT_FUNCTION_CALL_NAME,
423+
REQUEST_CREDENTIAL_FUNCTION_CALL_NAME,
423424
):
424425
response_data = unwrap_response(response_data)
425426

‎src/google/adk/workflow/utils/_workflow_graph_utils.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ def build_node(
4545
rerun_on_resume: bool | None = None,
4646
retry_config: RetryConfig | None = None,
4747
timeout: float | None = None,
48+
auth_config: Any = None,
4849
) -> BaseNode:
4950
"""Converts a NodeLike to a BaseNode, wrapping async funcs in FunctionNode.
5051
@@ -56,6 +57,7 @@ def build_node(
5657
retry_config: If provided, overrides the retry_config property of the
5758
wrapped node.
5859
timeout: If provided, overrides the timeout property of the wrapped node.
60+
auth_config: If provided, passed to FunctionNode for authentication.
5961
6062
Returns:
6163
A BaseNode instance.
@@ -132,6 +134,7 @@ def build_node(
132134
rerun_on_resume=rerun_on_resume or False,
133135
retry_config=retry_config,
134136
timeout=timeout,
137+
auth_config=auth_config,
135138
)
136139
else:
137140
raise ValueError(

0 commit comments

Comments
 (0)