fix(utils): adapt callback keyword arguments to positional parameters

Callbacks typed as Callable[[CallbackContext, ...], ...] may define
positional parameters with arbitrary names (such as lambda ctx, resp:).
When invoked with keyword arguments, parameter binding fails with
unexpected keyword argument. _invoke_callback now inspects the
callback signature and adapts keyword arguments to positional arguments
when keyword binding fails, while preserving exceptions raised inside
the callback body.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 974083312

Author: DeanChensj

src/google/adk/utils/_callback_pipeline.py

@@ -18,6 +18,7 @@
 from collections.abc import Callable
 from collections.abc import Sequence
 import inspect
+from typing import Any
 from typing import TypeVar
 
 from typing_extensions import ParamSpec
@@ -27,6 +28,56 @@
 _TResult = TypeVar('_TResult')
 
 
+_CANONICAL_PARAM_ORDER: tuple[str, ...] = (
+    'tool',
+    'args',
+    'tool_context',
+    'response',
+    'tool_response',
+    'callback_context',
+    'llm_request',
+    'llm_response',
+    'error',
+)
+
+
+def _canonical_order_key(name: str) -> int:
+  try:
+    return _CANONICAL_PARAM_ORDER.index(name)
+  except ValueError:
+    return len(_CANONICAL_PARAM_ORDER)
+
+
+def _invoke_callback(
+    _callback: Callable[..., Any], /, *args: Any, **kwargs: Any
+) -> Any:
+  """Invokes callback, adapting keyword arguments to positional if needed."""
+  if kwargs:
+    use_pos = False
+    pos_args: tuple[Any, ...] = ()
+    try:
+      sig = inspect.signature(_callback)
+      try:
+        sig.bind(*args, **kwargs)
+      except TypeError:
+        sorted_kwarg_values = [
+            val
+            for _, val in sorted(
+                kwargs.items(), key=lambda item: _canonical_order_key(item[0])
+            )
+        ]
+        pos_args = args + tuple(sorted_kwarg_values)
+        sig.bind(*pos_args)
+        use_pos = True
+    except (ValueError, TypeError):
+      pass
+
+    if use_pos:
+      return _callback(*pos_args)
+
+  return _callback(*args, **kwargs)
+
+
 async def _run_callbacks(
     callbacks: Sequence[
         Callable[
@@ -41,7 +92,7 @@ async def _run_callbacks(
   """Runs callbacks in order while preserving their stop semantics."""
   result: _TResult | None = None
   for callback in callbacks:
-    callback_result = callback(*args, **kwargs)
+    callback_result = _invoke_callback(callback, *args, **kwargs)
     if inspect.isawaitable(callback_result):
       result = await callback_result
     else:

tests/unittests/flows/llm_flows/test_model_callbacks.py

@@ -159,6 +159,22 @@ async def test_after_model_callback_noop():
   ) == [('root_agent', 'model_response')]
 
 
+def test_after_model_callback_lambda_with_arbitrary_param_names():
+  """Test that after_model_callback works with lambda having non-matching param names."""
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_model_callback=lambda ctx, resp: None,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', 'model_response'),
+  ]
+
+
 @pytest.mark.asyncio
 async def test_on_model_callback_model_error_noop():
   """Test that the on_model_error_callback is a no-op when the model returns an error."""

tests/unittests/flows/llm_flows/test_tool_callbacks.py

@@ -476,3 +476,24 @@ async def async_on_tool_error_callback(
       ),
       ('root_agent', 'response1'),
   ]
+
+
+def test_before_tool_callback_lambda_with_arbitrary_param_names():
+  """Test that before_tool_callback works with lambda having non-matching param names."""
+  captured = []
+  responses = [
+      types.Part.from_function_call(name='simple_function', args={}),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_tool_callback=lambda t, a, tc: captured.append((t.name, a)),
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  runner.run('test')
+  assert len(captured) == 1
+  assert captured[0] == ('simple_function', {})

tests/unittests/utils/test_callback_pipeline.py

@@ -255,3 +255,197 @@ def empty_callback() -> dict[str, str]:
 
   assert result == {}
   assert calls == ['none', 'empty']
+
+
+@pytest.mark.asyncio
+async def test_callback_adapts_keyword_args_to_positional_params():
+  """Callbacks with positional param names accept keyword arguments."""
+
+  def callback(ctx: str, resp: str) -> str:
+    return f'{ctx}:{resp}'
+
+  result = await _run_callbacks(
+      [callback],
+      _stop_on_truthy,
+      callback_context='context_val',
+      llm_response='response_val',
+  )
+
+  assert result == 'context_val:response_val'
+
+
+@pytest.mark.asyncio
+async def test_async_callback_adapts_keyword_args_to_positional_params():
+  """Async callbacks with positional param names accept keyword arguments."""
+
+  async def callback(a: str, b: str) -> str:
+    return f'{a}:{b}'
+
+  result = await _run_callbacks(
+      [callback],
+      _stop_on_truthy,
+      callback_context='arg_a',
+      llm_response='arg_b',
+  )
+
+  assert result == 'arg_a:arg_b'
+
+
+@pytest.mark.asyncio
+async def test_lambda_callback_adapts_keyword_args():
+  """Lambda callbacks with non-matching param names accept keyword arguments."""
+  callback = lambda c, r: f'lambda_{c}_{r}'
+
+  result = await _run_callbacks(
+      [callback],
+      _stop_on_truthy,
+      callback_context='c1',
+      llm_response='r1',
+  )
+
+  assert result == 'lambda_c1_r1'
+
+
+@pytest.mark.asyncio
+async def test_callback_with_body_type_error_is_not_masked():
+  """A TypeError raised inside the callback body is preserved."""
+
+  def callback(ctx: str, resp: str) -> None:
+    raise TypeError('error inside callback body')
+
+  with pytest.raises(TypeError, match='error inside callback body'):
+    await _run_callbacks(
+        [callback],
+        _stop_on_truthy,
+        callback_context='ctx',
+        llm_response='resp',
+    )
+
+
+@pytest.mark.asyncio
+async def test_callback_with_reordered_kwargs_maintains_canonical_order():
+  """Reordering kwargs at call site still binds positional parameters in canonical order."""
+  callback = lambda ctx, resp: f'{ctx}:{resp}'
+
+  result = await _run_callbacks(
+      [callback],
+      _stop_on_truthy,
+      llm_response='resp_val',
+      callback_context='ctx_val',
+  )
+
+  assert result == 'ctx_val:resp_val'
+
+
+@pytest.mark.asyncio
+async def test_callback_with_callback_kwarg_does_not_collide():
+  """A kwarg named 'callback' does not collide with _callback parameter."""
+  callback = lambda c, cb: f'{c}:{cb}'
+
+  result = await _run_callbacks(
+      [callback],
+      _stop_on_truthy,
+      callback_context='ctx_val',
+      callback='custom_callback',
+  )
+
+  assert result == 'ctx_val:custom_callback'
+
+
+@pytest.mark.asyncio
+async def test_callback_with_incompatible_signature_raises_type_error():
+  """When both keyword and positional binding fail, original TypeError is raised."""
+
+  def callback(a: str, b: str, c: str, d: str) -> None:
+    pass
+
+  with pytest.raises(TypeError):
+    await _run_callbacks(
+        [callback],
+        _stop_on_truthy,
+        callback_context='ctx',
+        llm_response='resp',
+    )
+
+
+@pytest.mark.asyncio
+async def test_before_tool_callback_with_reordered_kwargs_maintains_canonical_order():
+  """Tool callbacks receive (tool, args, tool_context) in canonical order even if kwargs are reordered."""
+  callback = lambda t, a, tc: f'{t}:{a}:{tc}'
+
+  result = await _run_callbacks(
+      [callback],
+      _stop_on_non_none,
+      tool_context='ctx_val',
+      args='args_val',
+      tool='tool_val',
+  )
+
+  assert result == 'tool_val:args_val:ctx_val'
+
+
+@pytest.mark.asyncio
+async def test_after_tool_callback_with_reordered_kwargs_maintains_canonical_order():
+  """After-tool callbacks receive (tool, args, tool_context, response) in canonical order."""
+  callback = lambda t, a, tc, r: f'{t}:{a}:{tc}:{r}'
+
+  result = await _run_callbacks(
+      [callback],
+      _stop_on_non_none,
+      response='resp_val',
+      tool_context='ctx_val',
+      args='args_val',
+      tool='tool_val',
+  )
+
+  assert result == 'tool_val:args_val:ctx_val:resp_val'
+
+
+@pytest.mark.asyncio
+async def test_after_tool_callback_with_tool_response_kwarg_maintains_canonical_order():
+  """After-tool callbacks with tool_response kwarg receive parameters in canonical order."""
+  callback = lambda t, a, tc, tr: f'{t}:{a}:{tc}:{tr}'
+
+  result = await _run_callbacks(
+      [callback],
+      _stop_on_non_none,
+      tool_response='tr_val',
+      tool_context='ctx_val',
+      args='args_val',
+      tool='tool_val',
+  )
+
+  assert result == 'tool_val:args_val:ctx_val:tr_val'
+
+
+@pytest.mark.asyncio
+async def test_on_tool_error_callback_with_reordered_kwargs_maintains_canonical_order():
+  """Tool error callbacks receive (tool, args, tool_context, error) in canonical order."""
+  callback = lambda t, a, tc, err: f'{t}:{a}:{tc}:{err}'
+
+  result = await _run_callbacks(
+      [callback],
+      _stop_on_non_none,
+      error='err_val',
+      tool_context='ctx_val',
+      args='args_val',
+      tool='tool_val',
+  )
+
+  assert result == 'tool_val:args_val:ctx_val:err_val'
+
+
+@pytest.mark.asyncio
+async def test_on_model_error_callback_with_reordered_kwargs_maintains_canonical_order():
+  """Model error callbacks receive (callback_context, llm_request, error) in canonical order."""
+  callback = lambda ctx, req, err: f'{ctx}:{req}:{err}'
+
+  result = await _run_callbacks(
+      [callback],
+      _stop_on_truthy,
+      error='err_val',
+      llm_request='req_val',
+      callback_context='ctx_val',
+  )
+
+  assert result == 'ctx_val:req_val:err_val'