feat: add session-scoped client tool waits

This commit is contained in:
Xin Wang
2026-07-31 23:53:47 +08:00
parent f155f98e6e
commit ad5ff061bb
10 changed files with 339 additions and 26 deletions

View File

@@ -26,6 +26,7 @@ McpTransport = Literal["streamable_http", "sse"]
ToolParameterType = Literal["string", "number", "integer", "boolean", "object", "array"]
ToolParameterLocation = Literal["path", "query", "body", "header"]
ToolExecutionMode = Literal["immediate", "async"]
ClientToolResponseWaitMode = Literal["timeout", "session"]
DynamicVariableType = Literal["string", "number", "boolean"]
# 外部应用类型:其 config.apiKey 是该助手私有密钥,读时打码 / 写时哨兵
@@ -228,6 +229,7 @@ class ClientToolConfig(CamelModel):
allow_interruptions: bool = True
execution_mode: ToolExecutionMode = "async"
wait_for_response: bool = True
response_wait_mode: ClientToolResponseWaitMode = "timeout"
parameters: list[ToolParameter] = Field(default_factory=list)
timeout_seconds: int = Field(default=3, ge=1, le=30)
dynamic_variable_assignments: dict[str, str] = Field(default_factory=dict)

View File

@@ -533,7 +533,7 @@ class WorkflowBrain(BaseBrain):
cancel_on_interruption=policy.cancel_on_interruption,
timeout_secs=(
float(((tool.definition or {}).get("config") or {}).get("timeout_seconds") or 3)
if tool.type == "client"
if tool.type == "client" and policy.response_wait_mode == "timeout"
else None
),
)

View File

@@ -3,13 +3,15 @@
from __future__ import annotations
import asyncio
from typing import Any, Protocol
from typing import Any, Literal, Protocol
from loguru import logger
from pipecat.frames.frames import (
CancelFrame,
EndFrame,
InputTransportMessageFrame,
OutputTransportMessageUrgentFrame,
StopFrame,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
@@ -18,6 +20,9 @@ class ClientToolError(RuntimeError):
"""Raised when a client tool cannot be delivered or completed."""
ClientToolResponseWaitMode = Literal["timeout", "session"]
class ClientToolPort(Protocol):
async def call(
self,
@@ -26,6 +31,7 @@ class ClientToolPort(Protocol):
*,
timeout_seconds: float,
wait_for_response: bool = True,
response_wait_mode: ClientToolResponseWaitMode = "timeout",
) -> dict[str, Any]: ...
@@ -35,6 +41,7 @@ class ClientToolBroker(FrameProcessor):
def __init__(self) -> None:
super().__init__()
self._pending: dict[str, asyncio.Future[dict[str, Any]]] = {}
self._closed_message: str | None = None
async def call(
self,
@@ -43,9 +50,13 @@ class ClientToolBroker(FrameProcessor):
*,
timeout_seconds: float,
wait_for_response: bool = True,
response_wait_mode: ClientToolResponseWaitMode = "timeout",
) -> dict[str, Any]:
from uuid import uuid4
if self._closed_message is not None:
raise ClientToolError(self._closed_message)
tool_call_id = f"client_{uuid4().hex}"
message = {
"type": "client-tool-call",
@@ -75,6 +86,8 @@ class ClientToolBroker(FrameProcessor):
await self.push_frame(
OutputTransportMessageUrgentFrame(message=message)
)
if response_wait_mode == "session":
return await future
return await asyncio.wait_for(future, timeout=timeout_seconds)
except TimeoutError as exc:
raise ClientToolError(f"客户端工具调用超时: {function_name}") from exc
@@ -88,8 +101,17 @@ class ClientToolBroker(FrameProcessor):
async def process_frame(self, frame, direction: FrameDirection):
await super().process_frame(frame, direction)
if isinstance(frame, EndFrame):
self._fail_pending("会话已结束")
if isinstance(frame, (EndFrame, CancelFrame)):
self._close(
"会话已结束" if isinstance(frame, EndFrame) else "会话已取消"
)
await self.push_frame(frame, direction)
return
if isinstance(frame, StopFrame):
# StopFrame keeps processors reusable, so release only the calls
# owned by the stopped run without permanently closing the broker.
self._fail_pending("管线已停止")
await self.push_frame(frame, direction)
return
@@ -131,6 +153,12 @@ class ClientToolBroker(FrameProcessor):
future.set_exception(ClientToolError(message))
self._pending.clear()
def _close(self, message: str) -> None:
"""Close once and reject both current and future client-tool calls."""
if self._closed_message is None:
self._closed_message = message
self._fail_pending(self._closed_message)
async def cleanup(self) -> None:
self._fail_pending("客户端工具通道已关闭")
self._close("客户端工具通道已关闭")
await super().cleanup()

View File

@@ -213,6 +213,7 @@ class ToolExecutor:
arguments,
timeout_seconds=float(config.get("timeout_seconds") or 3),
wait_for_response=policy.wait_for_response,
response_wait_mode=policy.response_wait_mode,
)
except ClientToolError as exc:
raise ToolExecutionError(str(exc)) from exc

View File

@@ -19,6 +19,7 @@ class ToolRuntimePolicy:
allow_interruptions: bool
execution_mode: str
wait_for_response: bool
response_wait_mode: str
@property
def cancel_on_interruption(self) -> bool:
@@ -38,8 +39,12 @@ def policy_for_tool(tool: RuntimeTool) -> ToolRuntimePolicy:
mode = str(config.get("execution_mode") or default_mode)
if mode not in {"immediate", "async"}:
mode = default_mode
response_wait_mode = str(config.get("response_wait_mode") or "timeout")
if response_wait_mode not in {"timeout", "session"}:
response_wait_mode = "timeout"
return ToolRuntimePolicy(
allow_interruptions=bool(config.get("allow_interruptions", True)),
execution_mode=mode,
wait_for_response=bool(config.get("wait_for_response", True)),
response_wait_mode=response_wait_mode,
)

View File

@@ -518,6 +518,42 @@ class PromptBrainTests(unittest.IsolatedAsyncioTestCase):
class WorkflowBrainTests(unittest.IsolatedAsyncioTestCase):
def test_client_tool_session_wait_disables_flow_timeout(self):
brain = WorkflowBrain(
{
"specVersion": 3,
"settings": {},
"nodes": [{"id": "start", "type": "start", "data": {}}],
"edges": [],
}
)
timeout_tool = RuntimeTool(
id="client_timeout",
name="限时等待",
function_name="wait_with_timeout",
type="client",
definition={
"type": "client",
"config": {"timeout_seconds": 7},
},
)
session_tool = RuntimeTool(
id="client_session",
name="会话内等待",
function_name="wait_for_session",
type="client",
definition={
"type": "client",
"config": {
"timeout_seconds": 7,
"response_wait_mode": "session",
},
},
)
self.assertEqual(brain._flow_tool(timeout_tool, "start").timeout_secs, 7.0)
self.assertIsNone(brain._flow_tool(session_tool, "start").timeout_secs)
async def test_session_update_refreshes_current_agent_without_routing(self):
cfg = prepare_dynamic_config(
AssistantConfig(

View File

@@ -3,8 +3,11 @@ import unittest
from models import RuntimeTool
from pipecat.frames.frames import (
CancelFrame,
EndFrame,
InputTransportMessageFrame,
OutputTransportMessageUrgentFrame,
StopFrame,
)
from pipecat.processors.frame_processor import FrameDirection
from services.client_tools import ClientToolBroker, ClientToolError
@@ -51,14 +54,32 @@ class FakeClientTools:
*,
timeout_seconds,
wait_for_response=True,
response_wait_mode="timeout",
):
self.calls.append(
(function_name, arguments, timeout_seconds, wait_for_response)
(
function_name,
arguments,
timeout_seconds,
wait_for_response,
response_wait_mode,
)
)
return self.result
class ClientToolExecutorTests(unittest.IsolatedAsyncioTestCase):
def test_response_wait_mode_defaults_and_uses_camel_case(self):
legacy = ClientToolConfig()
session = ClientToolConfig(responseWaitMode="session")
self.assertEqual(legacy.response_wait_mode, "timeout")
self.assertEqual(session.response_wait_mode, "session")
self.assertEqual(
session.model_dump(by_alias=True)["responseWaitMode"],
"session",
)
def test_fire_and_forget_rejects_result_assignments(self):
with self.assertRaisesRegex(ValueError, "不能配置结果变量赋值"):
ClientToolConfig(
@@ -80,9 +101,27 @@ class ClientToolExecutorTests(unittest.IsolatedAsyncioTestCase):
self.assertIs(store.values["photo_button_visible"], True)
self.assertEqual(
port.calls,
[("set_photo_button_visible", {"visible": True}, 3.0, True)],
[
(
"set_photo_button_visible",
{"visible": True},
3.0,
True,
"timeout",
)
],
)
async def test_session_wait_mode_is_forwarded(self):
tool = client_tool()
tool.definition["config"]["response_wait_mode"] = "session"
port = FakeClientTools({"status": "ok", "data": {"visible": True}})
executor = ToolExecutor(DynamicVariableStore({}), client_tools=port)
await executor.execute(tool, {"visible": True})
self.assertEqual(port.calls[0][-1], "session")
async def test_failure_does_not_update_variable(self):
store = DynamicVariableStore({"photo_button_visible": False})
executor = ToolExecutor(
@@ -152,6 +191,160 @@ class ClientToolBrokerTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(broker._pending, {})
self.assertFalse(outbound[0][0].message["wait_for_response"])
async def test_session_wait_is_released_by_end_frame(self):
broker = ClientToolBroker()
async def push_frame(frame, direction=FrameDirection.DOWNSTREAM):
pass
broker.push_frame = push_frame
call = asyncio.create_task(
broker.call(
"show_message",
{},
timeout_seconds=0.001,
response_wait_mode="session",
)
)
await asyncio.sleep(0.01)
self.assertFalse(call.done())
await broker.process_frame(EndFrame(), FrameDirection.DOWNSTREAM)
with self.assertRaisesRegex(ClientToolError, "会话已结束"):
await call
self.assertEqual(broker._pending, {})
async def test_session_wait_is_released_by_cleanup(self):
broker = ClientToolBroker()
async def push_frame(frame, direction=FrameDirection.DOWNSTREAM):
pass
broker.push_frame = push_frame
call = asyncio.create_task(
broker.call(
"show_message",
{},
timeout_seconds=0.001,
response_wait_mode="session",
)
)
await asyncio.sleep(0.01)
self.assertFalse(call.done())
await broker.cleanup()
with self.assertRaisesRegex(ClientToolError, "通道已关闭"):
await call
self.assertEqual(broker._pending, {})
async def test_call_after_end_frame_fails_without_registering(self):
broker = ClientToolBroker()
outbound = []
async def push_frame(frame, direction=FrameDirection.DOWNSTREAM):
outbound.append(frame)
broker.push_frame = push_frame
await broker.process_frame(EndFrame(), FrameDirection.DOWNSTREAM)
with self.assertRaisesRegex(ClientToolError, "会话已结束"):
await broker.call(
"show_message",
{},
timeout_seconds=1,
response_wait_mode="session",
)
self.assertEqual(broker._pending, {})
self.assertEqual(len(outbound), 1)
async def test_call_after_cleanup_fails_without_registering(self):
broker = ClientToolBroker()
await broker.cleanup()
with self.assertRaisesRegex(ClientToolError, "通道已关闭"):
await broker.call(
"show_message",
{},
timeout_seconds=1,
response_wait_mode="session",
)
self.assertEqual(broker._pending, {})
async def test_cancel_frame_closes_broker(self):
broker = ClientToolBroker()
async def push_frame(frame, direction=FrameDirection.DOWNSTREAM):
pass
broker.push_frame = push_frame
await broker.process_frame(CancelFrame(), FrameDirection.DOWNSTREAM)
with self.assertRaisesRegex(ClientToolError, "会话已取消"):
await broker.call(
"show_message",
{},
timeout_seconds=1,
response_wait_mode="session",
)
async def test_stop_frame_releases_pending_but_keeps_broker_reusable(self):
broker = ClientToolBroker()
outbound = []
async def push_frame(frame, direction=FrameDirection.DOWNSTREAM):
outbound.append(frame)
broker.push_frame = push_frame
stopped_call = asyncio.create_task(
broker.call(
"show_message",
{},
timeout_seconds=1,
response_wait_mode="session",
)
)
await asyncio.sleep(0)
await broker.process_frame(StopFrame(), FrameDirection.DOWNSTREAM)
with self.assertRaisesRegex(ClientToolError, "管线已停止"):
await stopped_call
self.assertEqual(broker._pending, {})
next_call = asyncio.create_task(
broker.call(
"show_message",
{},
timeout_seconds=1,
response_wait_mode="session",
)
)
await asyncio.sleep(0)
message = next(
frame.message
for frame in reversed(outbound)
if isinstance(frame, OutputTransportMessageUrgentFrame)
)
await broker.process_frame(
InputTransportMessageFrame(
message={
"type": "client-tool-result",
"tool_call_id": message["tool_call_id"],
"status": "ok",
"data": {"action": "confirmed"},
}
),
FrameDirection.DOWNSTREAM,
)
self.assertEqual(
await next_call,
{"status": "ok", "data": {"action": "confirmed"}},
)
if __name__ == "__main__":
unittest.main()

View File

@@ -32,6 +32,10 @@ class ToolPolicyTests(unittest.TestCase):
policy_for_tool(runtime_tool("client")).execution_mode,
"async",
)
self.assertEqual(
policy_for_tool(runtime_tool("client")).response_wait_mode,
"timeout",
)
def test_explicit_policy_is_normalized(self):
policy = policy_for_tool(
@@ -41,12 +45,14 @@ class ToolPolicyTests(unittest.TestCase):
"allow_interruptions": False,
"execution_mode": "immediate",
"wait_for_response": False,
"response_wait_mode": "session",
},
)
)
self.assertFalse(policy.allow_interruptions)
self.assertTrue(policy.cancel_on_interruption)
self.assertFalse(policy.wait_for_response)
self.assertEqual(policy.response_wait_mode, "session")
class ToolInterruptionStrategyTests(unittest.IsolatedAsyncioTestCase):