mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-30 12:25:29 +08:00
276 lines
9.8 KiB
Python
276 lines
9.8 KiB
Python
import asyncio
|
|
import importlib
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from test_isolation import GATEWAYS, context, gateway_invocation, load
|
|
|
|
|
|
@pytest.mark.parametrize("name", GATEWAYS)
|
|
def test_gateway_actual_http_revocation_limits_and_close(name):
|
|
async def run(runner):
|
|
import httpx
|
|
|
|
runner.get_run_api = lambda ctx: SimpleNamespace()
|
|
reg = await runner._create_asset_gateway_registration(
|
|
context(),
|
|
dict(
|
|
asset_gateway_host="127.0.0.1",
|
|
asset_gateway_port=0,
|
|
asset_gateway_request_timeout=0.1,
|
|
asset_gateway_token_ttl=60,
|
|
),
|
|
)
|
|
|
|
class Tools:
|
|
async def call_mcp_tool(self, name, args):
|
|
return {"content": [{"type": "text", "text": "scope-A"}]}
|
|
|
|
reg.tools = Tools()
|
|
payload = {
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/call",
|
|
"params": {"name": "fixture", "arguments": {"run_token": reg.token}},
|
|
}
|
|
async with httpx.AsyncClient(trust_env=False) as client:
|
|
result = (await client.post(reg.http_mcp_endpoint, json=payload)).json()
|
|
assert result["result"]["content"][0]["text"] == "scope-A"
|
|
payload["params"]["arguments"]["run_token"] = "foreign"
|
|
assert (await client.post(reg.http_mcp_endpoint, json=payload)).json()["result"]["isError"]
|
|
assert (await client.post(reg.http_mcp_endpoint, content=b"x" * (1024 * 1024 + 1))).status_code == 413
|
|
# Slow HTTP peer must expire rather than retaining a handler/task.
|
|
port = reg.gateway.server.sockets[0].getsockname()[1]
|
|
reader, writer = await asyncio.open_connection("127.0.0.1", port)
|
|
writer.write(b"POST /mcp HTTP/1.1\r\nContent-Length: 100\r\n\r\n")
|
|
await writer.drain()
|
|
assert await asyncio.wait_for(reader.read(), 1) == b""
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
await reg.stop()
|
|
assert reg.gateway._registration_for_token(reg.token) is None
|
|
assert not reg.gateway.tasks
|
|
with pytest.raises(httpx.ConnectError):
|
|
await client.get(reg.endpoint)
|
|
|
|
with load(name) as (_, runner):
|
|
with gateway_invocation(runner):
|
|
asyncio.run(run(runner))
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["deerflow", "weknora", "dify", "langflow"])
|
|
def test_total_deadline_includes_trickle_stream(name, monkeypatch):
|
|
import httpx
|
|
|
|
async def run():
|
|
with load(name):
|
|
mod = importlib.import_module("pkg." + name + "_client")
|
|
cls = getattr(
|
|
mod,
|
|
{
|
|
"deerflow": "AsyncDeerFlowClient",
|
|
"weknora": "AsyncWeKnoraClient",
|
|
"dify": "AsyncDifyClient",
|
|
"langflow": "AsyncLangflowClient",
|
|
}[name],
|
|
)
|
|
|
|
class Trickle(httpx.AsyncByteStream):
|
|
closed = False
|
|
|
|
async def __aiter__(self):
|
|
while True:
|
|
await asyncio.sleep(0.015)
|
|
yield b": keepalive\n\n"
|
|
|
|
async def aclose(self):
|
|
self.closed = True
|
|
|
|
body = Trickle()
|
|
original = httpx.AsyncClient
|
|
client = original(
|
|
base_url="https://fixture.invalid",
|
|
transport=httpx.MockTransport(lambda request: httpx.Response(200, stream=body)),
|
|
)
|
|
monkeypatch.setattr(httpx, "AsyncClient", lambda *args, **kwargs: client)
|
|
if name == "deerflow":
|
|
c = cls(api_base="https://fixture.invalid")
|
|
gen = c.stream_run("thread", {}, timeout=0.08)
|
|
elif name == "weknora":
|
|
c = cls(api_key="fixture", base_url="https://fixture.invalid")
|
|
gen = c.knowledge_chat(session_id="session", query="hi", user="user", timeout=0.08)
|
|
elif name == "dify":
|
|
c = cls(api_key="fixture", base_url="https://fixture.invalid", timeout=0.08)
|
|
gen = c.chat_messages(inputs={}, query="hi", user="user")
|
|
else:
|
|
c = cls(api_key="fixture", base_url="https://fixture.invalid", timeout=0.08)
|
|
gen = c.run_flow(flow_id="flow", input_value="hi")
|
|
started = asyncio.get_running_loop().time()
|
|
with pytest.raises(TimeoutError):
|
|
await asyncio.wait_for(anext(gen), 0.35)
|
|
elapsed = asyncio.get_running_loop().time() - started
|
|
await gen.aclose()
|
|
assert elapsed < 0.25, "per-read timeout permits indefinite trickle"
|
|
assert body.closed
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_dify_continuation_rejects_cross_conversation():
|
|
with load("dify") as (m, runner):
|
|
|
|
async def run():
|
|
data = {}
|
|
|
|
async def put(k, v):
|
|
data[k] = v
|
|
|
|
async def get(k):
|
|
return data[k]
|
|
|
|
runner.get_run_api = lambda ctx: SimpleNamespace(set_plugin_storage=put, get_plugin_storage=get)
|
|
a = context()
|
|
b = context()
|
|
b.conversation.conversation_id = "other"
|
|
await runner._store_interaction_continuation(a, {"interaction_id": "i", "version": 1})
|
|
with pytest.raises(m.DifyAPIError):
|
|
await runner._load_interaction_continuation(b, "i")
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_dify_failed_resume_is_not_replayed():
|
|
with load("dify") as (m, runner):
|
|
|
|
async def run():
|
|
data = {}
|
|
calls = []
|
|
|
|
async def put(k, v):
|
|
data[k] = v
|
|
|
|
async def get(k):
|
|
return data[k]
|
|
|
|
runner.get_run_api = lambda ctx: SimpleNamespace(set_plugin_storage=put, get_plugin_storage=get)
|
|
|
|
class Client:
|
|
async def workflow_submit(self, **kwargs):
|
|
calls.append(kwargs)
|
|
raise ConnectionError("ambiguous remote failure")
|
|
yield
|
|
|
|
ctx = context()
|
|
await runner._store_interaction_continuation(
|
|
ctx,
|
|
{
|
|
"interaction_id": "i",
|
|
"version": 1,
|
|
"phase": "action",
|
|
"action_map": {"yes": "approve"},
|
|
"form_token": "fixture",
|
|
"workflow_run_id": "w",
|
|
"user": "u",
|
|
},
|
|
)
|
|
submission = SimpleNamespace(interaction_id="i", values={}, action_id="yes")
|
|
for _ in range(2):
|
|
with pytest.raises((ConnectionError, m.DifyAPIError)):
|
|
async for _ in runner._resume_workflow(ctx, Client(), submission, False):
|
|
pass
|
|
assert len(calls) == 1
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.parametrize("name", GATEWAYS)
|
|
def test_same_installation_gateway_overlap_keeps_other_token_alive(name):
|
|
async def run(runner):
|
|
runner.get_run_api = lambda ctx: SimpleNamespace()
|
|
cfg = dict(
|
|
asset_gateway_host="127.0.0.1",
|
|
asset_gateway_port=0,
|
|
asset_gateway_request_timeout=1,
|
|
asset_gateway_token_ttl=60,
|
|
)
|
|
a = await runner._create_asset_gateway_registration(context(), cfg)
|
|
# A fixed port is the documented externally configured MCP address.
|
|
cfg["asset_gateway_port"] = a.gateway.server.sockets[0].getsockname()[1]
|
|
await a.stop()
|
|
a = await runner._create_asset_gateway_registration(context(), cfg)
|
|
b = None
|
|
try:
|
|
b = await runner._create_asset_gateway_registration(context(), cfg)
|
|
await a.stop()
|
|
assert b.gateway._registration_for_token(b.token) is b
|
|
finally:
|
|
await a.stop()
|
|
if b:
|
|
await b.stop()
|
|
|
|
with load(name) as (_, runner):
|
|
with gateway_invocation(runner):
|
|
asyncio.run(run(runner))
|
|
|
|
|
|
@pytest.mark.parametrize("name", GATEWAYS)
|
|
def test_runner_cancellation_releases_gateway(name, monkeypatch):
|
|
async def run(module, runner):
|
|
entered = asyncio.Event()
|
|
|
|
class Client:
|
|
def __init__(self, **kwargs):
|
|
pass
|
|
|
|
async def close(self):
|
|
pass
|
|
|
|
async def hang(self, **kwargs):
|
|
entered.set()
|
|
await asyncio.Event().wait()
|
|
yield {}
|
|
|
|
chat_messages = iter_agent = iter_workflow = run_flow = call_webhook = hang
|
|
|
|
monkeypatch.setattr(
|
|
module,
|
|
{
|
|
"coze": "AsyncCozeClient",
|
|
"dashscope": "DashScopeClient",
|
|
"dify": "AsyncDifyClient",
|
|
"langflow": "AsyncLangflowClient",
|
|
"n8n": "AsyncN8nClient",
|
|
}[name],
|
|
Client,
|
|
)
|
|
runner.get_run_api = lambda ctx: SimpleNamespace()
|
|
ctx = context()
|
|
ctx.config.update(
|
|
{
|
|
"api-key": "fixture",
|
|
"app-id": "app",
|
|
"bot-id": "bot",
|
|
"flow-id": "flow",
|
|
"webhook-url": "https://fixture.invalid",
|
|
"langbot-assets-enabled": True,
|
|
"langbot-assets-gateway-host": "127.0.0.1",
|
|
"langbot-assets-gateway-port": 0,
|
|
}
|
|
)
|
|
|
|
async def consume():
|
|
async for _ in runner.invoke(ctx):
|
|
pass
|
|
|
|
task = asyncio.create_task(consume())
|
|
await asyncio.wait_for(entered.wait(), 1)
|
|
task.cancel()
|
|
await asyncio.sleep(0)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert not runner._asset_gateway_pool
|
|
|
|
with load(name) as (module, runner):
|
|
asyncio.run(run(module, runner))
|