From ad5ff061bbe11c6d2e9fb2bc032ffdbfa1173d4d Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Fri, 31 Jul 2026 23:53:47 +0800 Subject: [PATCH] feat: add session-scoped client tool waits --- backend/schemas.py | 2 + backend/services/brains/workflow_brain.py | 2 +- backend/services/client_tools.py | 36 +++- backend/services/tool_executor.py | 1 + backend/services/tool_policy.py | 5 + backend/tests/test_brains.py | 36 ++++ backend/tests/test_client_tools.py | 197 +++++++++++++++++- backend/tests/test_tool_policy.py | 6 + .../components/pages/ComponentsToolsPage.tsx | 78 +++++-- frontend/src/lib/api.ts | 2 + 10 files changed, 339 insertions(+), 26 deletions(-) diff --git a/backend/schemas.py b/backend/schemas.py index d9aaddb..109bed1 100644 --- a/backend/schemas.py +++ b/backend/schemas.py @@ -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) diff --git a/backend/services/brains/workflow_brain.py b/backend/services/brains/workflow_brain.py index 5d86086..25401c1 100644 --- a/backend/services/brains/workflow_brain.py +++ b/backend/services/brains/workflow_brain.py @@ -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 ), ) diff --git a/backend/services/client_tools.py b/backend/services/client_tools.py index f2f72ea..ca922b0 100644 --- a/backend/services/client_tools.py +++ b/backend/services/client_tools.py @@ -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() diff --git a/backend/services/tool_executor.py b/backend/services/tool_executor.py index c0858dc..f3c6725 100644 --- a/backend/services/tool_executor.py +++ b/backend/services/tool_executor.py @@ -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 diff --git a/backend/services/tool_policy.py b/backend/services/tool_policy.py index f7171d3..358cf99 100644 --- a/backend/services/tool_policy.py +++ b/backend/services/tool_policy.py @@ -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, ) diff --git a/backend/tests/test_brains.py b/backend/tests/test_brains.py index cd8975c..ef5d1f4 100644 --- a/backend/tests/test_brains.py +++ b/backend/tests/test_brains.py @@ -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( diff --git a/backend/tests/test_client_tools.py b/backend/tests/test_client_tools.py index 728d57c..622b137 100644 --- a/backend/tests/test_client_tools.py +++ b/backend/tests/test_client_tools.py @@ -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() diff --git a/backend/tests/test_tool_policy.py b/backend/tests/test_tool_policy.py index 7d0acbf..ae07d52 100644 --- a/backend/tests/test_tool_policy.py +++ b/backend/tests/test_tool_policy.py @@ -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): diff --git a/frontend/src/components/pages/ComponentsToolsPage.tsx b/frontend/src/components/pages/ComponentsToolsPage.tsx index d5c40fa..8d8bc91 100644 --- a/frontend/src/components/pages/ComponentsToolsPage.tsx +++ b/frontend/src/components/pages/ComponentsToolsPage.tsx @@ -54,6 +54,7 @@ import { Textarea } from "@/components/ui/textarea"; import { mcpServersApi, toolsApi, + type ClientToolResponseWaitMode, type HttpToolDefinition, type McpServer, type Tool, @@ -94,6 +95,7 @@ type ToolForm = { allowInterruptions: boolean; executionMode: ToolExecutionMode; waitForResponse: boolean; + responseWaitMode: ClientToolResponseWaitMode; headers: string; secretHeaders: string; parameters: string; @@ -122,6 +124,7 @@ function blankForm(): ToolForm { allowInterruptions: true, executionMode: "immediate", waitForResponse: true, + responseWaitMode: "timeout", headers: EMPTY_OBJECT, secretHeaders: EMPTY_OBJECT, parameters: EMPTY_ARRAY, @@ -156,6 +159,8 @@ function formFromTool(tool: Tool): ToolForm { tool.definition.config.allowInterruptions ?? true; base.executionMode = tool.definition.config.executionMode ?? "async"; base.waitForResponse = tool.definition.config.waitForResponse ?? true; + base.responseWaitMode = + tool.definition.config.responseWaitMode ?? "timeout"; base.timeoutSeconds = String(tool.definition.config.timeoutSeconds); base.parameters = pretty(tool.definition.config.parameters, EMPTY_ARRAY); base.dynamicVariableAssignments = pretty( @@ -249,14 +254,19 @@ function payloadFromForm(form: ToolForm): ToolUpsert { }; } if (form.type === "client") { - const timeoutSeconds = Number(form.timeoutSeconds); + const parsedTimeoutSeconds = Number(form.timeoutSeconds); + const hasValidTimeout = + Number.isInteger(parsedTimeoutSeconds) && + parsedTimeoutSeconds >= 1 && + parsedTimeoutSeconds <= 30; if ( - !Number.isInteger(timeoutSeconds) || - timeoutSeconds < 1 || - timeoutSeconds > 30 + form.waitForResponse && + form.responseWaitMode === "timeout" && + !hasValidTimeout ) { throw new Error("Client Tool 超时时间必须是 1 到 30 秒之间的整数"); } + const timeoutSeconds = hasValidTimeout ? parsedTimeoutSeconds : 15; const dynamicVariableAssignments = form.updateDynamicVariables ? parseObject(form.dynamicVariableAssignments, "变量赋值") : {}; @@ -272,6 +282,7 @@ function payloadFromForm(form: ToolForm): ToolUpsert { allowInterruptions: form.allowInterruptions, executionMode: form.executionMode, waitForResponse: form.waitForResponse, + responseWaitMode: form.responseWaitMode, parameters: parseParameters(form.parameters), timeoutSeconds, dynamicVariableAssignments: @@ -759,7 +770,11 @@ export function ComponentsToolsPage() { ...current, type, ...(type === "client" - ? { timeoutSeconds: "3", executionMode: "async" as const } + ? { + timeoutSeconds: "3", + executionMode: "async" as const, + responseWaitMode: "timeout" as const, + } : type === "http" ? { timeoutSeconds: "15", executionMode: "immediate" as const } : {}), @@ -926,20 +941,45 @@ function ClientToolFields({ /> {form.waitForResponse && ( - - - setForm((current) => ({ - ...current, - timeoutSeconds: event.target.value, - })) - } - /> - + <> + + + + {form.responseWaitMode === "timeout" + ? "超过设定时间后，将本次调用按失败处理。" + : "持续等待客户端返回，或直到本次会话结束。"} + + + {form.responseWaitMode === "timeout" && ( + + + setForm((current) => ({ + ...current, + timeoutSeconds: event.target.value, + })) + } + /> + + )} + )} ;