xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-09 04:40:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(cloud): harden multi-tenant runtime resources

Browse Source
This commit is contained in:
Junyan Qin
2026-07-29 11:32:26 +08:00
parent 32abbb636f
commit ae85ac2b16
211 changed files with 14963 additions and 1968 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/langbot/libs/coze_server_api/client.py
+42 -8
View File
@@ -6,6 +6,22 @@ from typing import Dict, List, Any, AsyncGenerator
import os
from pathlib import Path
from langbot.pkg.utils import httpclient
_MAX_COZE_RESPONSE_BYTES = 16 * 1024 * 1024
_MAX_COZE_EVENT_BYTES = 1024 * 1024
_MAX_COZE_MEDIA_BYTES = 10 * 1024 * 1024
def _read_local_media_limited(path: Path) -> bytes:
if path.stat().st_size > _MAX_COZE_MEDIA_BYTES:
raise ValueError('Coze upload exceeds the size limit')
with path.open('rb') as handle:
body = handle.read(_MAX_COZE_MEDIA_BYTES + 1)
if len(body) > _MAX_COZE_MEDIA_BYTES:
raise ValueError('Coze upload exceeds the size limit')
return body
class AsyncCozeAPIClient:
def __init__(self, api_key: str, api_base: str = 'https://api.coze.cn'):
@@ -58,19 +74,24 @@ class AsyncCozeAPIClient:
if isinstance(file, Path):
if not file.exists():
raise ValueError(f'File not found: {file}')
with open(file, 'rb') as f:
file = f.read()
file = await asyncio.to_thread(_read_local_media_limited, file)
# 处理文件路径字符串
elif isinstance(file, str):
if not os.path.isfile(file):
raise ValueError(f'File not found: {file}')
with open(file, 'rb') as f:
file = f.read()
file = await asyncio.to_thread(
_read_local_media_limited,
Path(file),
)
# 处理文件对象
elif hasattr(file, 'read'):
file = file.read()
file = await asyncio.to_thread(file.read, _MAX_COZE_MEDIA_BYTES + 1)
if not isinstance(file, (bytes, bytearray)):
raise ValueError('Unsupported Coze upload type')
if len(file) > _MAX_COZE_MEDIA_BYTES:
raise ValueError('Coze upload exceeds the size limit')
session = await self.coze_session()
url = f'{self.api_base}/v1/files/upload'
@@ -87,13 +108,18 @@ class AsyncCozeAPIClient:
if response.status == 401:
raise Exception('Coze API 认证失败，请检查 API Key 是否正确')
response_text = await response.text()
response_text = (
await httpclient.read_limited(
response,
max_bytes=_MAX_COZE_EVENT_BYTES,
)
).decode('utf-8', errors='replace')
if response.status != 200:
raise Exception(f'文件上传失败，状态码: {response.status}, 响应: {response_text}')
try:
result = await response.json()
except json.JSONDecodeError:
result = json.loads(response_text)
except (json.JSONDecodeError, TypeError):
raise Exception(f'文件上传响应解析失败: {response_text}')
if result.get('code') != 0:
@@ -158,7 +184,15 @@ class AsyncCozeAPIClient:
if response.status != 200:
raise Exception(f'Coze API 流式请求失败，状态码: {response.status}')
total_bytes = 0
chunk_type = 'message'
chunk_data = ''
async for chunk in response.content:
total_bytes += len(chunk)
if total_bytes > _MAX_COZE_RESPONSE_BYTES:
raise Exception('Coze API stream exceeds the runtime limit')
if len(chunk) > _MAX_COZE_EVENT_BYTES:
raise Exception('Coze API event exceeds the runtime limit')
chunk = chunk.decode('utf-8')
if chunk != '\n':
if chunk.startswith('event:'):
src/langbot/libs/deerflow_api/client.py
+33 -5
View File
@@ -12,10 +12,26 @@ from collections.abc import AsyncGenerator
import httpx
from langbot.pkg.utils import httpclient
from .errors import DeerFlowAPIError
SSE_MAX_BUFFER_CHARS = 1_048_576
SSE_MAX_TOTAL_BYTES = 16 * 1024 * 1024
ERROR_BODY_MAX_BYTES = 1024 * 1024
async def _read_error_body(response: httpx.Response) -> str:
body = bytearray()
async for chunk in response.aiter_bytes(8192):
body.extend(chunk)
if len(body) > ERROR_BODY_MAX_BYTES:
raise DeerFlowAPIError(
operation='read error response',
body='response exceeds the runtime limit',
)
return body.decode('utf-8', errors='replace')
def _normalize_sse_newlines(text: str) -> str:
@@ -94,6 +110,7 @@ class AsyncDeerFlowClient:
async with httpx.AsyncClient(
trust_env=True,
timeout=timeout,
event_hooks=httpclient.httpx_response_limit_hooks(),
) as client:
response = await client.post(
url,
@@ -101,13 +118,14 @@ class AsyncDeerFlowClient:
json=payload,
)
if response.status_code not in (200, 201):
body = await httpclient.response_text(response)
raise DeerFlowAPIError(
operation='create thread',
status=response.status_code,
body=response.text,
body=body,
url=url,
)
return response.json()
return await httpclient.parse_json_response(response)
async def delete_thread(self, thread_id: str, timeout: float = 20) -> None:
"""删除指定 thread"""
@@ -116,13 +134,15 @@ class AsyncDeerFlowClient:
async with httpx.AsyncClient(
trust_env=True,
timeout=timeout,
event_hooks=httpclient.httpx_response_limit_hooks(),
) as client:
response = await client.delete(url, headers=self.headers)
if response.status_code not in (200, 202, 204, 404):
body = await httpclient.response_text(response)
raise DeerFlowAPIError(
operation='delete thread',
status=response.status_code,
body=response.text,
body=body,
url=url,
thread_id=thread_id,
)
@@ -163,19 +183,27 @@ class AsyncDeerFlowClient:
json=payload,
) as resp:
if resp.status_code != 200:
body = await resp.aread()
raise DeerFlowAPIError(
operation='runs/stream request',
status=resp.status_code,
body=body.decode('utf-8', errors='replace'),
body=await _read_error_body(resp),
url=url,
thread_id=thread_id,
)
decoder = codecs.getincrementaldecoder('utf-8')('replace')
buffer = ''
total_bytes = 0
async for chunk in resp.aiter_bytes(8192):
total_bytes += len(chunk)
if total_bytes > SSE_MAX_TOTAL_BYTES:
raise DeerFlowAPIError(
operation='runs/stream response',
body='response exceeds the runtime limit',
url=url,
thread_id=thread_id,
)
buffer += _normalize_sse_newlines(decoder.decode(chunk))
while '\n\n' in buffer:
src/langbot/libs/dify_service_api/v1/client.py
+176 -111
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import httpx
import typing
import json
@@ -8,6 +9,75 @@ from .errors import DifyAPIError
from pathlib import Path
import os
_MAX_DIFY_RESPONSE_BYTES = 1024 * 1024
_MAX_DIFY_SSE_LINE_BYTES = 1024 * 1024
_MAX_DIFY_STREAM_BYTES = 16 * 1024 * 1024
_MAX_DIFY_UPLOAD_BYTES = 10 * 1024 * 1024
async def _read_limited_response(
response: httpx.Response,
*,
max_bytes: int = _MAX_DIFY_RESPONSE_BYTES,
) -> bytes:
content_length = response.headers.get('Content-Length')
if content_length is not None:
try:
if int(content_length) > max_bytes:
raise DifyAPIError(f'Remote response exceeds the {max_bytes}-byte limit')
except (TypeError, ValueError):
pass
body = bytearray()
async for chunk in response.aiter_bytes(chunk_size=8192):
body.extend(chunk)
if len(body) > max_bytes:
raise DifyAPIError(f'Remote response exceeds the {max_bytes}-byte limit')
return bytes(body)
async def _iter_sse_json(
response: httpx.Response,
) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
"""Parse Dify's one-JSON-per-data-line SSE without unbounded line buffering."""
buffer = bytearray()
total = 0
async for chunk in response.aiter_bytes(chunk_size=8192):
total += len(chunk)
if total > _MAX_DIFY_STREAM_BYTES:
raise DifyAPIError('Dify SSE stream exceeds the runtime limit')
buffer.extend(chunk)
while b'\n' in buffer:
raw_line, _, remainder = buffer.partition(b'\n')
buffer = bytearray(remainder)
if len(raw_line) > _MAX_DIFY_SSE_LINE_BYTES:
raise DifyAPIError('Dify SSE event exceeds the runtime limit')
line = raw_line.rstrip(b'\r').strip()
if not line or not line.startswith(b'data:'):
continue
payload = json.loads(line[5:].decode('utf-8', errors='replace'))
if isinstance(payload, dict):
yield payload
if len(buffer) > _MAX_DIFY_SSE_LINE_BYTES:
raise DifyAPIError('Dify SSE event exceeds the runtime limit')
line = bytes(buffer).rstrip(b'\r').strip()
if line.startswith(b'data:'):
payload = json.loads(line[5:].decode('utf-8', errors='replace'))
if isinstance(payload, dict):
yield payload
def _read_local_file_limited(path: Path) -> bytes:
if path.stat().st_size > _MAX_DIFY_UPLOAD_BYTES:
raise ValueError('Dify upload exceeds the size limit')
with path.open('rb') as handle:
body = handle.read(_MAX_DIFY_UPLOAD_BYTES + 1)
if len(body) > _MAX_DIFY_UPLOAD_BYTES:
raise ValueError('Dify upload exceeds the size limit')
return body
class AsyncDifyServiceClient:
"""Dify Service API 客户端"""
@@ -22,6 +92,21 @@ class AsyncDifyServiceClient:
) -> None:
self.api_key = api_key
self.base_url = base_url
self._client: httpx.AsyncClient | None = None
def _get_client(self) -> httpx.AsyncClient:
if self._client is None:
self._client = httpx.AsyncClient(
base_url=self.base_url,
trust_env=True,
)
return self._client
async def aclose(self) -> None:
client = self._client
self._client = None
if client is not None:
await client.aclose()
async def chat_messages(
self,
@@ -38,37 +123,32 @@ class AsyncDifyServiceClient:
if response_mode != 'streaming':
raise DifyAPIError('当前仅支持 streaming 模式')
async with httpx.AsyncClient(
base_url=self.base_url,
trust_env=True,
timeout=timeout,
) as client:
payload = {
'inputs': inputs,
'query': query,
'user': user,
'response_mode': response_mode,
'conversation_id': conversation_id,
'files': files,
'model_config': model_config or {},
}
client = self._get_client()
payload = {
'inputs': inputs,
'query': query,
'user': user,
'response_mode': response_mode,
'conversation_id': conversation_id,
'files': files,
'model_config': model_config or {},
}
async with client.stream(
'POST',
'/chat-messages',
headers={
'Authorization': f'Bearer {self.api_key}',
'Content-Type': 'application/json',
},
json=payload,
) as r:
async for chunk in r.aiter_lines():
if r.status_code != 200:
raise DifyAPIError(f'{r.status_code} {chunk}')
if chunk.strip() == '':
continue
if chunk.startswith('data:'):
yield json.loads(chunk[5:])
async with client.stream(
'POST',
'/chat-messages',
headers={
'Authorization': f'Bearer {self.api_key}',
'Content-Type': 'application/json',
},
json=payload,
timeout=timeout,
) as r:
if r.status_code != 200:
body = await _read_limited_response(r)
raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
async for event in _iter_sse_json(r):
yield event
async def workflow_run(
self,
@@ -82,32 +162,27 @@ class AsyncDifyServiceClient:
if response_mode != 'streaming':
raise DifyAPIError('当前仅支持 streaming 模式')
async with httpx.AsyncClient(
base_url=self.base_url,
trust_env=True,
client = self._get_client()
async with client.stream(
'POST',
'/workflows/run',
headers={
'Authorization': f'Bearer {self.api_key}',
'Content-Type': 'application/json',
},
json={
'inputs': inputs,
'user': user,
'response_mode': response_mode,
'files': files,
},
timeout=timeout,
) as client:
async with client.stream(
'POST',
'/workflows/run',
headers={
'Authorization': f'Bearer {self.api_key}',
'Content-Type': 'application/json',
},
json={
'inputs': inputs,
'user': user,
'response_mode': response_mode,
'files': files,
},
) as r:
async for chunk in r.aiter_lines():
if r.status_code != 200:
raise DifyAPIError(f'{r.status_code} {chunk}')
if chunk.strip() == '':
continue
if chunk.startswith('data:'):
yield json.loads(chunk[5:])
) as r:
if r.status_code != 200:
body = await _read_limited_response(r)
raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
async for event in _iter_sse_json(r):
yield event
async def workflow_submit(
self,
@@ -129,41 +204,38 @@ class AsyncDifyServiceClient:
'Content-Type': 'application/json',
}
async with httpx.AsyncClient(
base_url=self.base_url,
trust_env=True,
client = self._get_client()
# Step 1: Submit the form
payload: dict[str, typing.Any] = {
'inputs': inputs if isinstance(inputs, dict) else {},
'user': user,
'action': action,
}
async with client.stream(
'POST',
f'/form/human_input/{form_token}',
headers=headers,
json=payload,
timeout=timeout,
) as client:
# Step 1: Submit the form
payload: dict[str, typing.Any] = {
'inputs': inputs if isinstance(inputs, dict) else {},
'user': user,
'action': action,
}
submit_resp = await client.post(
f'/form/human_input/{form_token}',
headers=headers,
json=payload,
)
) as submit_resp:
submit_body = await _read_limited_response(submit_resp)
if submit_resp.status_code != 200:
raise DifyAPIError(f'{submit_resp.status_code} {submit_resp.text}')
raise DifyAPIError(f'{submit_resp.status_code} {submit_body.decode(errors="replace")}')
# Step 2: Stream resumed workflow events
async with client.stream(
'GET',
f'/workflow/{workflow_run_id}/events',
headers={'Authorization': f'Bearer {self.api_key}'},
params={'user': user},
) as r:
if r.status_code != 200:
body = (await r.aread()).decode(errors='replace')
raise DifyAPIError(f'{r.status_code} {body}')
async for chunk in r.aiter_lines():
if chunk.strip() == '':
continue
if chunk.startswith('data:'):
yield json.loads(chunk[5:])
# Step 2: Stream resumed workflow events
async with client.stream(
'GET',
f'/workflow/{workflow_run_id}/events',
headers={'Authorization': f'Bearer {self.api_key}'},
params={'user': user},
timeout=timeout,
) as r:
if r.status_code != 200:
body = await _read_limited_response(r)
raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
async for event in _iter_sse_json(r):
yield event
async def upload_file(
self,
@@ -175,37 +247,30 @@ class AsyncDifyServiceClient:
if isinstance(file, Path):
if not file.exists():
raise ValueError(f'File not found: {file}')
with open(file, 'rb') as f:
file = f.read()
file = await asyncio.to_thread(_read_local_file_limited, file)
# 处理文件路径字符串
elif isinstance(file, str):
if not os.path.isfile(file):
raise ValueError(f'File not found: {file}')
with open(file, 'rb') as f:
file = f.read()
file = await asyncio.to_thread(_read_local_file_limited, Path(file))
# 处理文件对象
elif hasattr(file, 'read'):
file = file.read()
async with httpx.AsyncClient(
base_url=self.base_url,
trust_env=True,
file = await asyncio.to_thread(file.read, _MAX_DIFY_UPLOAD_BYTES + 1)
if len(file) > _MAX_DIFY_UPLOAD_BYTES:
raise ValueError('Dify upload exceeds the size limit')
client = self._get_client()
# multipart/form-data
async with client.stream(
'POST',
'/files/upload',
headers={'Authorization': f'Bearer {self.api_key}'},
files={'file': file},
data={'user': user},
timeout=timeout,
) as client:
# multipart/form-data
response = await client.post(
'/files/upload',
headers={'Authorization': f'Bearer {self.api_key}'},
files={
'file': file,
},
data={
'user': user,
},
)
) as response:
body = await _read_limited_response(response)
if response.status_code != 201:
raise DifyAPIError(f'{response.status_code} {response.text}')
return response.json()
raise DifyAPIError(f'{response.status_code} {body.decode(errors="replace")}')
return json.loads(body)
src/langbot/libs/dingtalk_api/api.py
+133 -46
View File
@@ -7,6 +7,7 @@ import time
import typing
import uuid
import urllib.parse
from contextlib import asynccontextmanager
from typing import Awaitable, Callable, Optional
import dingtalk_stream # type: ignore
import websockets
@@ -15,12 +16,42 @@ from .card_callback import DingTalkCardActionHandler
from .dingtalkevent import DingTalkEvent
import httpx
import traceback
from langbot.pkg.utils import httpclient
_stdout_logger = logging.getLogger('langbot.dingtalk_api')
DINGTALK_OPENAPI_BASE = 'https://api.dingtalk.com'
_MAX_MEDIA_BYTES = 10 * 1024 * 1024
_MAX_GATEWAY_MESSAGE_BYTES = 1024 * 1024
def _read_local_media_limited(file_path: str) -> bytes:
if os.path.getsize(file_path) > _MAX_MEDIA_BYTES:
raise ValueError('DingTalk media exceeds the size limit')
with open(file_path, 'rb') as file:
body = file.read(_MAX_MEDIA_BYTES + 1)
if len(body) > _MAX_MEDIA_BYTES:
raise ValueError('DingTalk media exceeds the size limit')
return body
async def _read_httpx_media_limited(response: httpx.Response) -> bytes:
content_length = response.headers.get('Content-Length')
if content_length is not None:
try:
if int(content_length) > _MAX_MEDIA_BYTES:
raise ValueError('DingTalk media exceeds the size limit')
except (TypeError, ValueError) as exc:
if 'exceeds' in str(exc):
raise
body = bytearray()
async for chunk in response.aiter_bytes():
body.extend(chunk)
if len(body) > _MAX_MEDIA_BYTES:
raise ValueError('DingTalk media exceeds the size limit')
return bytes(body)
def _stringify_card_param_map(card_param_map: Optional[dict]) -> dict:
@@ -44,6 +75,8 @@ def _stringify_card_param_map(card_param_map: Optional[dict]) -> dict:
class DingTalkClient:
_MAX_INBOUND_TASKS = 100
def __init__(
self,
client_id: str,
@@ -86,6 +119,37 @@ class DingTalkClient:
self.legacy_access_token = ''
self.legacy_access_token_expiry_time: typing.Optional[float] = None
self._stopped = False # Flag to control the event loop
self._inbound_tasks: set[asyncio.Task] = set()
self._http_client: httpx.AsyncClient | None = None
@asynccontextmanager
async def _http_client_context(self):
"""Reuse one connection pool while preserving existing call structure."""
if self._http_client is None or self._http_client.is_closed:
self._http_client = httpx.AsyncClient(event_hooks=httpclient.httpx_response_limit_hooks())
yield self._http_client
def _start_inbound_task(self, coro: typing.Coroutine) -> bool:
"""Start one bounded inbound callback task."""
for task in tuple(self._inbound_tasks):
if task.done():
self._inbound_tasks.discard(task)
if len(self._inbound_tasks) >= self._MAX_INBOUND_TASKS:
coro.close()
return False
task = asyncio.create_task(coro)
self._inbound_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
self._inbound_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
return True
async def _on_card_action(self, payload: dict) -> None:
"""Dispatch a parsed card-action payload to the adapter callback."""
@@ -101,11 +165,11 @@ class DingTalkClient:
url = 'https://api.dingtalk.com/v1.0/oauth2/accessToken'
headers = {'Content-Type': 'application/json'}
data = {'appKey': self.key, 'appSecret': self.secret}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
try:
response = await client.post(url, json=data, headers=headers)
if response.status_code == 200:
response_data = response.json()
response_data = await httpclient.parse_json_response(response)
self.access_token = response_data.get('accessToken')
expires_in = int(response_data.get('expireIn', 7200))
self.access_token_expiry_time = time.time() + expires_in - 60
@@ -129,28 +193,28 @@ class DingTalkClient:
url = 'https://api.dingtalk.com/v1.0/robot/messageFiles/download'
params = {'downloadCode': download_code, 'robotCode': self.robot_code}
headers = {'x-acs-dingtalk-access-token': self.access_token}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=params)
if response.status_code == 200:
result = response.json()
result = await httpclient.parse_json_response(response)
download_url = result.get('downloadUrl')
else:
await self.logger.error(f'failed to get download url: {response.json()}')
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'failed to get download url: {error_payload}')
if download_url:
return await self.download_url_to_base64(download_url)
async def download_url_to_base64(self, download_url):
async with httpx.AsyncClient() as client:
response = await client.get(download_url)
if response.status_code == 200:
file_bytes = response.content
mime_type = response.headers.get('Content-Type', 'application/octet-stream')
base64_str = base64.b64encode(file_bytes).decode('utf-8')
return f'data:{mime_type};base64,{base64_str}'
else:
await self.logger.error(f'failed to get files: {response.json()}')
async with self._http_client_context() as client:
async with client.stream('GET', download_url) as response:
if response.status_code == 200:
file_bytes = await _read_httpx_media_limited(response)
mime_type = response.headers.get('Content-Type', 'application/octet-stream')
base64_str = (await asyncio.to_thread(base64.b64encode, file_bytes)).decode('utf-8')
return f'data:{mime_type};base64,{base64_str}'
error_body = await _read_httpx_media_limited(response)
await self.logger.error(f'failed to get files: {error_body[:300]!r}')
async def get_audio_url(self, download_code: str):
if not await self.check_access_token():
@@ -158,17 +222,19 @@ class DingTalkClient:
url = 'https://api.dingtalk.com/v1.0/robot/messageFiles/download'
params = {'downloadCode': download_code, 'robotCode': self.robot_code}
headers = {'x-acs-dingtalk-access-token': self.access_token}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=params)
if response.status_code == 200:
result = response.json()
result = await httpclient.parse_json_response(response)
download_url = result.get('downloadUrl')
if download_url:
return await self.download_url_to_base64(download_url)
else:
await self.logger.error(f'failed to get audio: {response.json()}')
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'failed to get audio: {error_payload}')
else:
raise Exception(f'Error: {response.status_code}, {response.text}')
body = await httpclient.response_text(response)
raise Exception(f'Error: {response.status_code}, {body}')
async def get_file_url(self, download_code: str):
if not await self.check_access_token():
@@ -176,17 +242,19 @@ class DingTalkClient:
url = 'https://api.dingtalk.com/v1.0/robot/messageFiles/download'
params = {'downloadCode': download_code, 'robotCode': self.robot_code}
headers = {'x-acs-dingtalk-access-token': self.access_token}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=params)
if response.status_code == 200:
result = response.json()
result = await httpclient.parse_json_response(response)
download_url = result.get('downloadUrl')
if download_url:
return download_url
else:
await self.logger.error(f'failed to get file: {response.json()}')
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'failed to get file: {error_payload}')
else:
raise Exception(f'Error: {response.status_code}, {response.text}')
body = await httpclient.response_text(response)
raise Exception(f'Error: {response.status_code}, {body}')
async def update_incoming_message(self, message):
"""异步更新 DingTalkClient 中的 incoming_message"""
@@ -503,12 +571,13 @@ class DingTalkClient:
len(content),
)
try:
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=data)
response_body = await httpclient.response_text(response, max_chars=500)
_stdout_logger.info(
'DingTalk send_proactive_message_to_one response: status=%d body=%s',
response.status_code,
response.text[:500],
response_body,
)
if response.status_code == 200:
return
@@ -535,7 +604,7 @@ class DingTalkClient:
'msgParam': json.dumps({'content': content}),
}
try:
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=data)
if response.status_code == 200:
return
@@ -667,22 +736,23 @@ class DingTalkClient:
'DingTalk createAndDeliver request body: %s',
json.dumps(body, ensure_ascii=False)[:1500],
)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=body, timeout=30.0)
response_body = await httpclient.response_text(response, max_chars=500)
if response.status_code == 200:
_stdout_logger.info(
'DingTalk createAndDeliver response: %s',
response.text[:500],
response_body,
)
return True
_stdout_logger.error(
'DingTalk createAndDeliver failed: status=%s body=%s',
response.status_code,
response.text,
response_body,
)
if self.logger:
await self.logger.error(
f'DingTalk createAndDeliver failed: status={response.status_code} body={response.text}'
f'DingTalk createAndDeliver failed: status={response.status_code} body={response_body}'
)
return False
except Exception:
@@ -725,13 +795,14 @@ class DingTalkClient:
'Content-Type': 'application/json',
}
try:
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.put(url, headers=headers, json=body, timeout=30.0)
if response.status_code == 200:
return True
if self.logger:
response_body = await httpclient.response_text(response)
await self.logger.error(
f'DingTalk card streaming failed: status={response.status_code} body={response.text}'
f'DingTalk card streaming failed: status={response.status_code} body={response_body}'
)
return False
except Exception:
@@ -768,18 +839,19 @@ class DingTalkClient:
out_track_id,
json.dumps(body, ensure_ascii=False)[:1500],
)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.put(url, headers=headers, json=body, timeout=30.0)
response_body = await httpclient.response_text(response, max_chars=300)
_stdout_logger.info(
'DingTalk update_card_data response: status=%d body=%s',
response.status_code,
response.text[:300],
response_body,
)
if response.status_code == 200:
return True
if self.logger:
await self.logger.error(
f'DingTalk update card failed: status={response.status_code} body={response.text}'
f'DingTalk update card failed: status={response.status_code} body={response_body}'
)
return False
except Exception:
@@ -808,17 +880,18 @@ class DingTalkClient:
url = 'https://oapi.dingtalk.com/gettoken'
try:
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.get(url, params={'appkey': self.key, 'appsecret': self.secret}, timeout=15.0)
data = response.json() if response.status_code == 200 else {}
data = await httpclient.parse_json_response(response) if response.status_code == 200 else {}
if data.get('errcode') == 0 and data.get('access_token'):
self.legacy_access_token = data['access_token']
expires_in = int(data.get('expires_in', 7200))
self.legacy_access_token_expiry_time = now + expires_in - 60
return self.legacy_access_token
if self.logger:
response_body = await httpclient.response_text(response, max_chars=200)
await self.logger.error(
f'DingTalk legacy gettoken failed: status={response.status_code} body={response.text[:200]}'
f'DingTalk legacy gettoken failed: status={response.status_code} body={response_body}'
)
except Exception:
_stdout_logger.exception('DingTalk legacy gettoken error')
@@ -848,8 +921,7 @@ class DingTalkClient:
url = 'https://oapi.dingtalk.com/media/upload'
try:
with open(file_path, 'rb') as f:
file_bytes = f.read()
file_bytes = await asyncio.to_thread(_read_local_media_limited, file_path)
file_name = os.path.basename(file_path)
# Best-effort content-type guess; DingTalk accepts the major image
# mime types and otherwise infers from the bytes.
@@ -857,20 +929,21 @@ class DingTalkClient:
mime = {'png': 'image/png', 'jpg': 'image/jpeg', 'jpeg': 'image/jpeg', 'gif': 'image/gif'}.get(
ext, 'application/octet-stream'
)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(
url,
params={'access_token': token, 'type': 'image'},
files={'media': (file_name, file_bytes, mime)},
timeout=30.0,
)
data = response.json() if response.status_code == 200 else {}
data = await httpclient.parse_json_response(response) if response.status_code == 200 else {}
if data.get('errcode') == 0 and data.get('media_id'):
_stdout_logger.info('DingTalk upload_image_media OK: media_id=%s', data['media_id'])
return data['media_id']
if self.logger:
response_body = await httpclient.response_text(response, max_chars=300)
await self.logger.error(
f'DingTalk upload_image_media failed: status={response.status_code} body={response.text[:300]}'
f'DingTalk upload_image_media failed: status={response.status_code} body={response_body}'
)
except Exception:
_stdout_logger.exception('DingTalk upload_image_media error')
@@ -897,15 +970,19 @@ class DingTalkClient:
continue
uri = '%s?ticket=%s' % (connection['endpoint'], urllib.parse.quote_plus(connection['ticket']))
async with websockets.connect(uri) as websocket:
async with websockets.connect(uri, max_size=_MAX_GATEWAY_MESSAGE_BYTES) as websocket:
self.client.websocket = websocket
keepalive_task = asyncio.create_task(self._keepalive(websocket))
try:
async for raw_message in websocket:
if self._stopped:
break
json_message = json.loads(raw_message)
asyncio.create_task(self.client.background_task(json_message))
json_message = await asyncio.to_thread(json.loads, raw_message)
if not self._start_inbound_task(self.client.background_task(json_message)):
if self.logger:
await self.logger.warning(
'DingTalk inbound task capacity reached; dropping message'
)
finally:
keepalive_task.cancel()
try:
@@ -948,5 +1025,15 @@ class DingTalkClient:
await self.client.websocket.close()
except Exception:
pass
inbound_tasks = list(self._inbound_tasks)
for task in inbound_tasks:
if not task.done():
task.cancel()
if inbound_tasks:
await asyncio.gather(*inbound_tasks, return_exceptions=True)
self._inbound_tasks.clear()
# Clear message handlers to prevent stale callbacks
self._message_handlers = {'example': []}
if self._http_client is not None:
await self._http_client.aclose()
self._http_client = None
src/langbot/libs/official_account_api/api.py
+100 -9
View File
@@ -21,8 +21,14 @@ xml_template = """
</xml>
"""
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
class OAClient:
_STATE_TTL_SECONDS = 600
_STATE_MAX = 4096
_MAX_CONTENT_CHARS = 200000
def __init__(
self,
token: str,
@@ -41,6 +47,7 @@ class OAClient:
self.access_token = ''
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
# 只有在非统一模式下才注册独立路由
if not self.unified_mode:
@@ -57,8 +64,38 @@ class OAClient:
self.access_token_expiry_time = None
self.msg_id_map = {}
self.generated_content = {}
self._msg_seen_at = {}
self._generated_at = {}
self._last_state_prune = 0.0
self.logger = logger
def _prune_state(self) -> None:
now = time.monotonic()
if now - self._last_state_prune >= 60:
self._last_state_prune = now
for message_id, seen_at in tuple(self._msg_seen_at.items()):
if now - seen_at > self._STATE_TTL_SECONDS:
self._msg_seen_at.pop(message_id, None)
self.msg_id_map.pop(message_id, None)
for message_id, generated_at in tuple(self._generated_at.items()):
if now - generated_at > self._STATE_TTL_SECONDS:
self._generated_at.pop(message_id, None)
self.generated_content.pop(message_id, None)
while len(self.msg_id_map) > self._STATE_MAX:
message_id = next(iter(self.msg_id_map))
self.msg_id_map.pop(message_id, None)
self._msg_seen_at.pop(message_id, None)
while len(self.generated_content) > self._STATE_MAX:
message_id = next(iter(self.generated_content))
self.generated_content.pop(message_id, None)
self._generated_at.pop(message_id, None)
def clear(self) -> None:
self.msg_id_map.clear()
self.generated_content.clear()
self._msg_seen_at.clear()
self._generated_at.clear()
async def handle_callback_request(self):
"""处理回调请求（独立端口模式，使用全局 request）。"""
return await self._handle_callback_internal(request)
@@ -104,8 +141,16 @@ class OAClient:
raise Exception('拒绝请求')
elif req.method == 'POST':
encryt_msg = await req.data
if len(encryt_msg) > _MAX_CALLBACK_BODY_BYTES:
raise ValueError('Official Account callback body exceeds the size limit')
wxcpt = WXBizMsgCrypt(self.token, self.aes, self.appid)
ret, xml_msg = wxcpt.DecryptMsg(encryt_msg, msg_signature, timestamp, nonce)
ret, xml_msg = await asyncio.to_thread(
wxcpt.DecryptMsg,
encryt_msg,
msg_signature,
timestamp,
nonce,
)
xml_msg = xml_msg.decode('utf-8')
if ret != 0:
@@ -118,7 +163,7 @@ class OAClient:
if event:
await self._handle_message(event)
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
from_user = root.find('FromUserName').text # 发送者
to_user = root.find('ToUserName').text # 机器人
@@ -126,6 +171,7 @@ class OAClient:
interval = 0.1
while True:
content = self.generated_content.pop(message_data['MsgId'], None)
self._generated_at.pop(message_data['MsgId'], None)
if content:
response_xml = xml_template.format(
to_user=from_user,
@@ -156,7 +202,7 @@ class OAClient:
traceback.print_exc()
async def get_message(self, xml_msg: str):
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
message_data = {
'ToUserName': root.find('ToUserName').text,
@@ -193,21 +239,30 @@ class OAClient:
处理消息事件。
"""
message_id = event.message_id
self._prune_state()
if message_id in self.msg_id_map.keys():
self.msg_id_map[message_id] += 1
self._msg_seen_at[message_id] = time.monotonic()
return
self.msg_id_map[message_id] = 1
self._msg_seen_at[message_id] = time.monotonic()
msg_type = event.type
if msg_type in self._message_handlers:
for handler in self._message_handlers[msg_type]:
await handler(event)
async def set_message(self, msg_id: int, content: str):
self.generated_content[msg_id] = content
self.generated_content[msg_id] = str(content)[: self._MAX_CONTENT_CHARS]
self._generated_at[msg_id] = time.monotonic()
self._prune_state()
class OAClientForLongerResponse:
_MAX_USERS = 4096
_MAX_MESSAGES_PER_USER = 20
_MAX_CONTENT_CHARS = 200000
def __init__(
self,
token: str,
@@ -227,6 +282,7 @@ class OAClientForLongerResponse:
self.access_token = ''
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
# 只有在非统一模式下才注册独立路由
if not self.unified_mode:
@@ -244,8 +300,28 @@ class OAClientForLongerResponse:
self.loading_message = LoadingMessage
self.msg_queue = {}
self.user_msg_queue = {}
self._last_queue_cleanup = 0.0
self.logger = logger
def _prune_queues(self) -> None:
now = time.monotonic()
if now - self._last_queue_cleanup >= 60:
self._last_queue_cleanup = now
for user_id, queue in tuple(self.msg_queue.items()):
if not queue:
self.msg_queue.pop(user_id, None)
for user_id, queue in tuple(self.user_msg_queue.items()):
if not queue:
self.user_msg_queue.pop(user_id, None)
while len(self.msg_queue) > self._MAX_USERS:
self.msg_queue.pop(next(iter(self.msg_queue)), None)
while len(self.user_msg_queue) > self._MAX_USERS:
self.user_msg_queue.pop(next(iter(self.user_msg_queue)), None)
def clear(self) -> None:
self.msg_queue.clear()
self.user_msg_queue.clear()
async def handle_callback_request(self):
"""处理回调请求（独立端口模式，使用全局 request）。"""
return await self._handle_callback_internal(request)
@@ -285,8 +361,16 @@ class OAClientForLongerResponse:
elif req.method == 'POST':
encryt_msg = await req.data
if len(encryt_msg) > _MAX_CALLBACK_BODY_BYTES:
raise ValueError('Official Account callback body exceeds the size limit')
wxcpt = WXBizMsgCrypt(self.token, self.aes, self.appid)
ret, xml_msg = wxcpt.DecryptMsg(encryt_msg, msg_signature, timestamp, nonce)
ret, xml_msg = await asyncio.to_thread(
wxcpt.DecryptMsg,
encryt_msg,
msg_signature,
timestamp,
nonce,
)
xml_msg = xml_msg.decode('utf-8')
if ret != 0:
@@ -294,7 +378,7 @@ class OAClientForLongerResponse:
raise Exception('消息解密失败')
# 解析 XML
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
from_user = root.find('FromUserName').text
to_user = root.find('ToUserName').text
@@ -305,6 +389,7 @@ class OAClientForLongerResponse:
# 弹出用户消息
if self.user_msg_queue.get(from_user) and self.user_msg_queue[from_user]:
self.user_msg_queue[from_user].pop(0)
self._prune_queues()
response_xml = xml_template.format(
to_user=from_user,
@@ -332,9 +417,13 @@ class OAClientForLongerResponse:
if event:
self.user_msg_queue.setdefault(from_user, []).append(
{
'content': event.message,
'content': str(event.message)[: self._MAX_CONTENT_CHARS],
}
)
self.user_msg_queue[from_user] = self.user_msg_queue[from_user][
-self._MAX_MESSAGES_PER_USER :
]
self._prune_queues()
await self._handle_message(event)
return response_xml
@@ -344,7 +433,7 @@ class OAClientForLongerResponse:
traceback.print_exc()
async def get_message(self, xml_msg: str):
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
message_data = {
'ToUserName': root.find('ToUserName').text,
@@ -393,6 +482,8 @@ class OAClientForLongerResponse:
self.msg_queue[from_user].append(
{
'msg_id': message_id,
'content': content,
'content': str(content)[: self._MAX_CONTENT_CHARS],
}
)
self.msg_queue[from_user] = self.msg_queue[from_user][-self._MAX_MESSAGES_PER_USER :]
self._prune_queues()
src/langbot/libs/openclaw_weixin_api/client.py
+70 -42
View File
@@ -10,7 +10,9 @@ from __future__ import annotations
import asyncio
import base64
import hashlib
import io
import json
import logging
import os
import struct
@@ -21,6 +23,8 @@ from urllib.parse import quote
import aiohttp
from langbot.pkg.utils import httpclient
from .types import (
ApiError,
CDNMedia,
@@ -58,6 +62,51 @@ DEFAULT_BOT_TYPE = '3'
# Maximum text length per message chunk (WeChat limit)
MAX_TEXT_CHUNK_SIZE = 2000
MAX_CDN_MEDIA_BYTES = 10 * 1024 * 1024
async def _response_text(response: aiohttp.ClientResponse) -> str:
body = await httpclient.read_limited(
response,
max_bytes=MAX_CDN_MEDIA_BYTES,
)
return body.decode('utf-8', errors='replace')
async def _response_json(response: aiohttp.ClientResponse) -> dict:
payload = json.loads(await _response_text(response))
if not isinstance(payload, dict):
raise ApiError('OpenClaw API returned a non-object response', status=response.status)
return payload
def _decrypt_cdn_payload(encrypted: bytes, aes_key: bytes) -> bytes:
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7
cipher = Cipher(algorithms.AES(aes_key), modes.ECB())
decryptor = cipher.decryptor()
padded = decryptor.update(encrypted) + decryptor.finalize()
unpadder = PKCS7(128).unpadder()
return unpadder.update(padded) + unpadder.finalize()
def _encrypt_cdn_payload(
file_bytes: bytes,
) -> tuple[str, str, bytes, str]:
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7
raw_key = os.urandom(16)
aes_key_hex = raw_key.hex()
encoded_key = base64.b64encode(aes_key_hex.encode('utf-8')).decode('utf-8')
padder = PKCS7(128).padder()
padded = padder.update(file_bytes) + padder.finalize()
cipher = Cipher(algorithms.AES(raw_key), modes.ECB())
encryptor = cipher.encryptor()
encrypted = encryptor.update(padded) + encryptor.finalize()
raw_md5 = hashlib.md5(file_bytes).hexdigest()
return aes_key_hex, encoded_key, encrypted, raw_md5
def _random_wechat_uin() -> str:
@@ -125,12 +174,12 @@ class OpenClawWeixinClient:
url, json=payload, headers=headers, timeout=aiohttp.ClientTimeout(total=timeout)
) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
raise ApiError(
f'OpenClaw API error {resp.status}: {text}',
status=resp.status,
)
data = await resp.json(content_type=None)
data = await _response_json(resp)
# Check for application-level errors in the response body
errcode = data.get('errcode') or data.get('ret')
@@ -170,12 +219,12 @@ class OpenClawWeixinClient:
timeout=aiohttp.ClientTimeout(total=timeout),
) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
raise ApiError(
f'OpenClaw API error {resp.status}: {text}',
status=resp.status,
)
data = await resp.json(content_type=None)
data = await _response_json(resp)
except (asyncio.TimeoutError, aiohttp.ServerTimeoutError):
return GetUpdatesResponse(ret=0, msgs=[], get_updates_buf=get_updates_buf)
@@ -258,9 +307,6 @@ class OpenClawWeixinClient:
Returns:
Decrypted file bytes.
"""
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7
if not media.encrypt_query_param:
raise ApiError('CDN media has no encrypt_query_param', status=0)
if not media.aes_key:
@@ -285,17 +331,14 @@ class OpenClawWeixinClient:
async with session.get(cdn_url, timeout=aiohttp.ClientTimeout(total=120)) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
raise ApiError(f'CDN download failed: {resp.status} {text}', status=resp.status)
encrypted = await resp.read()
encrypted = await httpclient.read_limited(
resp,
max_bytes=MAX_CDN_MEDIA_BYTES,
)
# Decrypt AES-128-ECB with PKCS7 padding
cipher = Cipher(algorithms.AES(aes_key), modes.ECB())
decryptor = cipher.decryptor()
padded = decryptor.update(encrypted) + decryptor.finalize()
unpadder = PKCS7(128).unpadder()
return unpadder.update(padded) + unpadder.finalize()
return await asyncio.to_thread(_decrypt_cdn_payload, encrypted, aes_key)
async def upload_media(
self,
@@ -313,28 +356,13 @@ class OpenClawWeixinClient:
Returns:
CDNMedia with encrypt_query_param and aes_key for use in sendMessage.
"""
import hashlib
if len(file_bytes) > MAX_CDN_MEDIA_BYTES:
raise ApiError('CDN media exceeds the size limit', status=0)
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7
# 1. Generate random 16-byte AES key
raw_key = os.urandom(16)
aes_key_hex = raw_key.hex() # 32-char hex string
# 2. Encode key for CDNMedia: base64(hex_string) — same for all media types
# Matches official SDK: Buffer.from(aeskey_hex).toString("base64")
encoded_key = base64.b64encode(aes_key_hex.encode('utf-8')).decode('utf-8')
# 3. Encrypt file with AES-128-ECB + PKCS7
padder = PKCS7(128).padder()
padded = padder.update(file_bytes) + padder.finalize()
cipher = Cipher(algorithms.AES(raw_key), modes.ECB())
encryptor = cipher.encryptor()
encrypted = encryptor.update(padded) + encryptor.finalize()
# 4. Get upload URL
raw_md5 = hashlib.md5(file_bytes).hexdigest()
aes_key_hex, encoded_key, encrypted, raw_md5 = await asyncio.to_thread(
_encrypt_cdn_payload,
file_bytes,
)
filekey = os.urandom(16).hex() # 32-char hex, matches official SDK
upload_resp = await self.get_upload_url(
@@ -370,7 +398,7 @@ class OpenClawWeixinClient:
timeout=aiohttp.ClientTimeout(total=120),
) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
logger.error('CDN upload failed: status=%d url=%s body=%s', resp.status, cdn_url, text[:500])
raise ApiError(f'CDN upload failed: {resp.status} {text}', status=resp.status)
download_param = resp.headers.get('x-encrypted-param', '')
@@ -491,12 +519,12 @@ class OpenClawWeixinClient:
async with session.get(url, timeout=aiohttp.ClientTimeout(total=DEFAULT_API_TIMEOUT)) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
raise ApiError(
f'Failed to fetch QR code: {resp.status} {text}',
status=resp.status,
)
data = await resp.json(content_type=None)
data = await _response_json(resp)
logger.debug(
'fetch_qrcode response: qrcode=%s, img=%s', data.get('qrcode'), bool(data.get('qrcode_img_content'))
@@ -536,12 +564,12 @@ class OpenClawWeixinClient:
url, headers=headers, timeout=aiohttp.ClientTimeout(total=DEFAULT_QR_POLL_TIMEOUT)
) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
raise ApiError(
f'Failed to poll QR status: {resp.status} {text}',
status=resp.status,
)
data = await resp.json(content_type=None)
data = await _response_json(resp)
logger.debug('QR status poll response: %s', data)
except (asyncio.TimeoutError, aiohttp.ServerTimeoutError):
return QRStatusResponse(status='wait')
src/langbot/libs/qq_official_api/api.py
+79 -36
View File
@@ -9,10 +9,13 @@ import langbot_plugin.api.entities.builtin.platform.events as platform_events
from .qqofficialevent import QQOfficialEvent
import json
import traceback
from contextlib import asynccontextmanager
from cryptography.hazmat.primitives.asymmetric import ed25519
from langbot.pkg.utils import httpclient
QQ_SELECT_ACTION_PREFIX = '__langbot_select__:'
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
def get_select_field_options(form_data: dict) -> tuple[str, list[str]]:
@@ -152,6 +155,7 @@ class QQOfficialClient:
def __init__(self, secret: str, token: str, app_id: str, logger: None, unified_mode: bool = False):
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
# 只有在非统一模式下才注册独立路由
if not self.unified_mode:
@@ -176,6 +180,32 @@ class QQOfficialClient:
self.logger = logger
self._msg_seq_counter = 0
self._token_refresh_task: Optional[asyncio.Task] = None
self._http_clients: dict[float | None, httpx.AsyncClient] = {}
@asynccontextmanager
async def _http_client_context(self, timeout: float | None = None):
client = self._http_clients.get(timeout)
if client is None or client.is_closed:
response_hooks = httpclient.httpx_response_limit_hooks()
client = (
httpx.AsyncClient(event_hooks=response_hooks)
if timeout is None
else httpx.AsyncClient(timeout=timeout, event_hooks=response_hooks)
)
self._http_clients[timeout] = client
yield client
async def close(self) -> None:
"""Stop client-owned background work."""
if self._token_refresh_task and not self._token_refresh_task.done():
self._token_refresh_task.cancel()
await asyncio.gather(self._token_refresh_task, return_exceptions=True)
self._token_refresh_task = None
clients = list(self._http_clients.values())
self._http_clients.clear()
if clients:
await asyncio.gather(*(client.aclose() for client in clients), return_exceptions=True)
async def check_access_token(self):
"""检查access_token是否存在"""
@@ -186,7 +216,7 @@ class QQOfficialClient:
async def get_access_token(self):
"""获取access_token"""
url = 'https://bots.qq.com/app/getAppAccessToken'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'appId': self.app_id,
'clientSecret': self.secret,
@@ -196,8 +226,9 @@ class QQOfficialClient:
}
response = await client.post(url, json=params, headers=headers)
if response.status_code != 200:
raise Exception(f'Failed to get access_token: HTTP {response.status_code} {response.text}')
response_data = response.json()
body = await httpclient.response_text(response)
raise Exception(f'Failed to get access_token: HTTP {response.status_code} {body}')
response_data = await httpclient.parse_json_response(response)
access_token = response_data.get('access_token')
expires_in = int(response_data.get('expires_in', 7200))
self.access_token_expiry_time = time.time() + expires_in - 60
@@ -236,8 +267,10 @@ class QQOfficialClient:
if not body or len(body) == 0:
await self.logger.info('Received empty body, might be health check or GET request')
return {'code': 0, 'message': 'ok'}, 200
if len(body) > _MAX_CALLBACK_BODY_BYTES:
return {'error': 'callback body exceeds the size limit'}, 413
payload = json.loads(body)
payload = await asyncio.to_thread(json.loads, body)
if payload.get('op') == 13:
validation_data = payload.get('d')
@@ -367,7 +400,7 @@ class QQOfficialClient:
await self.get_access_token()
url = self.base_url + '/v2/users/' + user_openid + '/messages'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -382,7 +415,7 @@ class QQOfficialClient:
if event_id:
data['event_id'] = event_id
response = await client.post(url, headers=headers, json=data)
response_data = response.json()
response_data = await httpclient.parse_json_response(response)
if response.status_code == 200:
return
else:
@@ -406,7 +439,7 @@ class QQOfficialClient:
await self.get_access_token()
url = self.base_url + '/v2/groups/' + group_openid + '/messages'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -424,8 +457,9 @@ class QQOfficialClient:
if response.status_code == 200:
return
else:
await self.logger.error(f'Failed to send group message: {response.json()}')
raise Exception(response.read().decode())
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'Failed to send group message: {error_payload}')
raise Exception(str(error_payload))
async def send_channle_group_text_msg(self, channel_id: str, content: str, msg_id: str):
"""发送频道群聊消息"""
@@ -433,7 +467,7 @@ class QQOfficialClient:
await self.get_access_token()
url = self.base_url + '/channels/' + channel_id + '/messages'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -447,7 +481,8 @@ class QQOfficialClient:
if response.status_code == 200:
return True
else:
await self.logger.error(f'Failed to send channel group message: {response.json()}')
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'Failed to send channel group message: {error_payload}')
raise Exception(response)
async def send_channle_private_text_msg(self, guild_id: str, content: str, msg_id: str):
@@ -456,7 +491,7 @@ class QQOfficialClient:
await self.get_access_token()
url = self.base_url + '/dms/' + guild_id + '/messages'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -470,7 +505,8 @@ class QQOfficialClient:
if response.status_code == 200:
return True
else:
await self.logger.error(f'Failed to send channel private message: {response.json()}')
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'Failed to send channel private message: {error_payload}')
raise Exception(response)
# ---- 富媒体消息 ----
@@ -532,20 +568,21 @@ class QQOfficialClient:
if file_type == self.MEDIA_TYPE_FILE and file_name:
body['file_name'] = file_name
async with httpx.AsyncClient(timeout=120) as client:
async with self._http_client_context(timeout=120) as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
}
response = await client.post(url, headers=headers, json=body)
if response.status_code == 200:
data = response.json()
data = await httpclient.parse_json_response(response)
file_info = data.get('file_info', '')
preview = file_info[:80] + '...' if len(file_info) > 80 else file_info
await self.logger.info(f'Upload media success, file_info={preview}')
return file_info
else:
raise Exception(f'Failed to upload media: HTTP {response.status_code} {response.text}')
body = await httpclient.response_text(response)
raise Exception(f'Failed to upload media: HTTP {response.status_code} {body}')
async def _send_media_msg(
self,
@@ -578,7 +615,7 @@ class QQOfficialClient:
if msg_id:
body['msg_id'] = msg_id
async with httpx.AsyncClient(timeout=120) as client:
async with self._http_client_context(timeout=120) as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -586,7 +623,8 @@ class QQOfficialClient:
await self.logger.info(f'Sending rich media: {json.dumps(body, ensure_ascii=False)[:200]}')
response = await client.post(url, headers=headers, json=body)
if response.status_code != 200:
raise Exception(f'Failed to send rich media message: HTTP {response.status_code} {response.text}')
response_body = await httpclient.response_text(response)
raise Exception(f'Failed to send rich media message: HTTP {response.status_code} {response_body}')
async def send_image_msg(
self,
@@ -678,15 +716,16 @@ class QQOfficialClient:
if stream_msg_id:
body['stream_msg_id'] = stream_msg_id
async with httpx.AsyncClient(timeout=120) as client:
async with self._http_client_context(timeout=120) as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
}
response = await client.post(url, headers=headers, json=body)
if response.status_code != 200:
raise Exception(f'Failed to send stream message: HTTP {response.status_code} {response.text}')
return response.json()
response_body = await httpclient.response_text(response)
raise Exception(f'Failed to send stream message: HTTP {response.status_code} {response_body}')
return await httpclient.parse_json_response(response)
async def send_markdown_keyboard(
self,
@@ -743,18 +782,19 @@ class QQOfficialClient:
if event_id:
body['event_id'] = event_id
async with httpx.AsyncClient(timeout=30) as client:
async with self._http_client_context(timeout=30) as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
}
response = await client.post(url, headers=headers, json=body)
if response.status_code != 200:
response_body = await httpclient.response_text(response)
await self.logger.error(
f'Failed to send markdown+keyboard: HTTP {response.status_code} {response.text}'
f'Failed to send markdown+keyboard: HTTP {response.status_code} {response_body}'
)
raise Exception(f'Failed to send markdown+keyboard: HTTP {response.status_code} {response.text}')
return response.json()
raise Exception(f'Failed to send markdown+keyboard: HTTP {response.status_code} {response_body}')
return await httpclient.parse_json_response(response)
async def ack_interaction(self, interaction_id: str, code: int = 0) -> None:
"""Acknowledge a button-click INTERACTION_CREATE event.
@@ -775,7 +815,7 @@ class QQOfficialClient:
await self.get_access_token()
url = f'{self.base_url}/interactions/{interaction_id}'
async with httpx.AsyncClient(timeout=10) as client:
async with self._http_client_context(timeout=10) as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -783,8 +823,9 @@ class QQOfficialClient:
try:
response = await client.put(url, headers=headers, json={'code': code})
if response.status_code >= 400:
response_body = await httpclient.response_text(response)
await self.logger.warning(
f'ack_interaction non-success: HTTP {response.status_code} {response.text}'
f'ack_interaction non-success: HTTP {response.status_code} {response_body}'
)
except Exception as e:
await self.logger.warning(f'ack_interaction error (non-fatal): {e}')
@@ -796,10 +837,11 @@ class QQOfficialClient:
return time.time() > self.access_token_expiry_time
async def repeat_seed(self, bot_secret: str, target_size: int = 32) -> bytes:
seed = bot_secret
while len(seed) < target_size:
seed *= 2
return seed[:target_size].encode('utf-8')
if not bot_secret:
raise ValueError('QQ bot secret must not be empty')
target_size = max(int(target_size), 1)
repeats = (target_size + len(bot_secret) - 1) // len(bot_secret)
return (bot_secret * repeats)[:target_size].encode('utf-8')
async def verify(self, validation_payload: dict):
seed = await self.repeat_seed(self.secret)
@@ -843,19 +885,20 @@ class QQOfficialClient:
await self.get_access_token()
url = f'{self.base_url}/gateway'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
}
response = await client.get(url, headers=headers)
if response.status_code == 200:
data = response.json()
data = await httpclient.parse_json_response(response)
ws_url = data.get('url', '')
if not ws_url:
raise Exception('Gateway URL is empty')
return ws_url
else:
raise Exception(f'Failed to get Gateway URL: HTTP {response.status_code} {response.text}')
body = await httpclient.response_text(response)
raise Exception(f'Failed to get Gateway URL: HTTP {response.status_code} {body}')
async def _background_token_refresh(self):
"""在 token 到期前主动刷新"""
@@ -935,7 +978,7 @@ class QQOfficialClient:
try:
await self.logger.info('Connecting to WebSocket gateway...')
ws = await websockets.connect(ws_url)
ws = await websockets.connect(ws_url, max_size=_MAX_CALLBACK_BODY_BYTES)
await self.logger.info('WebSocket connected')
except Exception as e:
await self.logger.error(f'WebSocket connection failed: {e}')
@@ -948,7 +991,7 @@ class QQOfficialClient:
try:
async for raw_msg in ws:
try:
payload = json.loads(raw_msg)
payload = await asyncio.to_thread(json.loads, raw_msg)
except json.JSONDecodeError:
await self.logger.error(f'Failed to parse message: {raw_msg}')
continue
src/langbot/libs/slack_api/api.py
+7 -1
View File
@@ -1,3 +1,4 @@
import asyncio
import json
import traceback
from quart import Quart, jsonify, request
@@ -6,6 +7,8 @@ from .slackevent import SlackEvent
from typing import Callable
import langbot_plugin.api.entities.builtin.platform.events as platform_events
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
class SlackClient:
def __init__(self, bot_token: str, signing_secret: str, logger: None, unified_mode: bool = False):
@@ -13,6 +16,7 @@ class SlackClient:
self.signing_secret = signing_secret
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
self.client = AsyncWebClient(self.bot_token)
# 只有在非统一模式下才注册独立路由
@@ -50,7 +54,9 @@ class SlackClient:
"""
try:
body = await req.get_data()
data = json.loads(body)
if len(body) > _MAX_CALLBACK_BODY_BYTES:
raise ValueError('Slack callback body exceeds the size limit')
data = await asyncio.to_thread(json.loads, body)
if 'type' in data:
if data['type'] == 'url_verification':
return data['challenge']
src/langbot/libs/wechatpad_api/api/downloadpai.py
+37 -11
View File
@@ -1,7 +1,32 @@
from langbot.libs.wechatpad_api.util.http_util import post_json
import httpx
import asyncio
import base64
import httpx
from langbot.libs.wechatpad_api.util.http_util import post_json
from langbot.pkg.utils import httpclient
_MAX_WECHATPAD_MEDIA_BYTES = 16 * 1024 * 1024
async def _read_media_limited(response: httpx.Response) -> bytes:
content_length = response.headers.get('content-length')
if content_length is not None:
try:
declared_size = int(content_length)
except ValueError:
declared_size = None
if declared_size is not None and declared_size > _MAX_WECHATPAD_MEDIA_BYTES:
raise RuntimeError('WeChatPad media exceeds the runtime limit')
body = bytearray()
async for chunk in response.aiter_bytes(chunk_size=64 * 1024):
body.extend(chunk)
if len(body) > _MAX_WECHATPAD_MEDIA_BYTES:
raise RuntimeError('WeChatPad media exceeds the runtime limit')
return bytes(body)
class DownloadApi:
def __init__(self, base_url, token):
@@ -19,12 +44,13 @@ class DownloadApi:
return post_json(url, token=self.token, data=json_data)
async def download_url_to_base64(self, download_url):
async with httpx.AsyncClient() as client:
response = await client.get(download_url)
if response.status_code == 200:
file_bytes = response.content
base64_str = base64.b64encode(file_bytes).decode('utf-8') # 返回字符串格式
return base64_str
else:
raise Exception('获取文件失败')
async with httpx.AsyncClient(
timeout=30,
event_hooks=httpclient.httpx_response_limit_hooks(_MAX_WECHATPAD_MEDIA_BYTES),
) as client:
async with client.stream('GET', download_url) as response:
if response.status_code != 200:
raise RuntimeError('获取文件失败')
file_bytes = await _read_media_limited(response)
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
return encoded.decode('utf-8')
src/langbot/libs/wechatpad_api/util/http_util.py
+48 -11
View File
@@ -1,6 +1,29 @@
import json as json_module
import requests
from langbot.pkg.utils import httpclient
_MAX_WECHATPAD_RESPONSE_BYTES = 16 * 1024 * 1024
def _read_requests_response_limited(response: requests.Response) -> dict:
content_length = response.headers.get('Content-Length')
if content_length is not None:
try:
if int(content_length) > _MAX_WECHATPAD_RESPONSE_BYTES:
raise RuntimeError('WeChatPad response exceeds the runtime limit')
except (TypeError, ValueError):
pass
body = bytearray()
for chunk in response.iter_content(chunk_size=64 * 1024):
body.extend(chunk)
if len(body) > _MAX_WECHATPAD_RESPONSE_BYTES:
raise RuntimeError('WeChatPad response exceeds the runtime limit')
result = json_module.loads(body)
if not isinstance(result, dict):
raise RuntimeError('WeChatPad returned a non-object response')
return result
def post_json(base_url, token, data=None):
headers = {'Content-Type': 'application/json'}
@@ -8,16 +31,21 @@ def post_json(base_url, token, data=None):
url = base_url + f'?key={token}'
try:
response = requests.post(url, json=data, headers=headers, timeout=60)
response.raise_for_status()
result = response.json()
with requests.post(
url,
json=data,
headers=headers,
timeout=60,
stream=True,
) as response:
response.raise_for_status()
result = _read_requests_response_limited(response)
if result:
return result
else:
raise RuntimeError(response.text)
raise RuntimeError('WeChatPad returned an empty response')
except Exception as e:
print(f'http请求失败, url={url}, exception={e}')
raise RuntimeError(str(e))
@@ -27,16 +55,20 @@ def get_json(base_url, token):
url = base_url + f'?key={token}'
try:
response = requests.get(url, headers=headers, timeout=60)
response.raise_for_status()
result = response.json()
with requests.get(
url,
headers=headers,
timeout=60,
stream=True,
) as response:
response.raise_for_status()
result = _read_requests_response_limited(response)
if result:
return result
else:
raise RuntimeError(response.text)
raise RuntimeError('WeChatPad returned an empty response')
except Exception as e:
print(f'http请求失败, url={url}, exception={e}')
raise RuntimeError(str(e))
@@ -68,7 +100,12 @@ async def async_request(
method=method, url=url, params=params, headers=headers, data=data, json=json
) as response:
response.raise_for_status() # 如果状态码不是200，抛出异常
result = await response.json()
result = json_module.loads(
await httpclient.read_limited(
response,
max_bytes=_MAX_WECHATPAD_RESPONSE_BYTES,
)
)
# print(result)
return result
# if result.get('Code') == 200:
src/langbot/libs/wecom_ai_bot_api/api.py
+124 -27
View File
@@ -10,13 +10,16 @@ import re
from typing import Any, Callable, Optional, Tuple
from urllib.parse import unquote
import httpx
from Crypto.Cipher import AES
from quart import Quart, request, Response, jsonify
from langbot.libs.wecom_ai_bot_api import wecombotevent
from langbot.libs.wecom_ai_bot_api.WXBizMsgCrypt3 import WXBizMsgCrypt
from langbot.pkg.platform.logger import EventLogger
from langbot.pkg.utils import httpclient
_CLIENT_TRANSIENT_CACHE_MAX = 4096
_MAX_STREAM_CONTENT_CHARS = 200000
@dataclass
@@ -56,7 +59,7 @@ class StreamSession:
last_access: float = field(default_factory=time.time)
# 将流水线增量结果缓存到队列，刷新请求逐条消费
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
queue: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=1))
# 是否已经完成（收到最终片段）
finished: bool = False
@@ -85,6 +88,7 @@ class StreamSessionManager:
# full like → cancel → dislike feedback flow. Must align with the adapter's
# _stream_to_monitoring_msg TTL (wecombot.py).
_FEEDBACK_SESSION_TTL = 600 # 10 minutes
_MAX_SESSIONS = 4096
def __init__(self, logger: EventLogger, ttl: int = 60) -> None:
self.logger = logger
@@ -165,6 +169,26 @@ class StreamSessionManager:
if task_id:
self._task_index.pop(task_id, None)
def clear(self) -> None:
"""Release every retained stream and reverse index."""
self._sessions.clear()
self._msg_index.clear()
self._feedback_index.clear()
self._task_index.clear()
def _drop_session(self, stream_id: str) -> StreamSession | None:
session = self._sessions.pop(stream_id, None)
if session is None:
return None
if session.msg_id and self._msg_index.get(session.msg_id) == stream_id:
self._msg_index.pop(session.msg_id, None)
if session.feedback_id:
self._feedback_index.pop(session.feedback_id, None)
if session.pending_form_task_id:
self._task_index.pop(session.pending_form_task_id, None)
return session
def create_or_get(self, msg_json: dict[str, Any]) -> tuple[StreamSession, bool]:
"""根据企业微信回调创建或获取会话。
@@ -185,6 +209,14 @@ class StreamSessionManager:
session.last_access = time.time()
return session, False
self.cleanup()
while len(self._sessions) >= self._MAX_SESSIONS:
oldest_stream_id = min(
self._sessions,
key=lambda candidate: self._sessions[candidate].last_access,
)
self._drop_session(oldest_stream_id)
stream_id = str(uuid.uuid4())
session = StreamSession(
stream_id=stream_id,
@@ -221,8 +253,13 @@ class StreamSessionManager:
try:
session.queue.put_nowait(chunk)
except asyncio.QueueFull:
# 默认无界队列，此处兜底防御
await session.queue.put(chunk)
# Each chunk is a complete snapshot. Coalesce a slow consumer to
# the newest value instead of retaining every intermediate body.
try:
session.queue.get_nowait()
except asyncio.QueueEmpty:
pass
session.queue.put_nowait(chunk)
if chunk.is_final:
session.finished = True
@@ -265,7 +302,7 @@ class StreamSessionManager:
session.finished = True
session.last_access = time.time()
def cleanup(self) -> None:
def cleanup(self) -> list[str]:
"""定期清理过期会话，防止队列与映射无上限累积。
已注册 feedback_id 的会话使用更长的 TTL，确保用户在点赞/取消/点踩流程中
@@ -279,16 +316,14 @@ class StreamSessionManager:
if now - session.last_access > effective_ttl:
expired.append(stream_id)
removed_msg_ids: list[str] = []
for stream_id in expired:
session = self._sessions.pop(stream_id, None)
session = self._drop_session(stream_id)
if not session:
continue
msg_id = session.msg_id
if msg_id and self._msg_index.get(msg_id) == stream_id:
self._msg_index.pop(msg_id, None)
# Clean up feedback index for expired sessions
if session.feedback_id:
self._feedback_index.pop(session.feedback_id, None)
if session.msg_id:
removed_msg_ids.append(session.msg_id)
return removed_msg_ids
def _decrypt_file(encrypted_data: bytes, aes_key_str: str) -> bytes:
@@ -405,19 +440,19 @@ async def download_encrypted_file(
filename: Optional[str] = None
try:
async with httpx.AsyncClient(timeout=30.0) as client:
response = await client.get(download_url)
if response.status_code != 200:
await logger.error(f'Failed to download file (HTTP {response.status_code}): {response.text[:200]}')
client = httpclient.get_session()
async with client.get(download_url, timeout=30.0) as response:
if response.status != 200:
await logger.error(f'Failed to download file (HTTP {response.status})')
return None, None
encrypted_bytes = response.content
encrypted_bytes = await httpclient.read_limited(response)
filename = _extract_filename(response.headers.get('content-disposition', ''))
except Exception:
await logger.error(f'Failed to download file: {traceback.format_exc()}')
return None, None
try:
decrypted = _decrypt_file(encrypted_bytes, aes_key)
decrypted = await asyncio.to_thread(_decrypt_file, encrypted_bytes, aes_key)
return decrypted, filename
except Exception:
await logger.error(f'Failed to decrypt file: {traceback.format_exc()}')
@@ -466,7 +501,7 @@ async def parse_wecom_bot_message(
"""Download, decrypt, and convert to data URI for backward compatibility."""
data, _filename = await _safe_download(url, per_msg_aeskey)
if data:
return _bytes_to_data_uri(data)
return await asyncio.to_thread(_bytes_to_data_uri, data)
return None
if msg_type == 'text':
@@ -579,7 +614,10 @@ async def parse_wecom_bot_message(
if (file_data.get('filesize') or 0) <= max_inline_file_size:
file_bytes, dl_filename = await _safe_download(download_url, item_aeskey)
if file_bytes:
file_data['base64'] = _bytes_to_data_uri(file_bytes)
file_data['base64'] = await asyncio.to_thread(
_bytes_to_data_uri,
file_bytes,
)
if dl_filename and not file_data.get('filename'):
file_data['filename'] = dl_filename
files.append(file_data)
@@ -1567,6 +1605,8 @@ def build_multiple_interaction_update_card(
class WecomBotClient:
_MAX_DISPATCH_TASKS = 100
def __init__(
self,
Token: str,
@@ -1613,6 +1653,7 @@ class WecomBotClient:
self._feedback_callback: Optional[Callable] = None
self._card_action_callback: Optional[Callable] = None
self._stream_last_content: dict[str, str] = {}
self._dispatch_tasks: set[asyncio.Task] = set()
# Optional `source` block injected into every interactive template_card
# the client builds. Set via `set_card_source` from the adapter after
# reading config. Format: {icon_url, desc, desc_color}.
@@ -1695,7 +1736,12 @@ class WecomBotClient:
"""
reply_plain_str = json.dumps(payload, ensure_ascii=False)
reply_timestamp = str(int(time.time()))
ret, encrypt_text = self.wxcpt.EncryptMsg(reply_plain_str, nonce, reply_timestamp)
ret, encrypt_text = await asyncio.to_thread(
self.wxcpt.EncryptMsg,
reply_plain_str,
nonce,
reply_timestamp,
)
if ret != 0:
await self.logger.error(f'加密失败: {ret}')
return jsonify({'error': 'encrypt_failed'}), 500
@@ -1718,6 +1764,41 @@ class WecomBotClient:
except Exception:
await self.logger.error(traceback.format_exc())
def _start_dispatch_task(self, event: wecombotevent.WecomBotEvent) -> bool:
"""Start one bounded pipeline dispatch task."""
for task in tuple(self._dispatch_tasks):
if task.done():
self._dispatch_tasks.discard(task)
if len(self._dispatch_tasks) >= self._MAX_DISPATCH_TASKS:
return False
task = asyncio.create_task(self._dispatch_event(event))
self._dispatch_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
self._dispatch_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
return True
async def close(self) -> None:
"""Cancel callbacks and release retained webhook state."""
dispatch_tasks = list(self._dispatch_tasks)
for task in dispatch_tasks:
if not task.done():
task.cancel()
if dispatch_tasks:
await asyncio.gather(*dispatch_tasks, return_exceptions=True)
self._dispatch_tasks.clear()
self.generated_content.clear()
self.msg_id_map.clear()
self._stream_last_content.clear()
self.stream_sessions.clear()
async def _handle_post_initial_response(self, msg_json: dict[str, Any], nonce: str) -> tuple[Response, int]:
"""处理企业微信首次推送的消息，返回 stream_id 并开启流水线。
@@ -1747,7 +1828,8 @@ class WecomBotClient:
await self.logger.error(traceback.format_exc())
else:
if is_new:
asyncio.create_task(self._dispatch_event(event))
if not self._start_dispatch_task(event):
await self.logger.warning('WeCom webhook dispatch capacity reached; dropping message')
payload = self._build_stream_payload(session.stream_id, '', False, feedback_id)
return await self._encrypt_and_reply(payload, nonce)
@@ -1870,7 +1952,10 @@ class WecomBotClient:
async def _handle_post_callback(self, req) -> tuple[Response, int] | Response:
"""处理企业微信的 POST 回调请求。"""
self.stream_sessions.cleanup()
for expired_msg_id in self.stream_sessions.cleanup():
self.generated_content.pop(expired_msg_id, None)
self._stream_last_content.pop(expired_msg_id, None)
self.msg_id_map.pop(expired_msg_id, None)
msg_signature = unquote(req.args.get('msg_signature', ''))
timestamp = unquote(req.args.get('timestamp', ''))
@@ -1883,12 +1968,18 @@ class WecomBotClient:
return Response('Bad Request', status=400)
xml_post_data = f'<xml><Encrypt><![CDATA[{encrypted_msg}]]></Encrypt></xml>'
ret, decrypted_xml = self.wxcpt.DecryptMsg(xml_post_data, msg_signature, timestamp, nonce)
ret, decrypted_xml = await asyncio.to_thread(
self.wxcpt.DecryptMsg,
xml_post_data,
msg_signature,
timestamp,
nonce,
)
if ret != 0:
await self.logger.error('解密失败')
return Response('解密失败', status=400)
msg_json = json.loads(decrypted_xml)
msg_json = await asyncio.to_thread(json.loads, decrypted_xml)
event_type = extract_wecom_event_type(msg_json)
@@ -2014,6 +2105,8 @@ class WecomBotClient:
self.msg_id_map[message_id] += 1
return
self.msg_id_map[message_id] = 1
while len(self.msg_id_map) > _CLIENT_TRANSIENT_CACHE_MAX:
self.msg_id_map.pop(next(iter(self.msg_id_map)), None)
msg_type = event.type
if msg_type in self._message_handlers:
for handler in self._message_handlers[msg_type]:
@@ -2047,6 +2140,8 @@ class WecomBotClient:
next_content = previous_content
else:
next_content = previous_content + content if previous_content else content
if len(next_content) > _MAX_STREAM_CONTENT_CHARS:
next_content = next_content[-_MAX_STREAM_CONTENT_CHARS:]
if not is_final and next_content == previous_content:
return True
@@ -2096,7 +2191,9 @@ class WecomBotClient:
"""
handled = await self.push_stream_chunk(msg_id, content, is_final=True)
if not handled:
self.generated_content[msg_id] = content
self.generated_content[msg_id] = content[-_MAX_STREAM_CONTENT_CHARS:]
while len(self.generated_content) > _CLIENT_TRANSIENT_CACHE_MAX:
self.generated_content.pop(next(iter(self.generated_content)), None)
def on_message(self, msg_type: str):
def decorator(func: Callable[[wecombotevent.WecomBotEvent], None]):
@@ -2119,7 +2216,7 @@ class WecomBotClient:
async def download_url_to_base64(self, download_url, encoding_aes_key):
data, _filename = await download_encrypted_file(download_url, encoding_aes_key, self.logger)
if data:
return _bytes_to_data_uri(data)
return await asyncio.to_thread(_bytes_to_data_uri, data)
return None
async def run_task(self, host: str, port: int, *args, **kwargs):
src/langbot/libs/wecom_ai_bot_api/ws_client.py
+134 -16
View File
@@ -47,6 +47,17 @@ CMD_RESPOND_WELCOME = 'aibot_respond_welcome_msg'
CMD_RESPOND_UPDATE = 'aibot_respond_update_msg'
CMD_SEND_MSG = 'aibot_send_msg'
_DEDUP_CACHE_MAX = 4096
_STREAM_CACHE_MAX = 1024
_FEEDBACK_CACHE_MAX = 4096
_PENDING_FORM_MAX = 1024
_PENDING_FORM_TTL_SECONDS = 1800
_MAX_STREAM_CONTENT_CHARS = 200000
_MAX_CALLBACK_TASKS = 100
_MAX_REPLY_WORKERS = 100
_MAX_REPLY_QUEUE_SIZE = 100
_MAX_PENDING_ACKS = 256
def _generate_req_id(prefix: str) -> str:
"""Generate a unique request ID in the format: {prefix}_{timestamp}_{random}."""
@@ -106,6 +117,7 @@ class WecomBotWsClient:
# Per-req_id serial reply queues
self._reply_queues: dict[str, asyncio.Queue] = {}
self._reply_workers: dict[str, asyncio.Task] = {}
self._callback_tasks: set[asyncio.Task] = set()
self._reply_ack_timeout = 5.0
# Stream ID tracking for WebSocket mode
@@ -135,6 +147,31 @@ class WecomBotWsClient:
# `set_card_source` from the adapter after reading config.
self.card_source: Optional[dict] = None
@staticmethod
def _cap_mapping(mapping: dict, max_entries: int) -> None:
while len(mapping) > max_entries:
mapping.pop(next(iter(mapping)), None)
def _prune_stream_state(self) -> None:
while len(self._stream_sessions) > _STREAM_CACHE_MAX:
msg_id = next(iter(self._stream_sessions))
self._stream_sessions.pop(msg_id, None)
self._stream_ids.pop(msg_id, None)
self._stream_last_content.pop(msg_id, None)
task_id = self._task_id_by_msg.pop(msg_id, None)
if task_id:
self._pending_forms_by_task.pop(task_id, None)
def _prune_pending_forms(self) -> None:
cutoff = time.monotonic() - _PENDING_FORM_TTL_SECONDS
for task_id, pending in tuple(self._pending_forms_by_task.items()):
if float(pending.get('created_at', 0.0)) <= cutoff:
self._drop_pending_form_task(task_id, pending)
while len(self._pending_forms_by_task) > _PENDING_FORM_MAX:
task_id = next(iter(self._pending_forms_by_task))
pending = self._pending_forms_by_task.get(task_id, {})
self._drop_pending_form_task(task_id, pending)
# ── Public API ──────────────────────────────────────────────────
async def connect(self):
@@ -173,17 +210,40 @@ class WecomBotWsClient:
async def disconnect(self):
"""Gracefully disconnect from the WebSocket server."""
self._running = False
heartbeat_tasks = []
if self._heartbeat_task and not self._heartbeat_task.done():
self._heartbeat_task.cancel()
for task in self._reply_workers.values():
heartbeat_tasks.append(self._heartbeat_task)
reply_workers = list(self._reply_workers.values())
for task in reply_workers:
if not task.done():
task.cancel()
callback_tasks = list(self._callback_tasks)
for task in callback_tasks:
if not task.done():
task.cancel()
shutdown_tasks = [*heartbeat_tasks, *reply_workers, *callback_tasks]
if shutdown_tasks:
await asyncio.gather(*shutdown_tasks, return_exceptions=True)
self._clear_pending_acks('Connection closed')
if self._ws and not self._ws.closed:
await self._ws.close()
self._ws = None
if self._session and not self._session.closed:
await self._session.close()
self._session = None
self._heartbeat_task = None
self._reply_queues.clear()
self._reply_workers.clear()
self._callback_tasks.clear()
self._stream_ids.clear()
self._stream_last_content.clear()
self._stream_sessions.clear()
self._feedback_sessions.clear()
self._msg_feedback_ids.clear()
self._pending_forms_by_task.clear()
self._task_id_by_msg.clear()
self._msg_id_map.clear()
def on_message(self, msg_type: str) -> Callable:
"""Decorator to register a message handler.
@@ -366,8 +426,10 @@ class WecomBotWsClient:
'chat_id': session_info.get('chat_id', ''),
'stream_id': stream_id,
'req_id': req_id,
'created_at': time.monotonic(),
}
self._task_id_by_msg[msg_id] = task_id
self._prune_pending_forms()
card_payload = build_human_input_template_card_payload(
form_data,
@@ -458,6 +520,8 @@ class WecomBotWsClient:
next_content = previous_content
else:
next_content = previous_content + content if previous_content else content
if len(next_content) > _MAX_STREAM_CONTENT_CHARS:
next_content = next_content[-_MAX_STREAM_CONTENT_CHARS:]
# Skip sending if content hasn't changed (e.g. during tool call argument streaming)
if not is_final and next_content == previous_content:
@@ -485,6 +549,8 @@ class WecomBotWsClient:
session_info = self._stream_sessions.get(msg_id)
if session_info:
self._feedback_sessions[feedback_id] = session_info
self._cap_mapping(self._feedback_sessions, _FEEDBACK_CACHE_MAX)
self._cap_mapping(self._msg_feedback_ids, _FEEDBACK_CACHE_MAX)
# WeCom replaces the displayed stream content on each refresh, so
# every frame must contain the complete snapshot, not only a delta.
@@ -516,7 +582,7 @@ class WecomBotWsClient:
self._session = aiohttp.ClientSession()
try:
self._ws = await self._session.ws_connect(self.ws_url)
self._ws = await self._session.ws_connect(self.ws_url, max_msg_size=1024 * 1024)
self._missed_pong_count = 0
self._reconnect_attempts = 0
await self.logger.info('WebSocket connected, sending auth...')
@@ -539,6 +605,8 @@ class WecomBotWsClient:
finally:
if self._heartbeat_task and not self._heartbeat_task.done():
self._heartbeat_task.cancel()
await asyncio.gather(self._heartbeat_task, return_exceptions=True)
self._heartbeat_task = None
self._clear_pending_acks('Connection closed')
finally:
if self._ws and not self._ws.closed:
@@ -565,7 +633,7 @@ class WecomBotWsClient:
try:
msg = await asyncio.wait_for(self._ws.receive(), timeout=10.0)
if msg.type in (aiohttp.WSMsgType.TEXT,):
frame = json.loads(msg.data)
frame = await asyncio.to_thread(json.loads, msg.data)
req_id = frame.get('headers', {}).get('req_id', '')
if req_id.startswith(CMD_SUBSCRIBE) and frame.get('errcode') == 0:
return True
@@ -614,7 +682,7 @@ class WecomBotWsClient:
break
if msg.type == aiohttp.WSMsgType.TEXT:
try:
frame = json.loads(msg.data)
frame = await asyncio.to_thread(json.loads, msg.data)
await self._handle_frame(frame)
except json.JSONDecodeError:
await self.logger.error(f'Failed to parse WebSocket message: {str(msg.data)[:200]}')
@@ -622,7 +690,7 @@ class WecomBotWsClient:
await self.logger.error(f'Error handling frame: {traceback.format_exc()}')
elif msg.type == aiohttp.WSMsgType.BINARY:
try:
frame = json.loads(msg.data)
frame = await asyncio.to_thread(json.loads, msg.data)
await self._handle_frame(frame)
except Exception:
await self.logger.error(f'Error handling binary frame: {traceback.format_exc()}')
@@ -638,12 +706,14 @@ class WecomBotWsClient:
# Message push
if cmd == CMD_MSG_CALLBACK:
asyncio.create_task(self._handle_message_callback(frame))
if not self._start_callback_task(self._handle_message_callback(frame)):
await self.logger.warning('WeCom WebSocket callback capacity reached; dropping message')
return
# Event push
if cmd == CMD_EVENT_CALLBACK:
asyncio.create_task(self._handle_event_callback(frame))
if not self._start_callback_task(self._handle_event_callback(frame)):
await self.logger.warning('WeCom WebSocket callback capacity reached; dropping event')
return
# No cmd → response/ACK frame, dispatch by req_id prefix
@@ -665,6 +735,27 @@ class WecomBotWsClient:
# Unknown frame
await self.logger.warning(f'Unknown frame: {_frame_snippet(frame)}')
def _start_callback_task(self, coro) -> bool:
"""Start one bounded inbound frame callback."""
for task in tuple(self._callback_tasks):
if task.done():
self._callback_tasks.discard(task)
if len(self._callback_tasks) >= _MAX_CALLBACK_TASKS:
coro.close()
return False
task = asyncio.create_task(coro)
self._callback_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
self._callback_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
return True
async def _handle_message_callback(self, frame: dict):
"""Handle an incoming message callback frame."""
try:
@@ -697,6 +788,7 @@ class WecomBotWsClient:
'chat_id': message_data.get('chatid', ''),
'chat_type': message_data.get('type', 'single'),
}
self._prune_stream_state()
message_data['stream_id'] = stream_id
message_data['req_id'] = req_id
@@ -748,7 +840,7 @@ class WecomBotWsClient:
)
# Look up session by feedback_id
session_info = self._feedback_sessions.get(feedback_id)
session_info = self._feedback_sessions.pop(feedback_id, None)
session = None
if session_info:
session = StreamSession(
@@ -806,6 +898,10 @@ class WecomBotWsClient:
if pending is None:
await self.logger.warning(f'No pending_form found for task_id={task_id} (ws); card event ignored')
return
if time.monotonic() - float(pending.get('created_at', 0.0)) > _PENDING_FORM_TTL_SECONDS:
self._drop_pending_form_task(task_id, pending)
await self.logger.warning(f'Pending form expired for task_id={task_id} (ws)')
return
req_id_for_update = frame.get('headers', {}).get('req_id', '')
form_data = pending.get('form_data', {}) or {}
@@ -868,6 +964,7 @@ class WecomBotWsClient:
self._msg_id_map[message_id] += 1
return
self._msg_id_map[message_id] = 1
self._cap_mapping(self._msg_id_map, _DEDUP_CACHE_MAX)
msg_type = event.type
if msg_type in self._message_handlers:
@@ -899,40 +996,61 @@ class WecomBotWsClient:
# Ensure serial delivery per req_id
if req_id not in self._reply_queues:
self._reply_queues[req_id] = asyncio.Queue()
if len(self._reply_queues) >= _MAX_REPLY_WORKERS:
await self.logger.warning('WeCom WebSocket reply worker capacity reached; dropping reply')
return None
self._reply_queues[req_id] = asyncio.Queue(maxsize=_MAX_REPLY_QUEUE_SIZE)
self._reply_workers[req_id] = asyncio.create_task(self._reply_queue_worker(req_id))
future: asyncio.Future = asyncio.get_event_loop().create_future()
await self._reply_queues[req_id].put((frame, future))
try:
self._reply_queues[req_id].put_nowait((frame, future))
except asyncio.QueueFull:
await self.logger.warning(f'WeCom WebSocket reply queue full for req_id={req_id}; dropping reply')
return None
return await future
async def _reply_queue_worker(self, req_id: str):
"""Process reply queue items serially for a given req_id."""
queue = self._reply_queues[req_id]
current_future: asyncio.Future | None = None
try:
while self._running:
try:
frame, future = await asyncio.wait_for(queue.get(), timeout=60.0)
frame, current_future = await asyncio.wait_for(queue.get(), timeout=60.0)
except asyncio.TimeoutError:
# Queue idle, clean up worker
break
try:
ack = await self._send_and_wait_ack(frame)
if not future.done():
future.set_result(ack)
if not current_future.done():
current_future.set_result(ack)
except Exception as e:
if not future.done():
future.set_exception(e)
if not current_future.done():
current_future.set_exception(e)
finally:
current_future = None
except asyncio.CancelledError:
pass
if current_future is not None and not current_future.done():
current_future.set_exception(ConnectionError('Connection closed'))
finally:
while True:
try:
_, future = queue.get_nowait()
except asyncio.QueueEmpty:
break
if not future.done():
future.set_exception(ConnectionError('Reply worker stopped'))
self._reply_queues.pop(req_id, None)
self._reply_workers.pop(req_id, None)
async def _send_and_wait_ack(self, frame: dict) -> Optional[dict]:
"""Send a frame and wait for the corresponding ACK."""
req_id = frame['headers']['req_id']
if len(self._pending_acks) >= _MAX_PENDING_ACKS:
await self.logger.warning('WeCom WebSocket pending ACK capacity reached; dropping frame')
return None
ack_future: asyncio.Future = asyncio.get_event_loop().create_future()
self._pending_acks[req_id] = ack_future
src/langbot/libs/wecom_api/api.py
+150 -41
View File
@@ -1,16 +1,82 @@
from quart import request
from .WXBizMsgCrypt3 import WXBizMsgCrypt
import asyncio
import base64
import binascii
import contextvars
import functools
import httpx
import os
import traceback
from urllib.parse import quote
from quart import Quart
import xml.etree.ElementTree as ET
from contextlib import asynccontextmanager
from typing import Callable, Dict, Any
from .wecomevent import WecomEvent
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import aiofiles
from langbot.pkg.utils import httpclient
_MAX_MEDIA_BYTES = 10 * 1024 * 1024
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
_EXTENDED_HTTP_TIMEOUT_SECONDS = 120
async def _read_httpx_media_limited(response: httpx.Response) -> bytes:
content_length = response.headers.get('Content-Length')
if content_length is not None:
try:
if int(content_length) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
except (TypeError, ValueError) as exc:
if 'exceeds' in str(exc):
raise
content = bytearray()
async for chunk in response.aiter_bytes():
content.extend(chunk)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
return bytes(content)
async def _read_local_media_limited(path: str) -> bytes:
if await asyncio.to_thread(os.path.getsize, path) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
async with aiofiles.open(path, 'rb') as file:
content = await file.read(_MAX_MEDIA_BYTES + 1)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
return content
async def _decode_media_base64_limited(value: str) -> bytes:
max_encoded_chars = 4 * ((_MAX_MEDIA_BYTES + 2) // 3) + 4
if len(value) > max_encoded_chars:
raise ValueError('WeCom media exceeds the size limit')
content = await asyncio.to_thread(base64.b64decode, value)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
return content
def _bounded_token_retry(method):
"""Allow one token-refresh retry without unbounded async recursion."""
depth = contextvars.ContextVar(f'{method.__name__}_token_retry_depth', default=0)
@functools.wraps(method)
async def wrapped(*args, **kwargs):
current_depth = depth.get()
if current_depth >= 2:
raise RuntimeError(f'{method.__name__} exceeded the token refresh retry limit')
token = depth.set(current_depth + 1)
try:
return await method(*args, **kwargs)
finally:
depth.reset(token)
return wrapped
class WecomClient:
@@ -36,6 +102,7 @@ class WecomClient:
self.logger = logger
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
# 只有在非统一模式下才注册独立路由
if not self.unified_mode:
@@ -49,6 +116,29 @@ class WecomClient:
self._message_handlers = {
'example': [],
}
self._http_clients: dict[bool, httpx.AsyncClient] = {}
@asynccontextmanager
async def _http_client_context(self, *, unbounded_timeout: bool = False):
client = self._http_clients.get(unbounded_timeout)
if client is None or client.is_closed:
response_hooks = httpclient.httpx_response_limit_hooks()
client = (
httpx.AsyncClient(
timeout=_EXTENDED_HTTP_TIMEOUT_SECONDS,
event_hooks=response_hooks,
)
if unbounded_timeout
else httpx.AsyncClient(event_hooks=response_hooks)
)
self._http_clients[unbounded_timeout] = client
yield client
async def close(self) -> None:
clients = list(self._http_clients.values())
self._http_clients.clear()
if clients:
await asyncio.gather(*(client.aclose() for client in clients), return_exceptions=True)
# access——token操作
async def check_access_token(self):
@@ -59,15 +149,16 @@ class WecomClient:
async def get_access_token(self, secret):
url = f'{self.base_url}/gettoken?corpid={self.corpid}&corpsecret={secret}'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.get(url)
data = response.json()
data = await httpclient.parse_json_response(response)
if 'access_token' in data:
return data['access_token']
else:
await self.logger.error(f'获取accesstoken失败:{response.json()}')
await self.logger.error(f'获取accesstoken失败:{data}')
raise Exception(f'未获取access token: {data}')
@_bounded_token_retry
async def get_user_info(self, userid: str) -> dict:
"""
Get user information by user ID using the application secret.
@@ -82,9 +173,9 @@ class WecomClient:
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/user/get?access_token=' + self.access_token + '&userid=' + quote(userid)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.get(url)
data = response.json()
data = await httpclient.parse_json_response(response)
if data.get('errcode') == 40014 or data.get('errcode') == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.get_user_info(userid)
@@ -98,13 +189,13 @@ class WecomClient:
self.access_token_for_contacts = await self.get_access_token(self.secret_for_contacts)
url = self.base_url + '/user/list_id?access_token=' + self.access_token_for_contacts
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'cursor': '',
'limit': 10000,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 0:
dept_users = data['dept_user']
userid = []
@@ -121,7 +212,7 @@ class WecomClient:
url = self.base_url + '/message/send?access_token=' + self.access_token_for_contacts
user_ids = await self.get_users()
user_ids_string = '|'.join(user_ids)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'touser': user_ids_string,
'msgtype': 'text',
@@ -135,16 +226,17 @@ class WecomClient:
'duplicate_check_interval': 1800,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] != 0:
raise Exception('Failed to send message: ' + str(data))
@_bounded_token_retry
async def send_image(self, user_id: str, agent_id: int, media_id: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/message/send?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'touser': user_id,
'msgtype': 'image',
@@ -158,7 +250,7 @@ class WecomClient:
'duplicate_check_interval': 1800,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.send_image(user_id, agent_id, media_id)
@@ -166,11 +258,12 @@ class WecomClient:
await self.logger.error(f'发送图片失败:{data}')
raise Exception('Failed to send image: ' + str(data))
@_bounded_token_retry
async def send_voice(self, user_id: str, agent_id: int, media_id: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/message/send?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'touser': user_id,
'msgtype': 'voice',
@@ -184,7 +277,7 @@ class WecomClient:
'duplicate_check_interval': 1800,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.send_voice(user_id, agent_id, media_id)
@@ -192,11 +285,12 @@ class WecomClient:
await self.logger.error(f'发送语音失败:{data}')
raise Exception('Failed to send voice: ' + str(data))
@_bounded_token_retry
async def send_file(self, user_id: str, agent_id: int, media_id: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/message/send?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'touser': user_id,
'msgtype': 'file',
@@ -210,7 +304,7 @@ class WecomClient:
'duplicate_check_interval': 1800,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.send_file(user_id, agent_id, media_id)
@@ -218,12 +312,13 @@ class WecomClient:
await self.logger.error(f'发送文件失败:{data}')
raise Exception('Failed to send file: ' + str(data))
@_bounded_token_retry
async def send_private_msg(self, user_id: str, agent_id: int, content: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/message/send?access_token=' + self.access_token
async with httpx.AsyncClient(timeout=None) as client:
async with self._http_client_context(unbounded_timeout=True) as client:
params = {
'touser': user_id,
'msgtype': 'text',
@@ -237,7 +332,7 @@ class WecomClient:
'duplicate_check_interval': 1800,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.send_private_msg(user_id, agent_id, content)
@@ -283,7 +378,15 @@ class WecomClient:
elif req.method == 'POST':
encrypt_msg = await req.data
ret, xml_msg = wxcpt.DecryptMsg(encrypt_msg, msg_signature, timestamp, nonce)
if len(encrypt_msg) > _MAX_CALLBACK_BODY_BYTES:
raise ValueError('WeCom callback body exceeds the size limit')
ret, xml_msg = await asyncio.to_thread(
wxcpt.DecryptMsg,
encrypt_msg,
msg_signature,
timestamp,
nonce,
)
if ret != 0:
await self.logger.error('消息解密失败')
raise Exception(f'消息解密失败，错误码: {ret}')
@@ -332,7 +435,7 @@ class WecomClient:
"""
解析微信返回的 XML 消息并转换为字典。
"""
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
message_data = {
'ToUserName': root.find('ToUserName').text,
'FromUserName': root.find('FromUserName').text,
@@ -366,6 +469,7 @@ class WecomClient:
return ext
return 'jpg' # 默认返回jpg
@_bounded_token_retry
async def upload_image_to_work(self, image: platform_message.Image):
"""
获取 media_id
@@ -379,9 +483,8 @@ class WecomClient:
# 获取文件的二进制数据
if image.path:
async with aiofiles.open(image.path, 'rb') as f:
file_bytes = await f.read()
file_name = image.path.split('/')[-1]
file_bytes = await _read_local_media_limited(image.path)
file_name = image.path.split('/')[-1]
elif image.url:
file_bytes = await self.download_media_to_bytes(image.url)
file_name = image.url.split('/')[-1]
@@ -392,7 +495,7 @@ class WecomClient:
base64_data = base64_data.split(',', 1)[1]
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
padded_base64 = base64_data + '=' * padding
file_bytes = base64.b64decode(padded_base64)
file_bytes = await _decode_media_base64_limited(padded_base64)
except binascii.Error as e:
raise ValueError(f'Invalid base64 string: {str(e)}')
else:
@@ -400,6 +503,8 @@ class WecomClient:
raise ValueError('image对象出错')
# 设置 multipart/form-data 格式的文件
if len(file_bytes) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
boundary = '-------------------------acebdf13572468'
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
body = (
@@ -413,9 +518,9 @@ class WecomClient:
)
# 上传文件
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, content=body)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
media_id = await self.upload_image_to_work(image)
@@ -426,6 +531,7 @@ class WecomClient:
media_id = data.get('media_id')
return media_id
@_bounded_token_retry
async def upload_voice_to_work(self, voice: platform_message.Voice):
"""
上传语音文件到企业微信
@@ -437,9 +543,8 @@ class WecomClient:
file_name = 'voice.mp3'
if voice.path:
async with aiofiles.open(voice.path, 'rb') as f:
file_bytes = await f.read()
file_name = voice.path.split('/')[-1]
file_bytes = await _read_local_media_limited(voice.path)
file_name = voice.path.split('/')[-1]
elif voice.url:
file_bytes = await self.download_media_to_bytes(voice.url)
file_name = voice.url.split('/')[-1]
@@ -450,13 +555,15 @@ class WecomClient:
base64_data = base64_data.split(',', 1)[1]
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
padded_base64 = base64_data + '=' * padding
file_bytes = base64.b64decode(padded_base64)
file_bytes = await _decode_media_base64_limited(padded_base64)
except binascii.Error as e:
raise ValueError(f'Invalid base64 string: {str(e)}')
else:
await self.logger.error('Voice对象出错')
raise ValueError('voice对象出错')
if len(file_bytes) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
boundary = '-------------------------acebdf13572468'
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
body = (
@@ -470,9 +577,9 @@ class WecomClient:
)
# print(body)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, content=body)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
media_id = await self.upload_voice_to_work(voice)
@@ -482,6 +589,7 @@ class WecomClient:
media_id = data.get('media_id')
return media_id
@_bounded_token_retry
async def upload_file_to_work(self, file: platform_message.File):
"""
上传文件到企业微信
@@ -492,9 +600,8 @@ class WecomClient:
file_bytes = None
file_name = 'file.txt'
if file.path:
async with aiofiles.open(file.path, 'rb') as f:
file_bytes = await f.read()
file_name = file.path.split('/')[-1]
file_bytes = await _read_local_media_limited(file.path)
file_name = file.path.split('/')[-1]
elif file.url:
file_bytes = await self.download_media_to_bytes(file.url)
file_name = file.url.split('/')[-1]
@@ -505,12 +612,14 @@ class WecomClient:
base64_data = base64_data.split(',', 1)[1]
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
padded_base64 = base64_data + '=' * padding
file_bytes = base64.b64decode(padded_base64)
file_bytes = await _decode_media_base64_limited(padded_base64)
except binascii.Error as e:
raise ValueError(f'Invalid base64 string: {str(e)}')
else:
await self.logger.error('File对象出错')
raise ValueError('file对象出错')
if len(file_bytes) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
boundary = '-------------------------acebdf13572468'
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
body = (
@@ -522,9 +631,9 @@ class WecomClient:
+ file_bytes
+ f'\r\n--{boundary}--\r\n'.encode('utf-8')
)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, content=body)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
media_id = await self.upload_file_to_work(file)
@@ -535,10 +644,10 @@ class WecomClient:
return media_id
async def download_media_to_bytes(self, url: str) -> bytes:
async with httpx.AsyncClient() as client:
response = await client.get(url)
response.raise_for_status()
return response.content
async with self._http_client_context() as client:
async with client.stream('GET', url) as response:
response.raise_for_status()
return await _read_httpx_media_limited(response)
# 进行media_id的获取
async def get_media_id(self, media: platform_message.Image | platform_message.Voice | platform_message.File):
src/langbot/libs/wecom_customer_service_api/api.py
+145 -39
View File
@@ -1,8 +1,13 @@
from quart import request
from ..wecom_api.WXBizMsgCrypt3 import WXBizMsgCrypt
import asyncio
import base64
import binascii
import contextvars
import functools
import httpx
import json
import os
import traceback
from quart import Quart
import xml.etree.ElementTree as ET
@@ -11,9 +16,72 @@ from .wecomcsevent import WecomCSEvent
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import aiofiles
import time
from contextlib import asynccontextmanager
from langbot.pkg.utils import httpclient
_MAX_MEDIA_BYTES = 10 * 1024 * 1024
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
async def _read_httpx_media_limited(response: httpx.Response) -> bytes:
content_length = response.headers.get('Content-Length')
if content_length is not None:
try:
if int(content_length) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
except (TypeError, ValueError) as exc:
if 'exceeds' in str(exc):
raise
content = bytearray()
async for chunk in response.aiter_bytes():
content.extend(chunk)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
return bytes(content)
async def _read_local_media_limited(path: str) -> bytes:
if await asyncio.to_thread(os.path.getsize, path) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
async with aiofiles.open(path, 'rb') as file:
content = await file.read(_MAX_MEDIA_BYTES + 1)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
return content
async def _decode_media_base64_limited(value: str) -> bytes:
max_encoded_chars = 4 * ((_MAX_MEDIA_BYTES + 2) // 3) + 4
if len(value) > max_encoded_chars:
raise ValueError('WeCom customer-service media exceeds the size limit')
content = await asyncio.to_thread(base64.b64decode, value)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
return content
def _bounded_token_retry(method):
"""Allow one token-refresh retry without unbounded async recursion."""
depth = contextvars.ContextVar(f'{method.__name__}_token_retry_depth', default=0)
@functools.wraps(method)
async def wrapped(*args, **kwargs):
current_depth = depth.get()
if current_depth >= 2:
raise RuntimeError(f'{method.__name__} exceeded the token refresh retry limit')
token = depth.set(current_depth + 1)
try:
return await method(*args, **kwargs)
finally:
depth.reset(token)
return wrapped
class WecomCSClient:
_CUSTOMER_CACHE_MAX = 4096
def __init__(
self,
corpid: str,
@@ -34,10 +102,12 @@ class WecomCSClient:
self.logger = logger
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
# Customer info cache: {external_userid: (info_dict, timestamp)}
self._customer_cache: dict[str, tuple[dict, float]] = {}
self._cache_ttl = 60 # Cache TTL in seconds (1 minute)
self._customer_cache_cleanup_at = 0.0
# 只有在非统一模式下才注册独立路由
if not self.unified_mode:
@@ -48,29 +118,40 @@ class WecomCSClient:
self._message_handlers = {
'example': [],
}
self._http_client: httpx.AsyncClient | None = None
@asynccontextmanager
async def _http_client_context(self):
if self._http_client is None or self._http_client.is_closed:
self._http_client = httpx.AsyncClient(event_hooks=httpclient.httpx_response_limit_hooks())
yield self._http_client
async def close(self) -> None:
if self._http_client is not None:
await self._http_client.aclose()
self._http_client = None
@_bounded_token_retry
async def get_pic_url(self, media_id: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = f'{self.base_url}/media/get?access_token={self.access_token}&media_id={media_id}'
async with httpx.AsyncClient() as client:
response = await client.get(url)
if response.headers.get('Content-Type', '').startswith('application/json'):
data = response.json()
if data.get('errcode') in [40014, 42001]:
self.access_token = await self.get_access_token(self.secret)
return await self.get_pic_url(media_id)
else:
async with self._http_client_context() as client:
async with client.stream('GET', url) as response:
image_bytes = await _read_httpx_media_limited(response)
content_type = response.headers.get('Content-Type', '')
if content_type.startswith('application/json'):
data = json.loads(image_bytes)
if data.get('errcode') in [40014, 42001]:
self.access_token = await self.get_access_token(self.secret)
return await self.get_pic_url(media_id)
raise Exception('Failed to get image: ' + str(data))
# 否则是图片，转成 base64
image_bytes = response.content
content_type = response.headers.get('Content-Type', '')
base64_str = base64.b64encode(image_bytes).decode('utf-8')
base64_str = f'data:{content_type};base64,{base64_str}'
return base64_str
# 否则是图片，转成 base64
base64_str = (await asyncio.to_thread(base64.b64encode, image_bytes)).decode('utf-8')
return f'data:{content_type};base64,{base64_str}'
# access——token操作
async def check_access_token(self):
@@ -81,19 +162,20 @@ class WecomCSClient:
async def get_access_token(self, secret):
url = f'{self.base_url}/gettoken?corpid={self.corpid}&corpsecret={secret}'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.get(url)
data = response.json()
data = await httpclient.parse_json_response(response)
if 'access_token' in data:
return data['access_token']
else:
raise Exception(f'未获取access token: {data}')
@_bounded_token_retry
async def get_detailed_message_list(self, xml_msg: str):
# 在本方法中解析消息，并且获得消息的具体内容
if isinstance(xml_msg, bytes):
xml_msg = xml_msg.decode('utf-8')
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
token = root.find('Token').text
open_kfid = root.find('OpenKfId').text
@@ -106,14 +188,14 @@ class WecomCSClient:
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/kf/sync_msg?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'token': token,
'voice_format': 0,
'open_kfid': open_kfid,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.get_detailed_message_list(xml_msg)
@@ -130,11 +212,12 @@ class WecomCSClient:
# await self.change_service_status(userid=external_userid,openkfid=open_kfid,servicer=servicer)
return last_msg_data
@_bounded_token_retry
async def change_service_status(self, userid: str, openkfid: str, servicer: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/kf/service_state/get?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'open_kfid': openkfid,
'external_userid': userid,
@@ -142,18 +225,19 @@ class WecomCSClient:
'servicer_userid': servicer,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.change_service_status(userid, openkfid)
return await self.change_service_status(userid, openkfid, servicer)
if data['errcode'] != 0:
raise Exception('Failed to change service status: ' + str(data))
@_bounded_token_retry
async def send_image(self, user_id: str, agent_id: int, media_id: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/media/upload?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'touser': user_id,
'toparty': '',
@@ -170,7 +254,7 @@ class WecomCSClient:
}
try:
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
except Exception as e:
raise Exception('Failed to send image: ' + str(e))
@@ -182,6 +266,7 @@ class WecomCSClient:
if data['errcode'] != 0:
raise Exception('Failed to send image: ' + str(data))
@_bounded_token_retry
async def send_text_msg(self, open_kfid: str, external_userid: str, msgid: str, content: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
@@ -198,10 +283,10 @@ class WecomCSClient:
},
}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, json=payload)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.send_text_msg(open_kfid, external_userid, msgid, content)
@@ -250,7 +335,15 @@ class WecomCSClient:
elif req.method == 'POST':
encrypt_msg = await req.data
ret, xml_msg = wxcpt.DecryptMsg(encrypt_msg, msg_signature, timestamp, nonce)
if len(encrypt_msg) > _MAX_CALLBACK_BODY_BYTES:
raise ValueError('WeCom customer-service callback body exceeds the size limit')
ret, xml_msg = await asyncio.to_thread(
wxcpt.DecryptMsg,
encrypt_msg,
msg_signature,
timestamp,
nonce,
)
if ret != 0:
raise Exception(f'消息解密失败，错误码: {ret}')
@@ -315,6 +408,7 @@ class WecomCSClient:
return ext
return 'jpg' # 默认返回jpg
@_bounded_token_retry
async def upload_to_work(self, image: platform_message.Image):
"""
获取 media_id
@@ -328,9 +422,8 @@ class WecomCSClient:
# 获取文件的二进制数据
if image.path:
async with aiofiles.open(image.path, 'rb') as f:
file_bytes = await f.read()
file_name = image.path.split('/')[-1]
file_bytes = await _read_local_media_limited(image.path)
file_name = image.path.split('/')[-1]
elif image.url:
file_bytes = await self.download_image_to_bytes(image.url)
file_name = image.url.split('/')[-1]
@@ -341,13 +434,15 @@ class WecomCSClient:
base64_data = base64_data.split(',', 1)[1]
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
padded_base64 = base64_data + '=' * padding
file_bytes = base64.b64decode(padded_base64)
file_bytes = await _decode_media_base64_limited(padded_base64)
except binascii.Error as e:
raise ValueError(f'Invalid base64 string: {str(e)}')
else:
raise ValueError('image对象出错')
# 设置 multipart/form-data 格式的文件
if len(file_bytes) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
boundary = '-------------------------acebdf13572468'
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
body = (
@@ -361,9 +456,9 @@ class WecomCSClient:
)
# 上传文件
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, content=body)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
media_id = await self.upload_to_work(image)
@@ -374,16 +469,17 @@ class WecomCSClient:
return media_id
async def download_image_to_bytes(self, url: str) -> bytes:
async with httpx.AsyncClient() as client:
response = await client.get(url)
response.raise_for_status()
return response.content
async with self._http_client_context() as client:
async with client.stream('GET', url) as response:
response.raise_for_status()
return await _read_httpx_media_limited(response)
# 进行media_id的获取
async def get_media_id(self, image: platform_message.Image):
media_id = await self.upload_to_work(image=image)
return media_id
@_bounded_token_retry
async def get_customer_info(self, external_userid: str) -> dict | None:
"""
Get customer information by external_userid with caching.
@@ -398,6 +494,11 @@ class WecomCSClient:
"""
# Check cache first
current_time = time.time()
if current_time - self._customer_cache_cleanup_at >= 30:
self._customer_cache_cleanup_at = current_time
for user_id, (_, cached_time) in tuple(self._customer_cache.items()):
if current_time - cached_time >= self._cache_ttl:
self._customer_cache.pop(user_id, None)
if external_userid in self._customer_cache:
cached_info, cached_time = self._customer_cache[external_userid]
if current_time - cached_time < self._cache_ttl:
@@ -413,9 +514,9 @@ class WecomCSClient:
'external_userid_list': [external_userid],
}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, json=payload)
data = response.json()
data = await httpclient.parse_json_response(response)
if data.get('errcode') in [40014, 42001]:
self.access_token = await self.get_access_token(self.secret)
@@ -431,5 +532,10 @@ class WecomCSClient:
customer_info = customer_list[0]
# Store in cache
self._customer_cache[external_userid] = (customer_info, current_time)
while len(self._customer_cache) > self._CUSTOMER_CACHE_MAX:
self._customer_cache.pop(next(iter(self._customer_cache)), None)
return customer_info
return None
def clear(self) -> None:
self._customer_cache.clear()
src/langbot/libs/weknora_api/client.py
+71 -35
View File
@@ -6,6 +6,56 @@ import json
from .errors import WeKnoraAPIError
_MAX_WENKORA_RESPONSE_BYTES = 1024 * 1024
_MAX_WENKORA_STREAM_BYTES = 16 * 1024 * 1024
_MAX_WENKORA_SSE_LINE_BYTES = 1024 * 1024
async def _read_limited_response(response: httpx.Response) -> bytes:
body = bytearray()
async for chunk in response.aiter_bytes(chunk_size=8192):
body.extend(chunk)
if len(body) > _MAX_WENKORA_RESPONSE_BYTES:
raise WeKnoraAPIError('WeKnora response exceeds the runtime limit')
return bytes(body)
async def _iter_sse_json(
response: httpx.Response,
) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
buffer = bytearray()
total = 0
async for chunk in response.aiter_bytes(chunk_size=8192):
total += len(chunk)
if total > _MAX_WENKORA_STREAM_BYTES:
raise WeKnoraAPIError('WeKnora stream exceeds the runtime limit')
buffer.extend(chunk)
while b'\n' in buffer:
raw_line, _, remainder = buffer.partition(b'\n')
buffer = bytearray(remainder)
if len(raw_line) > _MAX_WENKORA_SSE_LINE_BYTES:
raise WeKnoraAPIError('WeKnora SSE event exceeds the runtime limit')
line = raw_line.rstrip(b'\r').strip()
if not line.startswith(b'data:'):
continue
try:
data = json.loads(line[5:].strip())
except json.JSONDecodeError:
continue
if isinstance(data, dict):
yield data
if len(buffer) > _MAX_WENKORA_SSE_LINE_BYTES:
raise WeKnoraAPIError('WeKnora SSE event exceeds the runtime limit')
line = bytes(buffer).rstrip(b'\r').strip()
if line.startswith(b'data:'):
try:
data = json.loads(line[5:].strip())
except json.JSONDecodeError:
return
if isinstance(data, dict):
yield data
class AsyncWeKnoraClient:
"""WeKnora API 客户端"""
@@ -39,19 +89,19 @@ class AsyncWeKnoraClient:
if description:
payload['description'] = description
response = await client.post(
async with client.stream(
'POST',
'/sessions',
headers={
'X-API-Key': self.api_key,
'Content-Type': 'application/json',
},
json=payload,
)
if response.status_code not in (200, 201):
raise WeKnoraAPIError(f'{response.status_code} {response.text}')
data = response.json()
) as response:
body = await _read_limited_response(response)
if response.status_code not in (200, 201):
raise WeKnoraAPIError(f'{response.status_code} {body.decode("utf-8", errors="replace")}')
data = json.loads(body)
return data['data']['id']
async def agent_chat(
@@ -107,20 +157,13 @@ class AsyncWeKnoraClient:
},
json=payload,
) as r:
async for chunk in r.aiter_lines():
if r.status_code != 200:
raise WeKnoraAPIError(f'{r.status_code} {chunk}')
if chunk.strip() == '':
continue
if chunk.startswith('data:'):
try:
data = json.loads(chunk[5:].strip())
except json.JSONDecodeError:
continue
yield data
# 收到 error 事件后主动结束流，避免上层未 raise 时持续等待
if data.get('response_type') == 'error':
return
if r.status_code != 200:
body = await _read_limited_response(r)
raise WeKnoraAPIError(f'{r.status_code} {body.decode("utf-8", errors="replace")}')
async for data in _iter_sse_json(r):
yield data
if data.get('response_type') == 'error':
return
async def knowledge_chat(
self,
@@ -164,17 +207,10 @@ class AsyncWeKnoraClient:
},
json=payload,
) as r:
async for chunk in r.aiter_lines():
if r.status_code != 200:
raise WeKnoraAPIError(f'{r.status_code} {chunk}')
if chunk.strip() == '':
continue
if chunk.startswith('data:'):
try:
data = json.loads(chunk[5:].strip())
except json.JSONDecodeError:
continue
yield data
# 收到 error 事件后主动结束流，避免上层未 raise 时持续等待
if data.get('response_type') == 'error':
return
if r.status_code != 200:
body = await _read_limited_response(r)
raise WeKnoraAPIError(f'{r.status_code} {body.decode("utf-8", errors="replace")}')
async for data in _iter_sse_json(r):
yield data
if data.get('response_type') == 'error':
return
src/langbot/pkg/api/http/controller/group.py
+34 -10
View File
@@ -10,9 +10,11 @@ import uuid
from quart.typing import RouteCallable
from ....utils import constants
from ....utils import bounded_executor
from ....workspace.collaboration import MembershipPermissionError, WorkspaceCollaborationError
from ....workspace.errors import WorkspaceNotFoundError
from ....cloud.entitlements import EntitlementUnavailableError
from ....core.errors import TaskCapacityError
from ..authz import AuthorizationError, Permission, permissions_for_role, require_permission
from ..context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext
@@ -185,16 +187,27 @@ class RouterGroup(abc.ABC):
try:
if request_context is not None:
persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
tenant_scope_descriptor = getattr(type(persistence_mgr), 'tenant_scope', None)
if callable(tenant_scope_descriptor):
# Authorization discovery is complete. Carry the
# trusted Workspace identity across the handler, but
# do not reserve a database connection while it waits
# on providers, runtimes, uploads, or streamed clients.
# Services that need atomic writes open a short UoW.
async with persistence_mgr.tenant_scope(request_context.workspace_uuid):
return await f(*args, **kwargs)
with bounded_executor.blocking_work_scope(request_context.workspace_uuid):
persistence_mgr = getattr(
self.ap,
'persistence_mgr',
None,
)
tenant_scope_descriptor = getattr(
type(persistence_mgr),
'tenant_scope',
None,
)
if callable(tenant_scope_descriptor):
# Authorization discovery is complete. Carry
# the trusted Workspace identity across the
# handler, but do not reserve a database
# connection while it waits on providers,
# runtimes, uploads, or streamed clients.
# Services that need atomic writes open a UoW.
async with persistence_mgr.tenant_scope(request_context.workspace_uuid):
return await f(*args, **kwargs)
return await f(*args, **kwargs)
return await f(*args, **kwargs)
except Exception as e: # 自动 500
@@ -206,6 +219,17 @@ class RouterGroup(abc.ABC):
return self.http_status(403, e.code, str(e))
if isinstance(e, WorkspaceCollaborationError):
return self.http_status(400, e.code, str(e))
if isinstance(e, TaskCapacityError):
return self.http_status(429, 'task_capacity_exceeded', str(e))
if isinstance(
e,
bounded_executor.BlockingWorkCapacityError,
):
return self.http_status(
429,
'blocking_work_capacity_exceeded',
str(e),
)
request_id = self.request_id()
logger = getattr(self.ap, 'logger', self.quart_app.logger)
logger.error(
src/langbot/pkg/api/http/controller/groups/knowledge/migration.py
+8 -2
View File
@@ -11,6 +11,7 @@ from ....context import ExecutionContext, RequestContext
from ......core import taskmgr
from ......entity.persistence import metadata as persistence_metadata
from ......workspace.errors import WorkspaceError, WorkspaceNotFoundError
from ......utils import httpclient
from langbot_plugin.runtime.plugin.mgr import PluginInstallSource
LANGRAG_PLUGIN_AUTHOR = 'langbot-team'
@@ -162,10 +163,15 @@ class KnowledgeMigrationRouterGroup(group.RouterGroup):
self.ap.logger.info(f'RAG migration: installing plugin {plugin_id} from marketplace...')
task_context.trace(f'Installing plugin {plugin_id} from marketplace...')
async with httpx.AsyncClient(trust_env=True, timeout=15) as client:
async with httpx.AsyncClient(
trust_env=True,
timeout=15,
event_hooks=httpclient.httpx_response_limit_hooks(),
) as client:
resp = await client.get(f'{space_url}/api/v1/marketplace/plugins/{p_author}/{p_name}')
resp.raise_for_status()
p_data = resp.json().get('data', {}).get('plugin', {})
response_data = await httpclient.parse_json_response(resp)
p_data = response_data.get('data', {}).get('plugin', {})
p_version = p_data.get('latest_version')
if not p_version:
raise Exception(f'Could not determine latest version for {plugin_id}')
src/langbot/pkg/api/http/controller/groups/pipelines/embed.py
+39 -10
View File
@@ -20,8 +20,9 @@ import httpx
import quart
from ... import group
from ......utils import paths
from ......utils import httpclient, paths
from ......platform.sources.websocket_manager import WebSocketScope, is_valid_session_id, ws_connection_manager
from .websocket_chat import create_scoped_duplex_tasks, wait_for_duplex_tasks
logger = logging.getLogger(__name__)
_AUTH_TIMEOUT_SECONDS = 10.0
@@ -103,7 +104,7 @@ class EmbedRouterGroup(group.RouterGroup):
"""Require the embed session token as the first WebSocket frame."""
raw_message = await asyncio.wait_for(quart.websocket.receive(), timeout=_AUTH_TIMEOUT_SECONDS)
payload = json.loads(raw_message)
payload = await asyncio.to_thread(json.loads, raw_message)
if not isinstance(payload, dict) or payload.get('type') != 'authenticate':
raise ValueError('Authentication is required')
token = str(payload.get('token') or '')
@@ -160,12 +161,12 @@ class EmbedRouterGroup(group.RouterGroup):
ts = time.time()
return self.success(data={'token': f'{ts}.dummy'})
async with httpx.AsyncClient() as client:
async with httpx.AsyncClient(event_hooks=httpclient.httpx_response_limit_hooks()) as client:
resp = await client.post(
'https://challenges.cloudflare.com/turnstile/v0/siteverify',
data={'secret': secret, 'response': token},
)
result = resp.json()
result = await httpclient.parse_json_response(resp)
if not result.get('success'):
return self.http_status(403, -1, 'Turnstile verification failed')
@@ -371,6 +372,21 @@ class EmbedRouterGroup(group.RouterGroup):
session_type=session_type,
session_id=session_id,
metadata={'user_agent': quart.websocket.headers.get('User-Agent', '')},
send_queue_size=(
self.ap.instance_config.data.get('system', {})
.get('websocket_retention', {})
.get('send_queue_size', 100)
),
max_connections=(
self.ap.instance_config.data.get('system', {})
.get('websocket_retention', {})
.get('max_connections', 1024)
),
max_connections_per_workspace=(
self.ap.instance_config.data.get('system', {})
.get('websocket_retention', {})
.get('max_connections_per_workspace', 32)
),
)
await quart.websocket.send(
@@ -390,13 +406,19 @@ class EmbedRouterGroup(group.RouterGroup):
f'(bot={bot_uuid}, pipeline={pipeline_uuid}, session_type={session_type})'
)
receive_task = asyncio.create_task(
self._handle_receive(connection, websocket_adapter, runtime_bot, pipeline_uuid)
receive_task, send_task = create_scoped_duplex_tasks(
self._handle_receive(
connection,
websocket_adapter,
runtime_bot,
pipeline_uuid,
),
self._handle_send(connection),
runtime_bot.execution_context.workspace_uuid,
)
send_task = asyncio.create_task(self._handle_send(connection))
try:
await asyncio.gather(receive_task, send_task)
await wait_for_duplex_tasks(receive_task, send_task)
except Exception as e:
logger.error(f'Embed WebSocket task error: {e}')
finally:
@@ -418,7 +440,7 @@ class EmbedRouterGroup(group.RouterGroup):
await ws_connection_manager.update_activity(connection.connection_id)
try:
data = json.loads(message)
data = await asyncio.to_thread(json.loads, message)
message_type = data.get('type', 'message')
if message_type == 'ping':
@@ -442,13 +464,20 @@ class EmbedRouterGroup(group.RouterGroup):
logger.error(f'Embed receive error: {e}', exc_info=True)
finally:
connection.is_active = False
try:
connection.send_queue.put_nowait(None)
except asyncio.QueueFull:
pass
async def _handle_send(self, connection):
try:
while connection.is_active or not connection.send_queue.empty():
try:
message = await asyncio.wait_for(connection.send_queue.get(), timeout=1.0)
await quart.websocket.send(json.dumps(message))
if message is None:
break
encoded = await asyncio.to_thread(json.dumps, message)
await quart.websocket.send(encoded)
except asyncio.TimeoutError:
continue
except Exception as e:
src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py
+86 -7
View File
@@ -6,6 +6,7 @@ import asyncio
import datetime
import json
import logging
import typing
import uuid
import quart
@@ -15,9 +16,64 @@ from ....context import PrincipalContext, PrincipalType, RequestContext, Workspa
from ... import group
from ......core.task_boundary import run_in_workspace_uow
from ......platform.sources.websocket_manager import WebSocketScope, ws_connection_manager
from ......utils import bounded_executor
logger = logging.getLogger(__name__)
_AUTH_TIMEOUT_SECONDS = 10.0
_DUPLEX_DRAIN_TIMEOUT_SECONDS = 0.25
def create_scoped_duplex_tasks(
receive_coro: typing.Coroutine[typing.Any, typing.Any, None],
send_coro: typing.Coroutine[typing.Any, typing.Any, None],
workspace_uuid: str,
) -> tuple[asyncio.Task[None], asyncio.Task[None]]:
"""Create both socket directions under one trusted Workspace budget."""
return (
asyncio.create_task(
bounded_executor.run_in_blocking_work_scope(
receive_coro,
workspace_uuid,
)
),
asyncio.create_task(
bounded_executor.run_in_blocking_work_scope(
send_coro,
workspace_uuid,
)
),
)
async def wait_for_duplex_tasks(
receive_task: asyncio.Task,
send_task: asyncio.Task,
) -> None:
"""Stop the peer direction as soon as either socket task terminates."""
try:
done, _ = await asyncio.wait(
{receive_task, send_task},
return_when=asyncio.FIRST_COMPLETED,
)
# A receive task may enqueue a terminal authorization/error frame and
# then finish. Give the sender a short deterministic drain window
# instead of cancelling it before that frame reaches the client.
if receive_task in done and not send_task.done():
await asyncio.wait(
{send_task},
timeout=_DUPLEX_DRAIN_TIMEOUT_SECONDS,
)
finally:
for task in (receive_task, send_task):
if not task.done():
task.cancel()
await asyncio.gather(
receive_task,
send_task,
return_exceptions=True,
)
@group.group_class('websocket_chat', '/api/v1/pipelines/<pipeline_uuid>/ws')
@@ -32,7 +88,7 @@ class WebSocketChatRouterGroup(group.RouterGroup):
"""
raw_message = await asyncio.wait_for(quart.websocket.receive(), timeout=_AUTH_TIMEOUT_SECONDS)
payload = json.loads(raw_message)
payload = await asyncio.to_thread(json.loads, raw_message)
if not isinstance(payload, dict) or payload.get('type') != 'authenticate':
raise ValueError('Authentication is required')
@@ -155,6 +211,21 @@ class WebSocketChatRouterGroup(group.RouterGroup):
pipeline_uuid=pipeline_uuid,
session_type=session_type,
metadata={'user_agent': quart.websocket.headers.get('User-Agent', '')},
send_queue_size=(
self.ap.instance_config.data.get('system', {})
.get('websocket_retention', {})
.get('send_queue_size', 100)
),
max_connections=(
self.ap.instance_config.data.get('system', {})
.get('websocket_retention', {})
.get('max_connections', 1024)
),
max_connections_per_workspace=(
self.ap.instance_config.data.get('system', {})
.get('websocket_retention', {})
.get('max_connections_per_workspace', 32)
),
)
await quart.websocket.send(
@@ -175,17 +246,18 @@ class WebSocketChatRouterGroup(group.RouterGroup):
f'session_type={session_type})'
)
receive_task = asyncio.create_task(
receive_task, send_task = create_scoped_duplex_tasks(
self._handle_receive(
connection,
websocket_adapter,
request_context,
token,
)
),
self._handle_send(connection),
request_context.workspace_uuid,
)
send_task = asyncio.create_task(self._handle_send(connection))
try:
await asyncio.gather(receive_task, send_task)
await wait_for_duplex_tasks(receive_task, send_task)
except Exception as exc:
logger.error(f'WebSocket task execution error: {exc}')
finally:
@@ -310,7 +382,7 @@ class WebSocketChatRouterGroup(group.RouterGroup):
await ws_connection_manager.update_activity(connection.connection_id)
try:
data = json.loads(message)
data = await asyncio.to_thread(json.loads, message)
message_type = data.get('type', 'message')
if message_type == 'ping':
await connection.send_queue.put(
@@ -334,13 +406,20 @@ class WebSocketChatRouterGroup(group.RouterGroup):
logger.error('Dashboard WebSocket receive error', exc_info=True)
finally:
connection.is_active = False
try:
connection.send_queue.put_nowait(None)
except asyncio.QueueFull:
pass
async def _handle_send(self, connection):
try:
while connection.is_active or not connection.send_queue.empty():
try:
message = await asyncio.wait_for(connection.send_queue.get(), timeout=1.0)
await quart.websocket.send(json.dumps(message))
if message is None:
break
encoded = await asyncio.to_thread(json.dumps, message)
await quart.websocket.send(encoded)
except asyncio.TimeoutError:
continue
except Exception:
src/langbot/pkg/api/http/controller/groups/platform/adapters.py
+122 -13
View File
@@ -6,7 +6,8 @@ import quart
from langbot.pkg.api.http.authz import Permission
from langbot.pkg.api.http.context import RequestContext
from langbot.pkg.utils import importutil
from langbot.pkg.core.errors import TaskCapacityError
from langbot.pkg.utils import httpclient, importutil
from ... import group
@@ -71,6 +72,64 @@ def _pop_owned_session(
return sessions.pop(session_id, None)
_MAX_ADAPTER_SESSIONS = 100
_MAX_ADAPTER_SESSIONS_PER_WORKSPACE = 10
def _start_adapter_session_task(
ap,
coro,
*,
adapter: str,
session_id: str,
request_context: RequestContext,
) -> asyncio.Task | None:
"""Attach one credential exchange to tenant admission and app shutdown."""
try:
wrapper = ap.task_mgr.create_user_task(
coro,
kind='platform-adapter-credential-exchange',
name=f'{adapter}-credential-{session_id}',
label=f'{adapter} credential exchange',
instance_uuid=request_context.instance_uuid,
workspace_uuid=request_context.workspace_uuid,
placement_generation=request_context.placement_generation,
)
except TaskCapacityError:
coro.close()
return None
return wrapper.task
def _make_room_for_session(
sessions: dict[str, dict],
request_context: RequestContext,
) -> None:
"""Bound credential-exchange sessions globally and per workspace."""
workspace_uuid = request_context.workspace_uuid
owned = [
(session_id, session)
for session_id, session in sessions.items()
if getattr(session.get('scope'), 'workspace_uuid', None) == workspace_uuid
]
evict_workspace_session = len(owned) >= _MAX_ADAPTER_SESSIONS_PER_WORKSPACE
evict_global_session = len(sessions) >= _MAX_ADAPTER_SESSIONS
if not evict_workspace_session and not evict_global_session:
return
candidates = owned if evict_workspace_session else list(sessions.items())
session_id, _ = min(
candidates,
key=lambda item: float(item[1].get('created_at', 0.0)),
)
session = sessions.pop(session_id, None)
task = session.get('task') if session is not None else None
if task is not None and not task.done():
task.cancel()
def _decrypt_qqofficial_secret(encrypted_b64: str, key: bytes) -> str:
"""Decrypt the AppSecret returned by the QQ Official QR binding endpoint.
@@ -173,6 +232,7 @@ class AdaptersRouterGroup(group.RouterGroup):
'created_at': time.time(),
}
_bind_session_scope(session, request_context)
_make_room_for_session(_create_app_sessions, request_context)
_create_app_sessions[session_id] = session
def on_qr_code(info):
@@ -204,7 +264,16 @@ class AdaptersRouterGroup(group.RouterGroup):
session['status'] = 'error'
session['error'] = str(e)
task = asyncio.create_task(run_registration())
task = _start_adapter_session_task(
self.ap,
run_registration(),
adapter='lark',
session_id=session_id,
request_context=request_context,
)
if task is None:
_create_app_sessions.pop(session_id, None)
return self.http_status(429, -1, 'Too many active credential exchanges')
session['task'] = task
# Wait for QR code to be ready (max 10 seconds)
@@ -308,6 +377,7 @@ class AdaptersRouterGroup(group.RouterGroup):
'created_at': time.time(),
}
_bind_session_scope(session, request_context)
_make_room_for_session(_weixin_login_sessions, request_context)
_weixin_login_sessions[session_id] = session
client = OpenClawWeixinClient(
@@ -346,7 +416,16 @@ class AdaptersRouterGroup(group.RouterGroup):
finally:
await client.close()
task = asyncio.create_task(run_login())
task = _start_adapter_session_task(
self.ap,
run_login(),
adapter='weixin',
session_id=session_id,
request_context=request_context,
)
if task is None:
_weixin_login_sessions.pop(session_id, None)
return self.http_status(429, -1, 'Too many active credential exchanges')
session['task'] = task
# Wait for QR code to be ready (max 10 seconds)
@@ -459,6 +538,7 @@ class AdaptersRouterGroup(group.RouterGroup):
'interval': 5,
}
_bind_session_scope(session, request_context)
_make_room_for_session(_dingtalk_sessions, request_context)
_dingtalk_sessions[session_id] = session
async def run_device_flow():
@@ -471,7 +551,7 @@ class AdaptersRouterGroup(group.RouterGroup):
json={'source': 'langbot'},
) as resp:
try:
data = await resp.json()
data = await httpclient.read_json_limited(resp)
except (aiohttp.ContentTypeError, ValueError):
session['status'] = 'error'
session['error'] = 'Invalid response from DingTalk service'
@@ -488,7 +568,7 @@ class AdaptersRouterGroup(group.RouterGroup):
json={'nonce': nonce},
) as resp:
try:
data = await resp.json()
data = await httpclient.read_json_limited(resp)
except (aiohttp.ContentTypeError, ValueError):
session['status'] = 'error'
session['error'] = 'Invalid response from DingTalk service'
@@ -519,7 +599,7 @@ class AdaptersRouterGroup(group.RouterGroup):
json={'device_code': device_code},
) as poll_resp:
try:
poll_data = await poll_resp.json()
poll_data = await httpclient.read_json_limited(poll_resp)
except (aiohttp.ContentTypeError, ValueError):
continue
@@ -555,7 +635,16 @@ class AdaptersRouterGroup(group.RouterGroup):
session['status'] = 'error'
session['error'] = str(e)
task = asyncio.create_task(run_device_flow())
task = _start_adapter_session_task(
self.ap,
run_device_flow(),
adapter='dingtalk',
session_id=session_id,
request_context=request_context,
)
if task is None:
_dingtalk_sessions.pop(session_id, None)
return self.http_status(429, -1, 'Too many active credential exchanges')
session['task'] = task
# Wait for QR code to be ready (max 10 seconds)
@@ -665,6 +754,7 @@ class AdaptersRouterGroup(group.RouterGroup):
'task': None,
}
_bind_session_scope(session, request_context)
_make_room_for_session(_wecombot_sessions, request_context)
_wecombot_sessions[session_id] = session
async def run_qr_flow():
@@ -676,7 +766,7 @@ class AdaptersRouterGroup(group.RouterGroup):
f'{WECOM_QC_GENERATE_URL}?source=langbot&plat=0',
) as resp:
try:
data = await resp.json()
data = await httpclient.read_json_limited(resp)
except (aiohttp.ContentTypeError, ValueError):
session['status'] = 'error'
session['error'] = 'Invalid response from WeCom service'
@@ -703,7 +793,7 @@ class AdaptersRouterGroup(group.RouterGroup):
f'{WECOM_QC_QUERY_URL}?scode={scode}',
) as poll_resp:
try:
poll_data = await poll_resp.json()
poll_data = await httpclient.read_json_limited(poll_resp)
except (aiohttp.ContentTypeError, ValueError):
continue
@@ -730,7 +820,16 @@ class AdaptersRouterGroup(group.RouterGroup):
session['status'] = 'error'
session['error'] = str(e)
task = asyncio.create_task(run_qr_flow())
task = _start_adapter_session_task(
self.ap,
run_qr_flow(),
adapter='wecombot',
session_id=session_id,
request_context=request_context,
)
if task is None:
_wecombot_sessions.pop(session_id, None)
return self.http_status(429, -1, 'Too many active credential exchanges')
session['task'] = task
# Wait for QR code to be ready (max 10 seconds)
@@ -852,6 +951,7 @@ class AdaptersRouterGroup(group.RouterGroup):
'interval': 2,
}
_bind_session_scope(session, request_context)
_make_room_for_session(_qqofficial_sessions, request_context)
_qqofficial_sessions[session_id] = session
async def run_qr_binding():
@@ -865,7 +965,7 @@ class AdaptersRouterGroup(group.RouterGroup):
headers={'Accept': 'application/json'},
) as resp:
try:
data = await resp.json(content_type=None)
data = await httpclient.read_json_limited(resp)
except (aiohttp.ContentTypeError, ValueError):
session['status'] = 'error'
session['error'] = 'Invalid response from QQ bind service'
@@ -903,7 +1003,7 @@ class AdaptersRouterGroup(group.RouterGroup):
headers={'Accept': 'application/json'},
) as poll_resp:
try:
poll_data = await poll_resp.json(content_type=None)
poll_data = await httpclient.read_json_limited(poll_resp)
except (aiohttp.ContentTypeError, ValueError):
continue
@@ -956,7 +1056,16 @@ class AdaptersRouterGroup(group.RouterGroup):
session['status'] = 'error'
session['error'] = str(e)
task = asyncio.create_task(run_qr_binding())
task = _start_adapter_session_task(
self.ap,
run_qr_binding(),
adapter='qqofficial',
session_id=session_id,
request_context=request_context,
)
if task is None:
_qqofficial_sessions.pop(session_id, None)
return self.http_status(429, -1, 'Too many active credential exchanges')
session['task'] = task
# Wait up to 10s for the QR URL to be ready before responding.
src/langbot/pkg/api/http/controller/groups/plugins.py
+34 -50
View File
@@ -1,16 +1,15 @@
from __future__ import annotations
import asyncio
import base64
import collections.abc
import copy
import io
import quart
import re
import httpx
import uuid
import os
import zipfile
import yaml
from urllib.parse import urlparse
import posixpath
import sqlalchemy
@@ -23,6 +22,8 @@ from ...context import ExecutionContext, RequestContext
from .. import group
from .....workspace.errors import WorkspaceNotFoundError
from .....plugin.github import validate_github_plugin_install_info
from .....plugin.archive import inspect_plugin_archive_metadata
from .....utils import httpclient
from langbot_plugin.runtime.plugin.mgr import PluginInstallSource
@@ -548,7 +549,7 @@ class PluginsRouterGroup(group.RouterGroup):
icon_base64 = icon_data['plugin_icon_base64']
mime_type = icon_data['mime_type']
icon_data = base64.b64decode(icon_base64)
icon_data = await asyncio.to_thread(base64.b64decode, icon_base64)
return quart.Response(icon_data, mimetype=mime_type)
@@ -566,7 +567,10 @@ class PluginsRouterGroup(group.RouterGroup):
asset_data = await self.ap.plugin_connector.get_plugin_assets(author, plugin_name, asset_path)
if not asset_data.get('asset_base64'):
return quart.Response('Asset not found', status=404)
asset_bytes = base64.b64decode(asset_data['asset_base64'])
asset_bytes = await asyncio.to_thread(
base64.b64decode,
asset_data['asset_base64'],
)
mime_type = asset_data['mime_type']
resp = quart.Response(asset_bytes, mimetype=mime_type)
# CSP for HTML pages served to sandboxed iframes (opaque origin).
@@ -662,10 +666,11 @@ class PluginsRouterGroup(group.RouterGroup):
trust_env=True,
follow_redirects=True,
timeout=10,
event_hooks=httpclient.httpx_response_limit_hooks(),
) as client:
response = await client.get(url)
response.raise_for_status()
releases = response.json()
releases = await httpclient.parse_json_response(response)
# Format releases data for frontend
formatted_releases = []
@@ -716,12 +721,13 @@ class PluginsRouterGroup(group.RouterGroup):
trust_env=True,
follow_redirects=True,
timeout=10,
event_hooks=httpclient.httpx_response_limit_hooks(),
) as client:
response = await client.get(
url,
)
response.raise_for_status()
release = response.json()
release = await httpclient.parse_json_response(response)
# Format assets data for frontend
formatted_assets = []
@@ -902,51 +908,29 @@ class PluginsRouterGroup(group.RouterGroup):
file_bytes = file.read()
try:
with zipfile.ZipFile(io.BytesIO(file_bytes)) as zf:
names = [name for name in zf.namelist() if not name.endswith('/')]
manifest_name = next(
(
name
for name in names
if name.replace('\\', '/').strip('/').lower() in ('manifest.yaml', 'manifest.yml')
),
None,
)
if manifest_name is None:
return self.http_status(400, -1, 'manifest.yaml is required')
manifest, requirements, names = await asyncio.to_thread(
inspect_plugin_archive_metadata,
file_bytes,
)
spec = manifest.get('spec') or {}
components = spec.get('components') or {}
component_counts = self._count_plugin_components(components, names)
component_types = list(component_counts.keys())
manifest = yaml.safe_load(zf.read(manifest_name).decode('utf-8')) or {}
requirements: list[str] = []
requirements_name = next(
(name for name in names if name.replace('\\', '/').strip('/').lower() == 'requirements.txt'),
None,
)
if requirements_name is not None:
requirements = [
line.strip()
for line in zf.read(requirements_name).decode('utf-8', errors='ignore').splitlines()
if line.strip() and not line.strip().startswith('#')
]
spec = manifest.get('spec') or {}
components = spec.get('components') or {}
component_counts = self._count_plugin_components(components, names)
component_types = list(component_counts.keys())
return self.success(
data={
'filename': file.filename or 'local plugin',
'size': len(file_bytes),
'manifest': manifest,
'metadata': manifest.get('metadata') or {},
'component_types': component_types,
'component_counts': component_counts,
'requirements': requirements,
'file_count': len(names),
}
)
except zipfile.BadZipFile:
return self.http_status(400, -1, 'invalid .lbpkg file')
return self.success(
data={
'filename': file.filename or 'local plugin',
'size': len(file_bytes),
'manifest': manifest,
'metadata': manifest.get('metadata') or {},
'component_types': component_types,
'component_counts': component_counts,
'requirements': requirements,
'file_count': len(names),
}
)
except (zipfile.BadZipFile, ValueError) as exc:
return self.http_status(400, -1, str(exc) or 'invalid .lbpkg file')
except Exception:
raise
src/langbot/pkg/api/http/controller/groups/survey.py
+9 -1
View File
@@ -1,3 +1,4 @@
import asyncio
import base64
import quart
@@ -59,7 +60,14 @@ class SurveyRouterGroup(group.RouterGroup):
continue
try:
payload = data_url.split(',', 1)[1]
if len(base64.b64decode(payload, validate=True)) > 1024 * 1024:
if len(payload) > 4 * ((1024 * 1024 + 2) // 3) + 4:
return self.fail(5, 'attachment too large')
decoded = await asyncio.to_thread(
base64.b64decode,
payload,
validate=True,
)
if len(decoded) > 1024 * 1024:
return self.fail(5, 'attachment too large')
except Exception:
return self.fail(5, 'attachment too large')
src/langbot/pkg/api/http/controller/groups/webhooks.py
+17 -9
View File
@@ -4,6 +4,7 @@ import quart
import traceback
from .. import group
from .....utils import bounded_executor
@group.group_class('webhooks', '/bots')
@@ -55,19 +56,26 @@ class WebhookRouterGroup(group.RouterGroup):
request=quart.request,
)
persistence_mgr = self.ap.persistence_mgr
cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
tenant_scope = getattr(persistence_mgr, 'tenant_scope', None)
if not callable(tenant_scope):
raise RuntimeError('Cloud webhook dispatch requires an explicit tenant scope')
async with tenant_scope(runtime_bot.workspace_uuid):
with bounded_executor.blocking_work_scope(runtime_bot.workspace_uuid):
persistence_mgr = self.ap.persistence_mgr
cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
tenant_scope = getattr(persistence_mgr, 'tenant_scope', None)
if not callable(tenant_scope):
raise RuntimeError('Cloud webhook dispatch requires an explicit tenant scope')
async with tenant_scope(runtime_bot.workspace_uuid):
response = await dispatch()
else:
response = await dispatch()
else:
response = await dispatch()
return response
except bounded_executor.BlockingWorkCapacityError as exc:
return self.http_status(
429,
'blocking_work_capacity_exceeded',
str(exc),
)
except Exception:
request_id = self.request_id()
self.ap.logger.error(
src/langbot/pkg/api/http/controller/main.py
+44 -1
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import os
import typing
import quart
import quart_cors
@@ -27,6 +28,37 @@ importutil.import_modules_in_pkg(groups_knowledge)
importutil.import_modules_in_pkg(groups_resources)
class BoundedJSONRequest(quart.Request):
"""Parse bounded HTTP JSON bodies outside the shared event loop."""
async def get_json(
self,
force: bool = False,
silent: bool = False,
cache: bool = True,
) -> typing.Any:
# Keep Quart's cache and error semantics, changing only where the
# potentially 10 MiB JSON decoder runs. The RouterGroup establishes a
# trusted Workspace blocking-work scope before calling route handlers.
if cache and self._cached_json[silent] is not Ellipsis:
return self._cached_json[silent]
if not (force or self.is_json):
return None
data = await self.get_data(cache=cache, as_text=False)
try:
result = await asyncio.to_thread(self.json_module.loads, data)
except ValueError as error:
if silent:
result = None
else:
result = self.on_json_loading_failed(error)
if cache:
self._cached_json[silent] = result
return result
class HTTPController:
ap: app.Application
@@ -35,6 +67,7 @@ class HTTPController:
def __init__(self, ap: app.Application) -> None:
self.ap = ap
self.quart_app = quart.Quart(__name__)
self.quart_app.request_class = BoundedJSONRequest
quart_cors.cors(self.quart_app, allow_origin='*')
# Set maximum content length to prevent large file uploads
@@ -103,6 +136,7 @@ class HTTPController:
config.accesslog = '-'
config.bind = [f'{host}:{port}']
config.errorlog = config.accesslog
config.websocket_max_message_size = group.MAX_FILE_SIZE
asgi_app = self.quart_app
if self.mcp_mount is not None:
@@ -113,7 +147,16 @@ class HTTPController:
async def register_routes(self) -> None:
@self.quart_app.route('/healthz')
async def healthz():
return {'code': 0, 'msg': 'ok'}
get_resource_stats = getattr(
self.ap,
'get_runtime_resource_stats',
None,
)
return {
'code': 0,
'msg': 'ok',
'resources': (get_resource_stats() if callable(get_resource_stats) else {}),
}
for g in group.preregistered_groups:
ginst = g(self.ap, self.quart_app)
src/langbot/pkg/api/http/service/maintenance.py
+74 -26
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import datetime
import functools
import os
@@ -69,7 +70,10 @@ class MaintenanceService:
return {
'uploaded_files': await self._cleanup_expired_uploaded_files(context, upload_retention_days),
'log_files': self._cleanup_expired_log_files(log_retention_days)
'log_files': await asyncio.to_thread(
self._cleanup_expired_log_files,
log_retention_days,
)
if await self._is_oss_singleton(context)
else 0,
}
@@ -108,22 +112,19 @@ class MaintenanceService:
scoped_storage_path = Path('data/storage') / self.ap.storage_mgr.scoped_prefix(context)
roots = [('storage', scoped_storage_path)]
sections = []
for key, path in roots:
sections.append(
{
'key': key,
'path': str(path) if path else '',
'exists': path.exists() if path else False,
'size_bytes': self._path_size(path) if path else 0,
'file_count': self._file_count(path) if path else 0,
}
)
sections = await asyncio.to_thread(self._collect_sections, roots)
monitoring_counts = await self._monitoring_counts(context)
binary_storage = await self._binary_storage_stats(context)
upload_candidates = await self._expired_uploaded_candidates(context, upload_retention_days)
log_candidates = self._expired_log_candidates(log_retention_days) if is_oss_singleton else []
log_candidates = (
await asyncio.to_thread(
self._expired_log_candidates,
log_retention_days,
)
if is_oss_singleton
else []
)
return {
'generated_at': datetime.datetime.now(datetime.timezone.utc).isoformat(),
@@ -144,6 +145,23 @@ class MaintenanceService:
'tasks': self.ap.task_mgr.get_stats() if is_oss_singleton and self.ap.task_mgr else {},
}
def _collect_sections(
self,
roots: list[tuple[str, Path | None]],
) -> list[dict[str, Any]]:
sections = []
for key, path in roots:
sections.append(
{
'key': key,
'path': str(path) if path else '',
'exists': path.exists() if path else False,
'size_bytes': self._path_size(path) if path else 0,
'file_count': self._file_count(path) if path else 0,
}
)
return sections
async def _is_oss_singleton(self, context: TenantContext) -> bool:
try:
await self.ap.workspace_service.get_local_execution_binding(
@@ -162,21 +180,16 @@ class MaintenanceService:
provider = self.ap.storage_mgr.storage_provider
provider_name = provider.__class__.__name__
if provider_name == 'LocalStorageProvider':
candidates = self._expired_local_upload_candidates(
candidates = await asyncio.to_thread(
self._expired_local_upload_candidates,
context,
retention_days,
include_paths=True,
True,
)
return await asyncio.to_thread(
self._delete_local_candidates,
candidates,
)
deleted = 0
for item in candidates:
try:
os.remove(item['path'])
deleted += 1
except FileNotFoundError:
pass
except Exception as e:
self.ap.logger.warning(f'Failed to delete expired uploaded file {item["key"]}: {e}')
return deleted
if provider_name == 'S3StorageProvider':
return await self._cleanup_expired_s3_uploaded_files(context, retention_days)
@@ -190,7 +203,11 @@ class MaintenanceService:
) -> list[dict[str, Any]]:
provider_name = self.ap.storage_mgr.storage_provider.__class__.__name__
if provider_name == 'LocalStorageProvider':
return self._expired_local_upload_candidates(context, retention_days)
return await asyncio.to_thread(
self._expired_local_upload_candidates,
context,
retention_days,
)
if provider_name == 'S3StorageProvider':
return await self._expired_s3_upload_candidates(context, retention_days)
return []
@@ -212,6 +229,25 @@ class MaintenanceService:
self,
context: TenantContext,
retention_days: int,
) -> list[dict[str, Any]]:
provider = self.ap.storage_mgr.storage_provider
run_io = getattr(provider, '_run_io', None)
if callable(run_io):
return await run_io(
self._expired_s3_upload_candidates_sync,
context,
retention_days,
)
return await asyncio.to_thread(
self._expired_s3_upload_candidates_sync,
context,
retention_days,
)
def _expired_s3_upload_candidates_sync(
self,
context: TenantContext,
retention_days: int,
) -> list[dict[str, Any]]:
provider = self.ap.storage_mgr.storage_provider
cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=retention_days)
@@ -241,6 +277,18 @@ class MaintenanceService:
return candidates
def _delete_local_candidates(self, candidates: list[dict[str, Any]]) -> int:
deleted = 0
for item in candidates:
try:
os.remove(item['path'])
deleted += 1
except FileNotFoundError:
pass
except Exception as e:
self.ap.logger.warning(f'Failed to delete expired uploaded file {item["key"]}: {e}')
return deleted
def _cleanup_expired_log_files(self, retention_days: int) -> int:
deleted = 0
for item in self._expired_log_candidates(retention_days, include_paths=True):
src/langbot/pkg/api/http/service/mcp.py
+38 -12
View File
@@ -252,8 +252,17 @@ class MCPService:
task = create_detached_task(
self.ap.tool_mgr.mcp_tool_loader.host_mcp_server(execution_context, created),
after_commit_manager=self.ap.persistence_mgr,
workspace_uuid=execution_context.workspace_uuid,
)
self.ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks.append(task)
tracker = getattr(
self.ap.tool_mgr.mcp_tool_loader,
'track_hosted_task',
None,
)
if callable(tracker):
tracker(task, execution_context)
else:
self.ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks.append(task)
return payload['uuid']
async def get_mcp_server_by_uuid(self, context: TenantContext, server_uuid: str) -> dict | None:
@@ -357,8 +366,13 @@ class MCPService:
task = create_detached_task(
loader.host_mcp_server(execution_context, updated),
after_commit_manager=self.ap.persistence_mgr,
workspace_uuid=execution_context.workspace_uuid,
)
loader._hosted_mcp_tasks.append(task)
tracker = getattr(loader, 'track_hosted_task', None)
if callable(tracker):
tracker(task, execution_context)
else:
loader._hosted_mcp_tasks.append(task)
async def delete_mcp_server(self, context: TenantContext, server_uuid: str) -> None:
execution_context = await self._execution_context(context)
@@ -420,6 +434,7 @@ class MCPService:
async def test_mcp_server(self, context: TenantContext, server_name: str, server_data: dict) -> int:
execution_context = await self._execution_context(context)
runtime_mcp_session: RuntimeMCPSession | None = None
test_session: RuntimeMCPSession | None = None
ctx = taskmgr.TaskContext.new()
if server_name != '_':
@@ -468,16 +483,27 @@ class MCPService:
coroutine = _run_and_cleanup()
wrapper = self.ap.task_mgr.create_user_task(
coroutine,
kind='mcp-operation',
name=f'mcp-test-{execution_context.workspace_uuid}-{server_name}',
label=f'Testing MCP server {server_name}',
context=ctx,
instance_uuid=execution_context.instance_uuid,
workspace_uuid=execution_context.workspace_uuid,
placement_generation=execution_context.placement_generation,
)
try:
wrapper = self.ap.task_mgr.create_user_task(
coroutine,
kind='mcp-operation',
name=f'mcp-test-{execution_context.workspace_uuid}-{server_name}',
label=f'Testing MCP server {server_name}',
context=ctx,
instance_uuid=execution_context.instance_uuid,
workspace_uuid=execution_context.workspace_uuid,
placement_generation=execution_context.placement_generation,
)
except taskmgr.TaskCapacityError:
if test_session is not None:
try:
await test_session.shutdown()
except Exception as exc:
self.ap.logger.warning(
f'Failed to tear down rejected transient MCP test session '
f'{test_session.server_name}: {type(exc).__name__}: {exc}'
)
raise
return wrapper.id
async def get_mcp_server_logs(
src/langbot/pkg/api/http/service/skill.py
+14 -1
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import io
import inspect
import os
@@ -13,6 +14,7 @@ import httpx
from ....core import app
from ....skill.utils import parse_frontmatter
from ....utils import httpclient
from ..context import ExecutionContext
from .tenant import TenantContext, require_workspace_uuid
@@ -328,7 +330,11 @@ class SkillService:
await result
async def _download_github_asset(self, asset_url: str) -> bytes:
async with httpx.AsyncClient(follow_redirects=True, timeout=120) as client:
async with httpx.AsyncClient(
follow_redirects=True,
timeout=120,
event_hooks=httpclient.httpx_response_limit_hooks(_MAX_GITHUB_ARCHIVE_BYTES),
) as client:
async with client.stream('GET', asset_url) as resp:
resp.raise_for_status()
content_length = resp.headers.get('content-length')
@@ -352,7 +358,14 @@ class SkillService:
info = self._parse_github_skill_md_url(asset_url, owner=owner, repo=repo)
archive_url = f'https://codeload.github.com/{owner}/{repo}/zip/{quote(info["ref"], safe="/")}'
archive_bytes = await self._download_github_asset(archive_url)
return await asyncio.to_thread(self._build_github_skill_directory_zip, archive_bytes, info)
def _build_github_skill_directory_zip(
self,
archive_bytes: bytes,
info: dict[str, str],
) -> tuple[bytes, str, str]:
"""Validate and repack a GitHub skill archive outside the event loop."""
try:
source_archive = zipfile.ZipFile(io.BytesIO(archive_bytes), 'r')
except zipfile.BadZipFile as exc:
src/langbot/pkg/api/http/service/space.py
+46 -13
View File
@@ -1,5 +1,7 @@
from __future__ import annotations
from collections import OrderedDict
from langbot.pkg.utils import httpclient
import typing
import datetime
@@ -11,6 +13,10 @@ from ....entity.persistence import user
from ....entity.dto.space_model import SpaceModel
_CREDITS_CACHE_TTL_SECONDS = 60
_CREDITS_CACHE_MAX_ENTRIES = 4096
class SpaceService:
"""Service for interacting with LangBot Space API"""
@@ -19,7 +25,24 @@ class SpaceService:
def __init__(self, ap: app.Application) -> None:
self.ap = ap
self._credits_cache = {}
self._credits_cache = OrderedDict()
def _ordered_credits_cache(
self,
) -> OrderedDict[str, tuple[int, float]]:
if not isinstance(self._credits_cache, OrderedDict):
# Preserve compatibility with tests and callers that seed the cache.
self._credits_cache = OrderedDict(self._credits_cache)
return self._credits_cache
def _prune_credits_cache(self, now: float) -> None:
cache = self._ordered_credits_cache()
while cache:
email = next(iter(cache))
_, cached_at = cache[email]
if now - cached_at < _CREDITS_CACHE_TTL_SECONDS:
break
cache.pop(email, None)
def _get_space_config(self) -> typing.Dict[str, str]:
"""Get Space configuration from config file"""
@@ -107,8 +130,9 @@ class SpaceService:
json={'code': code, 'instance_id': constants.instance_id},
) as response:
if response.status != 200:
raise ValueError(f'Failed to exchange OAuth code: {await response.text()}')
data = await response.json()
error = await httpclient.read_text_limited(response)
raise ValueError(f'Failed to exchange OAuth code: {error}')
data = await httpclient.read_json_limited(response)
if data.get('code') != 0:
raise ValueError(f'Failed to exchange OAuth code: {data.get("msg")}')
return data.get('data', {})
@@ -123,8 +147,9 @@ class SpaceService:
f'{space_url}/api/v1/accounts/token/refresh', json={'refresh_token': refresh_token}
) as response:
if response.status != 200:
raise ValueError(f'Failed to refresh token: {await response.text()}')
data = await response.json()
error = await httpclient.read_text_limited(response)
raise ValueError(f'Failed to refresh token: {error}')
data = await httpclient.read_json_limited(response)
if data.get('code') != 0:
raise ValueError(f'Failed to refresh token: {data.get("msg")}')
return data.get('data', {})
@@ -139,8 +164,9 @@ class SpaceService:
f'{space_url}/api/v1/accounts/me', headers={'Authorization': f'Bearer {access_token}'}
) as response:
if response.status != 200:
raise ValueError(f'Failed to get user info: {await response.text()}')
data = await response.json()
error = await httpclient.read_text_limited(response)
raise ValueError(f'Failed to get user info: {error}')
data = await httpclient.read_json_limited(response)
if data.get('code') != 0:
raise ValueError(f'Failed to get user info: {data.get("msg")}')
return data.get('data', {})
@@ -156,11 +182,13 @@ class SpaceService:
async def get_credits(self, user_email: str, force_refresh: bool = False) -> int | None:
"""Get Space credits for user with caching (60s TTL)"""
cache_ttl = 60
now = time.time()
cached_fallback = self._credits_cache.get(user_email)
self._prune_credits_cache(now)
if not force_refresh and user_email in self._credits_cache:
credits, ts = self._credits_cache[user_email]
if time.time() - ts < cache_ttl:
if now - ts < _CREDITS_CACHE_TTL_SECONDS:
return credits
try:
@@ -169,10 +197,14 @@ class SpaceService:
return None
credits = info.get('credits')
if credits is not None:
self._credits_cache[user_email] = (credits, time.time())
cache = self._ordered_credits_cache()
cache.pop(user_email, None)
if len(cache) >= _CREDITS_CACHE_MAX_ENTRIES:
cache.popitem(last=False)
cache[user_email] = (credits, time.time())
return credits
except Exception:
return self._credits_cache.get(user_email, (None, 0))[0]
return cached_fallback[0] if cached_fallback is not None else None
async def get_models(self) -> typing.List[SpaceModel]:
"""Get models from Space"""
@@ -183,8 +215,9 @@ class SpaceService:
session = httpclient.get_session()
async with session.get(f'{space_url}/api/v1/models', params={'page_size': 100}) as response:
if response.status != 200:
raise ValueError(f'Failed to get models: {await response.text()}')
data = await response.json()
error = await httpclient.read_text_limited(response)
raise ValueError(f'Failed to get models: {error}')
data = await httpclient.read_json_limited(response)
if data.get('code') != 0:
raise ValueError(f'Failed to get models: {data.get("msg")}')
models_data = data.get('data', {}).get('models', [])
src/langbot/pkg/api/http/service/user.py
+60 -7
View File
@@ -7,6 +7,7 @@ import datetime
import typing
import asyncio
import dataclasses
import heapq
import hashlib
import secrets
import time
@@ -19,11 +20,17 @@ from ....entity.persistence.workspace import MembershipRole, MembershipStatus, W
from ....utils import constants
from ....entity.errors import account as account_errors
from ....workspace.collaboration import normalize_email
from ....utils import bounded_executor
if typing.TYPE_CHECKING:
from ....core.app import Application
_SPACE_OAUTH_STATE_MAX_ENTRIES = 4096
_SPACE_OAUTH_STATE_HEAP_COMPACT_FLOOR = 64
_SPACE_OAUTH_STATE_HEAP_MAX_MULTIPLIER = 4
class AccountExistsLoginRequiredError(ValueError):
code = 'account_exists_login_required'
@@ -54,14 +61,45 @@ class UserService:
def __init__(self, ap: Application) -> None:
self.ap = ap
self._create_user_lock = asyncio.Lock()
self._password_hash_lock = asyncio.Semaphore(1)
self._password_hash_lock = asyncio.Lock()
self._space_oauth_state_lock = asyncio.Lock()
self._space_oauth_states: dict[str, tuple[str, str | None, float, str | None]] = {}
self._space_oauth_state_expiry_heap: list[tuple[float, str]] = []
@staticmethod
def _space_oauth_state_digest(state: str) -> str:
return hashlib.sha256(state.encode('utf-8')).hexdigest()
def _prune_space_oauth_states(self, now: float) -> None:
while self._space_oauth_state_expiry_heap:
expires_at, digest = self._space_oauth_state_expiry_heap[0]
entry = self._space_oauth_states.get(digest)
if entry is None or entry[2] != expires_at:
heapq.heappop(self._space_oauth_state_expiry_heap)
continue
if expires_at > now:
break
heapq.heappop(self._space_oauth_state_expiry_heap)
self._space_oauth_states.pop(digest, None)
max_heap_entries = max(
_SPACE_OAUTH_STATE_HEAP_COMPACT_FLOOR,
len(self._space_oauth_states) * _SPACE_OAUTH_STATE_HEAP_MAX_MULTIPLIER,
)
if len(self._space_oauth_state_expiry_heap) > max_heap_entries:
self._space_oauth_state_expiry_heap[:] = [
(entry[2], digest) for digest, entry in self._space_oauth_states.items()
]
heapq.heapify(self._space_oauth_state_expiry_heap)
def _evict_earliest_space_oauth_state(self) -> None:
while self._space_oauth_state_expiry_heap:
expires_at, digest = heapq.heappop(self._space_oauth_state_expiry_heap)
entry = self._space_oauth_states.get(digest)
if entry is not None and entry[2] == expires_at:
self._space_oauth_states.pop(digest, None)
return
async def issue_space_oauth_state(
self,
purpose: typing.Literal['login', 'bind'],
@@ -85,11 +123,14 @@ class UserService:
expires_at = time.monotonic() + min(ttl_seconds, 600)
async with self._space_oauth_state_lock:
now = time.monotonic()
self._space_oauth_states = {key: value for key, value in self._space_oauth_states.items() if value[2] > now}
if len(self._space_oauth_states) >= 4096:
oldest = min(self._space_oauth_states, key=lambda key: self._space_oauth_states[key][2])
self._space_oauth_states.pop(oldest, None)
self._prune_space_oauth_states(now)
if len(self._space_oauth_states) >= _SPACE_OAUTH_STATE_MAX_ENTRIES:
self._evict_earliest_space_oauth_state()
self._space_oauth_states[digest] = (purpose, account_uuid, expires_at, launch_workspace_uuid)
heapq.heappush(
self._space_oauth_state_expiry_heap,
(expires_at, digest),
)
return raw_state
async def consume_space_oauth_state_details(
@@ -129,8 +170,14 @@ class UserService:
return consumed.account
async def _hash_password(self, password: str) -> str:
if self._password_hash_lock.locked():
raise bounded_executor.BlockingWorkCapacityError(
'Password hashing capacity reached',
scope='system:authentication',
)
async with self._password_hash_lock:
return await asyncio.to_thread(argon2.PasswordHasher().hash, password)
with bounded_executor.blocking_work_scope('system:authentication'):
return await asyncio.to_thread(argon2.PasswordHasher().hash, password)
def _require_local_directory(self) -> None:
if self._uses_control_plane_directory():
@@ -143,8 +190,14 @@ class UserService:
return bool(workspace_service is not None and workspace_service.policy.multi_workspace_enabled)
async def _verify_password(self, hashed_password: str, password: str) -> None:
if self._password_hash_lock.locked():
raise bounded_executor.BlockingWorkCapacityError(
'Password hashing capacity reached',
scope='system:authentication',
)
async with self._password_hash_lock:
await asyncio.to_thread(argon2.PasswordHasher().verify, hashed_password, password)
with bounded_executor.blocking_work_scope('system:authentication'):
await asyncio.to_thread(argon2.PasswordHasher().verify, hashed_password, password)
async def _update_space_provider_for_account(self, account: typing.Any, api_key: str) -> None:
"""Refresh the OSS Workspace Space provider without guessing a SaaS Workspace.
src/langbot/pkg/box/admission.py
+2 -1
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import asyncio
import datetime as dt
import time
import weakref
from collections.abc import Callable
from typing import TYPE_CHECKING
@@ -52,7 +53,7 @@ class SandboxAdmissionController:
self.client = client
self.policy = policy
self._wall_time = wall_time
self._locks: dict[str, asyncio.Lock] = {}
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
self._highest_revisions: dict[str, int] = {}
def _workspace_lock(self, workspace_uuid: str) -> asyncio.Lock:
src/langbot/pkg/box/service.py
+92 -35
View File
@@ -21,6 +21,7 @@ from .admission import SandboxAdmissionController, require_cloud_admission_polic
from .connector import BoxRuntimeConnector, _get_box_config
from . import secure_fs
from ..telemetry import features as telemetry_features
from ..utils import httpclient
from ..api.http.context import ExecutionContext
from ..api.http.service.tenant import TenantContext, require_workspace_uuid
from langbot_plugin.box.errors import BoxAdmissionError, BoxError, BoxValidationError
@@ -39,6 +40,53 @@ _MAX_RECENT_ERRORS = 50
_MIB = 1024 * 1024
def _create_shared_workspace_probe(root: str, marker_name: str, payload: bytes) -> None:
"""Create and durably flush a no-follow probe without blocking the event loop."""
directory_flags = os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0)
nofollow = getattr(os, 'O_NOFOLLOW', 0)
root_fd: int | None = None
marker_fd: int | None = None
marker_created = False
try:
root_fd = os.open(root, directory_flags | nofollow)
marker_fd = os.open(
marker_name,
os.O_WRONLY | os.O_CREAT | os.O_EXCL | nofollow,
0o600,
dir_fd=root_fd,
)
marker_created = True
remaining = memoryview(payload)
while remaining:
written = os.write(marker_fd, remaining)
if written <= 0:
raise BoxValidationError('Failed to write Cloud Box shared-volume probe')
remaining = remaining[written:]
os.fsync(marker_fd)
except Exception:
if root_fd is not None and marker_created:
with contextlib.suppress(FileNotFoundError):
os.unlink(marker_name, dir_fd=root_fd)
raise
finally:
if marker_fd is not None:
os.close(marker_fd)
if root_fd is not None:
os.close(root_fd)
def _remove_shared_workspace_probe(root: str, marker_name: str) -> None:
directory_flags = os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0)
nofollow = getattr(os, 'O_NOFOLLOW', 0)
root_fd = os.open(root, directory_flags | nofollow)
try:
with contextlib.suppress(FileNotFoundError):
os.unlink(marker_name, dir_fd=root_fd)
finally:
os.close(root_fd)
def _is_path_under(path: str, root: str) -> bool:
"""Check whether *path* equals *root* or is a child of *root*."""
return path == root or path.startswith(f'{root}{os.sep}')
@@ -267,29 +315,15 @@ class BoxService:
marker_name = f'{BOX_SHARED_WORKSPACE_PROBE_PREFIX}{secrets.token_hex(16)}'
marker_payload = secrets.token_bytes(64)
expected_digest = hashlib.sha256(marker_payload).hexdigest()
directory_flags = os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0)
nofollow = getattr(os, 'O_NOFOLLOW', 0)
root_fd: int | None = None
marker_fd: int | None = None
marker_created = False
probe_created = False
try:
root_fd = os.open(self.default_workspace, directory_flags | nofollow)
marker_fd = os.open(
await asyncio.to_thread(
_create_shared_workspace_probe,
self.default_workspace,
marker_name,
os.O_WRONLY | os.O_CREAT | os.O_EXCL | nofollow,
0o600,
dir_fd=root_fd,
marker_payload,
)
marker_created = True
remaining = memoryview(marker_payload)
while remaining:
written = os.write(marker_fd, remaining)
if written <= 0:
raise BoxValidationError('Failed to write Cloud Box shared-volume probe')
remaining = remaining[written:]
os.fsync(marker_fd)
os.close(marker_fd)
marker_fd = None
probe_created = True
result = await self.client.verify_shared_workspace(marker_name)
if (
@@ -306,15 +340,12 @@ class BoxService:
except Exception as exc:
raise BoxValidationError('Cloud Box shared durable Workspace volume verification failed') from exc
finally:
if marker_fd is not None:
os.close(marker_fd)
if root_fd is not None:
if marker_created:
try:
os.unlink(marker_name, dir_fd=root_fd)
except FileNotFoundError:
pass
os.close(root_fd)
if probe_created:
await asyncio.to_thread(
_remove_shared_workspace_probe,
self.default_workspace,
marker_name,
)
@property
def available(self) -> bool:
@@ -886,7 +917,13 @@ class BoxService:
mime = data[5:split_index]
data = data[split_index + 8 :]
try:
return _b64.b64decode(data), mime
max_encoded_bytes = 4 * ((BoxService._ATTACHMENT_MAX_BYTES + 2) // 3)
if not isinstance(data, (str, bytes)) or len(data) > max_encoded_bytes:
return None
decoded = _b64.b64decode(data)
if len(decoded) > BoxService._ATTACHMENT_MAX_BYTES:
return None
return decoded, mime
except Exception:
return None
@@ -895,10 +932,25 @@ class BoxService:
try:
import httpx
async with httpx.AsyncClient(timeout=30) as client:
resp = await client.get(url)
resp.raise_for_status()
return resp.content, resp.headers.get('Content-Type', 'application/octet-stream')
async with httpx.AsyncClient(
timeout=30,
event_hooks=httpclient.httpx_response_limit_hooks(BoxService._ATTACHMENT_MAX_BYTES),
) as client:
async with client.stream('GET', url) as resp:
resp.raise_for_status()
declared_size = resp.headers.get('content-length')
if declared_size is not None:
try:
if int(declared_size) > BoxService._ATTACHMENT_MAX_BYTES:
return None
except ValueError:
pass
body = bytearray()
async for chunk in resp.aiter_bytes(chunk_size=64 * 1024):
body.extend(chunk)
if len(body) > BoxService._ATTACHMENT_MAX_BYTES:
return None
return bytes(body), resp.headers.get('Content-Type', 'application/octet-stream')
except Exception:
return None
@@ -907,8 +959,13 @@ class BoxService:
try:
import aiofiles
if await asyncio.to_thread(os.path.getsize, path) > BoxService._ATTACHMENT_MAX_BYTES:
return None
async with aiofiles.open(path, 'rb') as f:
return await f.read(), 'application/octet-stream'
data = await f.read(BoxService._ATTACHMENT_MAX_BYTES + 1)
if len(data) > BoxService._ATTACHMENT_MAX_BYTES:
return None
return data, 'application/octet-stream'
except Exception:
return None
src/langbot/pkg/cloud/directory_projection.py
+37
View File
@@ -34,6 +34,7 @@ from .directory import (
DirectorySnapshot,
DirectoryWorkspace,
)
from .entitlements import EntitlementResolver
if TYPE_CHECKING:
@@ -257,6 +258,7 @@ class DirectoryProjectionService:
await session.flush()
await self._reconcile_entitlement_snapshot_set(snapshot)
self._record_success()
self._consumer_cursor = snapshot.cursor
@@ -343,10 +345,45 @@ class DirectoryProjectionService:
await session.flush()
projection_caught_up = batch.cursor == batch.high_water_cursor and int(state.cursor) == batch.cursor
await self._update_entitlement_workspace_activity(
returned.values(),
requested_workspace_uuids=requested,
)
if projection_caught_up:
self._record_success()
self._consumer_cursor = batch.cursor
async def _reconcile_entitlement_snapshot_set(
self,
snapshot: DirectorySnapshot,
) -> None:
resolver = getattr(self.ap, 'entitlement_resolver', None)
if not isinstance(resolver, EntitlementResolver):
return
await resolver.reconcile_active_workspaces(
{workspace.uuid for workspace in snapshot.workspaces if workspace.status == WorkspaceStatus.ACTIVE.value}
)
async def _update_entitlement_workspace_activity(
self,
workspaces: Iterable[DirectoryWorkspace],
*,
requested_workspace_uuids: set[str],
) -> None:
resolver = getattr(self.ap, 'entitlement_resolver', None)
if not isinstance(resolver, EntitlementResolver):
return
returned = {workspace.uuid: workspace for workspace in workspaces}
active = {
workspace_uuid
for workspace_uuid, workspace in returned.items()
if workspace.status == WorkspaceStatus.ACTIVE.value
}
await resolver.update_workspace_activity(
active_workspace_uuids=active,
inactive_workspace_uuids=requested_workspace_uuids - active,
)
async def apply_event_batch(self, batch: DirectoryEventBatch) -> None:
"""Advance non-directory events after the adapter refreshes local caches."""
src/langbot/pkg/cloud/entitlements.py
+63
View File
@@ -122,6 +122,7 @@ class EntitlementResolver:
self._deployment_admission = deployment_admission
self._lock = asyncio.Lock()
self._snapshots: dict[str, tuple[int, str, EntitlementSnapshot]] = {}
self._active_workspace_uuids: frozenset[str] | None = None
@staticmethod
def _fingerprint(snapshot: EntitlementSnapshot) -> str:
@@ -136,6 +137,8 @@ class EntitlementResolver:
) -> EntitlementSnapshot:
if self._deployment_admission is not None:
self._deployment_admission()
async with self._lock:
self._require_projected_workspace_locked(workspace_uuid)
candidate = await self.provider.get_workspace_entitlement(workspace_uuid)
if self._deployment_admission is not None:
# A provider call may cross the Manifest expiry boundary.
@@ -154,6 +157,9 @@ class EntitlementResolver:
fingerprint = self._fingerprint(candidate)
async with self._lock:
# The directory may fence a Workspace while the provider call is
# in flight. Recheck before retaining or returning its snapshot.
self._require_projected_workspace_locked(workspace_uuid)
previous = self._snapshots.get(workspace_uuid)
if previous is not None:
previous_revision, previous_fingerprint, _ = previous
@@ -167,3 +173,60 @@ class EntitlementResolver:
candidate,
)
return candidate.model_copy(deep=True)
def _require_projected_workspace_locked(self, workspace_uuid: str) -> None:
active_workspace_uuids = self._active_workspace_uuids
if active_workspace_uuids is not None and workspace_uuid not in active_workspace_uuids:
raise EntitlementUnavailableError('Workspace is not active in the Cloud directory projection')
async def reconcile_active_workspaces(
self,
workspace_uuids: set[str] | frozenset[str],
) -> None:
"""Drop entitlement history for Workspaces fenced by the directory."""
active = frozenset(workspace_uuids)
async with self._lock:
self._active_workspace_uuids = active
self._snapshots = {
workspace_uuid: cached for workspace_uuid, cached in self._snapshots.items() if workspace_uuid in active
}
async def set_workspace_active(
self,
workspace_uuid: str,
*,
active: bool,
) -> None:
"""Apply one incremental directory activity change."""
await self.update_workspace_activity(
active_workspace_uuids={workspace_uuid} if active else set(),
inactive_workspace_uuids=set() if active else {workspace_uuid},
)
async def update_workspace_activity(
self,
*,
active_workspace_uuids: set[str] | frozenset[str],
inactive_workspace_uuids: set[str] | frozenset[str],
) -> None:
"""Apply one directory delta without copying the active set per item."""
active_updates = set(active_workspace_uuids)
inactive_updates = set(inactive_workspace_uuids)
if active_updates & inactive_updates:
raise ValueError('Workspace activity update contains conflicting entries')
async with self._lock:
current = set(self._active_workspace_uuids or ())
current.update(active_updates)
current.difference_update(inactive_updates)
for workspace_uuid in inactive_updates:
self._snapshots.pop(workspace_uuid, None)
self._active_workspace_uuids = frozenset(current)
def snapshot_counts(self) -> dict[str, int]:
return {
'active_workspaces': len(self._active_workspace_uuids or ()),
'cached_snapshots': len(self._snapshots),
}
src/langbot/pkg/cloud/launch.py
+34 -4
View File
@@ -4,6 +4,7 @@ import asyncio
import base64
import binascii
import hashlib
import heapq
import json
import os
import time
@@ -22,6 +23,9 @@ CONTROL_PLANE_TYP = 'langbot-control-plane+jwt'
LAUNCH_KIND = 'workspace.launch'
EXPECTED_ISSUER = 'langbot-space'
EXPECTED_AUDIENCE = 'langbot-cloud-runtime'
_CONSUMED_JTI_MAX_ENTRIES = 4096
_CONSUMED_JTI_HEAP_COMPACT_FLOOR = 64
_CONSUMED_JTI_HEAP_MAX_MULTIPLIER = 4
class SpaceLaunchError(ValueError):
@@ -107,6 +111,7 @@ class SpaceLaunchService:
self._wall_time = wall_time
self._replay_lock = asyncio.Lock()
self._consumed_jtis: dict[str, int] = {}
self._consumed_jti_expiry_heap: list[tuple[int, str]] = []
async def consume_assertion(
self,
@@ -210,13 +215,38 @@ class SpaceLaunchService:
digest = hashlib.sha256(jti.encode('utf-8')).hexdigest()
now = int(self._wall_time())
async with self._replay_lock:
self._consumed_jtis = {existing: expiry for existing, expiry in self._consumed_jtis.items() if expiry > now}
self._prune_consumed_jtis(now)
if digest in self._consumed_jtis:
raise SpaceLaunchError('Launch assertion has already been consumed')
if len(self._consumed_jtis) >= 4096:
oldest = min(self._consumed_jtis, key=lambda key: self._consumed_jtis[key])
self._consumed_jtis.pop(oldest, None)
if len(self._consumed_jtis) >= _CONSUMED_JTI_MAX_ENTRIES:
# Evicting a still-valid digest would make a signed launch
# assertion replayable. Bound memory by failing closed instead.
raise SpaceLaunchError('Launch assertion replay cache capacity reached')
self._consumed_jtis[digest] = expires_at
heapq.heappush(
self._consumed_jti_expiry_heap,
(expires_at, digest),
)
def _prune_consumed_jtis(self, now: int) -> None:
while self._consumed_jti_expiry_heap:
expires_at, digest = self._consumed_jti_expiry_heap[0]
current_expiry = self._consumed_jtis.get(digest)
if current_expiry != expires_at:
heapq.heappop(self._consumed_jti_expiry_heap)
continue
if expires_at > now:
break
heapq.heappop(self._consumed_jti_expiry_heap)
self._consumed_jtis.pop(digest, None)
max_heap_entries = max(
_CONSUMED_JTI_HEAP_COMPACT_FLOOR,
len(self._consumed_jtis) * _CONSUMED_JTI_HEAP_MAX_MULTIPLIER,
)
if len(self._consumed_jti_expiry_heap) > max_heap_entries:
self._consumed_jti_expiry_heap[:] = [(expiry, digest) for digest, expiry in self._consumed_jtis.items()]
heapq.heapify(self._consumed_jti_expiry_heap)
@staticmethod
def _bounded_float(
src/langbot/pkg/core/app.py
+102 -3
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import logging
import asyncio
import contextlib
import traceback
import os
@@ -18,7 +19,7 @@ from ..plugin import connector as plugin_connector
from ..pipeline import pool
from ..pipeline import controller, pipelinemgr
from ..pipeline import aggregator as message_aggregator
from ..utils import version as version_mgr, proxy as proxy_mgr
from ..utils import version as version_mgr, proxy as proxy_mgr, httpclient
from ..persistence import mgr as persistencemgr
from ..api.http.controller import main as http_controller
from ..api.http.service import user as user_service
@@ -36,7 +37,7 @@ from ..api.http.service import skill as skill_service
from ..api.http.service import maintenance as maintenance_service
from ..discover import engine as discover_engine
from ..storage import mgr as storagemgr
from ..utils import logcache
from ..utils import bounded_executor, event_loop_monitor, logcache
from . import taskmgr
from . import entities as core_entities
from ..rag.knowledge import kbmgr as rag_mgr
@@ -191,14 +192,78 @@ class Application:
maintenance_service: maintenance_service.MaintenanceService = None
blocking_executor: bounded_executor.BoundedThreadPoolExecutor | None = None
event_loop_monitor: event_loop_monitor.EventLoopLagMonitor
def __init__(self):
self._shutdown_lock = asyncio.Lock()
self._shutdown_complete = False
self._shutdown_task: asyncio.Task | None = None
self.event_loop_monitor = event_loop_monitor.EventLoopLagMonitor()
def get_runtime_resource_stats(self) -> dict[str, object]:
"""Return aggregate O(1) counters for liveness and soak validation."""
try:
asyncio_tasks = len(asyncio.all_tasks(self.event_loop))
except (RuntimeError, TypeError):
asyncio_tasks = 0
task_stats = self.task_mgr.get_stats() if self.task_mgr is not None else {}
query_pool_stats = {}
if self.query_pool is not None:
query_pool_stats = {
'queued': len(self.query_pool.queries),
'cached': len(self.query_pool.cached_queries),
'active_workspaces': len(self.query_pool.active_query_count_by_workspace),
}
model_stats = {}
if self.model_mgr is not None:
model_stats = {
'providers': len(self.model_mgr.provider_dict),
'llms': len(self.model_mgr.llm_model_dict),
'embeddings': len(self.model_mgr.embedding_model_dict),
'rerankers': len(self.model_mgr.rerank_model_dict),
}
runtime_stats = {
'bots': len(getattr(self.platform_mgr, '_bots_by_key', {})),
'pipelines': len(getattr(self.pipeline_mgr, '_pipelines_by_key', {})),
'knowledge_bases': len(getattr(self.rag_mgr, 'knowledge_bases', {})),
'plugin_installations': len(
getattr(
self.plugin_connector,
'_known_desired_states',
{},
)
),
}
mcp_loader = getattr(self.tool_mgr, 'mcp_tool_loader', None)
runtime_stats.update(
{
'mcp_sessions': len(getattr(mcp_loader, '_sessions', {})),
'mcp_host_tasks': len(getattr(mcp_loader, '_hosted_mcp_tasks', ())),
'mcp_dispatch_tasks': len(getattr(mcp_loader, '_host_dispatch_tasks', ())),
}
)
return {
'asyncio_tasks': asyncio_tasks,
'event_loop': self.event_loop_monitor.snapshot(),
'blocking_executor': (self.blocking_executor.snapshot() if self.blocking_executor is not None else {}),
'application_tasks': task_stats,
'query_pool': query_pool_stats,
'models': model_stats,
'runtimes': runtime_stats,
'telemetry_tasks': len(getattr(self.telemetry, 'send_tasks', ())),
}
async def initialize(self):
pass
async def run(self):
self.event_loop_monitor.start()
try:
if self.directory_projection_service is not None:
self.task_mgr.create_task(
@@ -392,18 +457,36 @@ class Application:
if self.task_mgr is not None:
self.task_mgr.cancel_by_scope(core_entities.LifecycleControlScope.APPLICATION)
with contextlib.suppress(Exception):
await self.event_loop_monitor.stop()
mcp_mount = getattr(self.http_ctrl, 'mcp_mount', None)
if mcp_mount is not None:
with contextlib.suppress(Exception):
await mcp_mount.stop_session_manager()
if self.platform_mgr is not None:
with contextlib.suppress(Exception):
await self.platform_mgr.shutdown()
if self.tool_mgr is not None:
with contextlib.suppress(Exception):
await self.tool_mgr.shutdown()
if self.model_mgr is not None:
with contextlib.suppress(Exception):
await self.model_mgr.shutdown()
if self.box_service is not None:
with contextlib.suppress(Exception):
await self.box_service.shutdown()
if self.plugin_connector is not None:
with contextlib.suppress(Exception):
await self.plugin_connector.aclose()
if self.telemetry is not None:
with contextlib.suppress(Exception):
await self.telemetry.shutdown()
if self.vector_db_mgr is not None:
with contextlib.suppress(Exception):
await self.vector_db_mgr.shutdown()
if self.storage_mgr is not None:
with contextlib.suppress(Exception):
await self.storage_mgr.shutdown()
manifest_provider = getattr(self.deployment, 'manifest_provider', None)
if manifest_provider is not None:
with contextlib.suppress(Exception):
@@ -413,13 +496,29 @@ class Application:
tasks = [wrapper.task for wrapper in self.task_mgr.tasks if not wrapper.task.done()]
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
with contextlib.suppress(Exception):
await httpclient.close_all()
persistence_shutdown = getattr(self.persistence_mgr, 'shutdown', None)
if callable(persistence_shutdown):
with contextlib.suppress(Exception):
await persistence_shutdown()
else:
# Compatibility for lightweight test/application doubles.
persistence_db = getattr(self.persistence_mgr, 'db', None)
persistence_engine = getattr(persistence_db, 'engine', None)
if persistence_engine is not None:
with contextlib.suppress(Exception):
await persistence_engine.dispose()
self._shutdown_complete = True
def dispose(self):
"""Compatibility wrapper for callers that cannot await shutdown."""
if self._shutdown_complete:
return
loop = self.event_loop
if loop is not None and not loop.is_closed():
loop.create_task(self.shutdown())
if self._shutdown_task is None or self._shutdown_task.done():
self._shutdown_task = loop.create_task(self.shutdown())
return
if self.plugin_connector is not None:
self.plugin_connector.dispose()
src/langbot/pkg/core/boot.py
+15 -6
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import traceback
import asyncio
import contextlib
import os
from . import app
@@ -32,14 +33,22 @@ async def make_app(loop: asyncio.AbstractEventLoop) -> app.Application:
ap.event_loop = loop
# Execute startup stage
for stage_name in stage_order:
stage_cls = stage.preregistered_stages[stage_name]
stage_inst = stage_cls()
try:
# Execute startup stage
for stage_name in stage_order:
stage_cls = stage.preregistered_stages[stage_name]
stage_inst = stage_cls()
await stage_inst.run(ap)
await stage_inst.run(ap)
await ap.initialize()
await ap.initialize()
except BaseException:
# ``main()`` cannot clean up a partially built application because
# ``make_app()`` has not returned it yet. Release managers, pools and
# child processes that earlier startup stages already attached.
with contextlib.suppress(BaseException):
await ap.shutdown()
raise
return ap
src/langbot/pkg/core/errors.py
+2
View File
@@ -0,0 +1,2 @@
class TaskCapacityError(RuntimeError):
"""Raised when the configured user-task admission limit is exhausted."""
src/langbot/pkg/core/stages/build_app.py
+17 -7
View File
@@ -139,8 +139,8 @@ class BuildAppStage(stage.BootingStage):
ap.log_cache = log_cache
storage_mgr_inst = storagemgr.StorageMgr(ap)
await storage_mgr_inst.initialize()
ap.storage_mgr = storage_mgr_inst
await storage_mgr_inst.initialize()
persistence_mgr_inst = persistencemgr.PersistenceManager(
ap,
@@ -166,6 +166,12 @@ class BuildAppStage(stage.BootingStage):
if not workspace_policy.multi_workspace_enabled:
await workspace_service_inst.ensure_singleton_workspace()
ap.workspace_service = workspace_service_inst
if workspace_policy.multi_workspace_enabled:
# Directory refresh starts in Application.run(), after this serial
# build graph. Share one validated immutable binding snapshot
# across model/platform/pipeline/RAG/plugin initialization instead
# of repeating tenant validation for every manager.
await workspace_service_inst.prime_startup_execution_bindings()
ap.workspace_collaboration_service = workspace_collaboration_module.WorkspaceCollaborationService(
ap,
@@ -188,14 +194,17 @@ class BuildAppStage(stage.BootingStage):
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
)
concurrency_config = ap.instance_config.data.get('concurrency', {})
ap.query_pool = pool.QueryPool(
singleton_context_resolver=resolve_singleton_execution_context,
max_queries=int(concurrency_config.get('pending_queries', 1000)),
max_queries_per_workspace=int(concurrency_config.get('pending_queries_per_workspace', 100)),
)
# Telemetry manager: attach to app so other components can call via self.ap.telemetry
telemetry_inst = telemetry_module.TelemetryManager(ap)
await telemetry_inst.initialize()
ap.telemetry = telemetry_inst
await telemetry_inst.initialize()
# Survey manager
survey_inst = survey_module.SurveyManager(ap)
@@ -215,16 +224,16 @@ class BuildAppStage(stage.BootingStage):
ap.sess_mgr = llm_session_mgr_inst
box_service_inst = box_service.BoxService(ap)
await box_service_inst.initialize()
ap.box_service = box_service_inst
await box_service_inst.initialize()
llm_tool_mgr_inst = llm_tool_mgr.ToolManager(ap)
await llm_tool_mgr_inst.initialize()
ap.tool_mgr = llm_tool_mgr_inst
await llm_tool_mgr_inst.initialize()
im_mgr_inst = im_mgr.PlatformManager(ap=ap)
await im_mgr_inst.initialize()
ap.platform_mgr = im_mgr_inst
await im_mgr_inst.initialize()
# Initialize webhook pusher
webhook_pusher_inst = WebhookPusher(ap)
@@ -252,12 +261,12 @@ class BuildAppStage(stage.BootingStage):
# 初始化向量数据库管理器
vectordb_mgr_inst = vectordb_mgr.VectorDBManager(ap)
await vectordb_mgr_inst.initialize()
ap.vector_db_mgr = vectordb_mgr_inst
await vectordb_mgr_inst.initialize()
http_ctrl = http_controller.HTTPController(ap)
await http_ctrl.initialize()
ap.http_ctrl = http_ctrl
await http_ctrl.initialize()
monitoring_service_inst = monitoring_service.MonitoringService(ap)
ap.monitoring_service = monitoring_service_inst
@@ -277,6 +286,7 @@ class BuildAppStage(stage.BootingStage):
ap.logger.warning(f'Plugin runtime unavailable during startup; reconnecting in background: {exc}')
plugin_connector_inst.schedule_reconnect()
ap.plugin_connector = plugin_connector_inst
workspace_service_inst.release_startup_execution_bindings()
ctrl = controller.Controller(ap)
ap.ctrl = ctrl
src/langbot/pkg/core/stages/load_config.py
+20 -1
View File
@@ -3,7 +3,7 @@ from __future__ import annotations
import os
import copy
from typing import Any
from langbot.pkg.utils import constants
from langbot.pkg.utils import bounded_executor, constants
import yaml
import importlib.resources as resources
import uuid
@@ -14,6 +14,13 @@ from ..bootutils import config
_RUNTIME_POLICY_DEFAULTS = {
'system': {
'blocking_executor': {
'max_workers': bounded_executor.DEFAULT_MAX_WORKERS,
'max_pending': bounded_executor.DEFAULT_MAX_PENDING,
'max_inflight_per_scope': (bounded_executor.DEFAULT_MAX_INFLIGHT_PER_SCOPE),
}
},
'plugin': {
'worker': {
'max_cpus': 1.0,
@@ -21,6 +28,10 @@ _RUNTIME_POLICY_DEFAULTS = {
'max_pids': 128,
'max_open_files': 256,
'max_file_size_mb': 512,
'max_workers': 16,
'max_total_cpus': 8.0,
'max_total_memory_mb': 8192,
'max_installations': 10000,
'require_hard_limits': False,
}
},
@@ -205,6 +216,14 @@ class LoadConfigStage(stage.BootingStage):
# Apply environment variable overrides to data/config.yaml
ap.instance_config.data = _apply_env_overrides_to_config(ap.instance_config.data)
blocking_config = ap.instance_config.data['system']['blocking_executor']
ap.blocking_executor = bounded_executor.configure_bounded_default_executor(
ap.event_loop,
max_workers=blocking_config['max_workers'],
max_pending=blocking_config['max_pending'],
max_inflight_per_scope=blocking_config['max_inflight_per_scope'],
)
await ap.instance_config.dump_config()
# load or generate instance id
src/langbot/pkg/core/stages/show_notes.py
+8 -3
View File
@@ -1,8 +1,7 @@
from __future__ import annotations
import asyncio
from .. import stage, app, note
from .. import entities as core_entities
from ...utils import importutil
from .. import notes
@@ -31,6 +30,12 @@ class ShowNotesStage(stage.BootingStage):
if msg:
ap.logger.log(level, msg)
asyncio.create_task(ayield_note(note_inst))
ap.task_mgr.create_task(
ayield_note(note_inst),
kind='launch-note',
name=f'launch-note-{note_cls.__name__}',
scopes=[core_entities.LifecycleControlScope.APPLICATION],
instance_uuid=ap.workspace_service.instance_uuid,
)
except Exception:
continue
src/langbot/pkg/core/task_boundary.py
+8
View File
@@ -4,6 +4,8 @@ import asyncio
import contextvars
import typing
from ..utils import bounded_executor
T = typing.TypeVar('T')
@@ -14,6 +16,7 @@ def create_detached_task(
loop: asyncio.AbstractEventLoop | None = None,
name: str | None = None,
after_commit_manager: typing.Any | None = None,
workspace_uuid: str | None = None,
) -> asyncio.Task[T]:
"""Create a task that inherits no request-local ContextVars.
@@ -33,6 +36,11 @@ def create_detached_task(
if callable(gate_factory):
gate = gate_factory(after_commit_manager)
task_coro = _wait_for_commit(coro, gate) if gate is not None else coro
if workspace_uuid is not None:
task_coro = bounded_executor.run_in_blocking_work_scope(
task_coro,
workspace_uuid,
)
return task_loop.create_task(task_coro, name=name, context=contextvars.Context())
src/langbot/pkg/core/taskmgr.py
+49 -1
View File
@@ -7,6 +7,7 @@ import time
from . import app
from . import entities as core_entities
from .errors import TaskCapacityError
from .task_boundary import create_detached_task
@@ -22,13 +23,18 @@ class TaskContext:
metadata: dict
"""Structured metadata for progress reporting"""
def __init__(self):
def __init__(self, max_log_chars: int = 200000):
self.current_action = 'default'
self.log = ''
self.metadata = {}
self.max_log_chars = max(int(max_log_chars), 1)
def _log(self, msg: str):
self.log += msg + '\n'
if len(self.log) > self.max_log_chars:
marker = '[older task output truncated]\n'
keep = max(self.max_log_chars - len(marker), 0)
self.log = marker + (self.log[-keep:] if keep else '')
def set_current_action(self, action: str):
self.current_action = action
@@ -131,6 +137,7 @@ class TaskWrapper:
loop=self.ap.event_loop,
name=name or None,
after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
workspace_uuid=workspace_uuid,
)
self.task_type = task_type
self.kind = kind
@@ -207,6 +214,39 @@ class AsyncTaskManager:
self.ap = ap
self.tasks = []
def _task_log_limit(self) -> int:
value = self.ap.instance_config.data.get('system', {}).get('task_retention', {}).get('max_log_chars', 200000)
try:
value = int(value)
except (TypeError, ValueError):
value = 200000
return max(value, 1)
def _user_task_limit(self, name: str, default: int) -> int:
value = self.ap.instance_config.data.get('system', {}).get('task_retention', {}).get(name, default)
try:
value = int(value)
except (TypeError, ValueError):
value = default
return max(value, 1)
def _admit_user_task(self, coro: typing.Coroutine, workspace_uuid: str | None) -> None:
active_user_tasks = [
wrapper for wrapper in self.tasks if wrapper.task_type == 'user' and not wrapper.task.done()
]
global_limit = self._user_task_limit('max_active_user_tasks', 256)
if len(active_user_tasks) >= global_limit:
coro.close()
raise TaskCapacityError('The instance has too many active user operations')
if workspace_uuid is None:
return
workspace_limit = self._user_task_limit('max_active_user_tasks_per_workspace', 8)
active_workspace_tasks = sum(1 for wrapper in active_user_tasks if wrapper.workspace_uuid == workspace_uuid)
if active_workspace_tasks >= workspace_limit:
coro.close()
raise TaskCapacityError('The Workspace has too many active user operations')
def create_task(
self,
coro: typing.Coroutine,
@@ -220,6 +260,13 @@ class AsyncTaskManager:
workspace_uuid: str | None = None,
placement_generation: int | None = None,
) -> TaskWrapper:
if context is None:
context = TaskContext(max_log_chars=self._task_log_limit())
else:
context.max_log_chars = self._task_log_limit()
if len(context.log) > context.max_log_chars:
context.log = context.log[-context.max_log_chars :]
wrapper = TaskWrapper(
self.ap,
coro,
@@ -250,6 +297,7 @@ class AsyncTaskManager:
workspace_uuid: str | None = None,
placement_generation: int | None = None,
) -> TaskWrapper:
self._admit_user_task(coro, workspace_uuid)
return self.create_task(
coro,
'user',
src/langbot/pkg/discover/engine.py
+1
View File
@@ -211,6 +211,7 @@ class ComponentDiscoveryEngine:
def __init__(self, ap: app.Application):
self.ap = ap
self.components = {}
def load_component_manifest(self, path: str, owner: str = 'builtin', no_save: bool = False) -> Component | None:
"""加载组件清单"""
src/langbot/pkg/persistence/databases/postgresql.py
+45 -2
View File
@@ -11,11 +11,27 @@ from ..postgresql_url import normalize_asyncpg_url
class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
"""PostgreSQL database manager"""
@staticmethod
def _pool_integer(
config: dict,
name: str,
default: int,
*,
minimum: int,
) -> int:
value = config.get(name, default)
if isinstance(value, bool) or not isinstance(value, int) or value < minimum:
comparator = 'non-negative' if minimum == 0 else 'positive'
raise ValueError(f'database.postgresql.{name} must be a {comparator} integer')
return value
async def initialize(self) -> None:
postgresql_config = self.ap.instance_config.data.get('database', {}).get('postgresql', {})
if not isinstance(postgresql_config, dict):
raise ValueError('database.postgresql must be an object')
if self.url_override is not None:
engine_url = self.url_override
else:
postgresql_config = self.ap.instance_config.data.get('database', {}).get('postgresql', {})
explicit_url = postgresql_config.get('url')
if explicit_url:
if not isinstance(explicit_url, str):
@@ -37,4 +53,31 @@ class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
port=postgresql_config.get('port', 5432),
database=postgresql_config.get('database', 'postgres'),
)
self.engine = sqlalchemy_asyncio.create_async_engine(engine_url)
self.engine = sqlalchemy_asyncio.create_async_engine(
engine_url,
pool_size=self._pool_integer(
postgresql_config,
'pool_size',
10,
minimum=1,
),
max_overflow=self._pool_integer(
postgresql_config,
'max_overflow',
10,
minimum=0,
),
pool_timeout=self._pool_integer(
postgresql_config,
'pool_timeout_seconds',
30,
minimum=1,
),
pool_recycle=self._pool_integer(
postgresql_config,
'pool_recycle_seconds',
1800,
minimum=1,
),
pool_pre_ping=True,
)
src/langbot/pkg/persistence/mgr.py
+8
View File
@@ -187,6 +187,14 @@ class PersistenceManager:
if self.mode == PersistenceMode.OSS_COMPAT:
await self.write_space_model_providers()
async def shutdown(self) -> None:
"""Dispose the owned database engine when initialization or runtime ends."""
db = getattr(self, 'db', None)
engine = getattr(db, 'engine', None)
if engine is not None:
await engine.dispose()
@contextlib.asynccontextmanager
async def _release_migration_lock(self) -> typing.AsyncIterator[None]:
"""Serialize the complete PostgreSQL migration and validation window."""
src/langbot/pkg/persistence/release_migration.py
+1 -4
View File
@@ -177,10 +177,7 @@ async def run_cloud_release_migration(
await manager.initialize()
ap.logger.info('Cloud PostgreSQL release migration reached and validated the exact release head.')
finally:
db = getattr(manager, 'db', None)
engine = getattr(db, 'engine', None)
if engine is not None:
await engine.dispose()
await manager.shutdown()
async def run_cloud_release_migration_from_config(loop: asyncio.AbstractEventLoop) -> None:
src/langbot/pkg/pipeline/aggregator.py
+47 -13
View File
@@ -72,6 +72,12 @@ class MessageAggregator:
self.ap = ap
self.buffers = {}
self.lock = asyncio.Lock()
concurrency = self.ap.instance_config.data.get('concurrency', {})
self.max_buffers = max(int(concurrency.get('pending_queries', 1000)), 1)
self.max_buffers_per_workspace = max(
int(concurrency.get('pending_queries_per_workspace', 100)),
1,
)
def _get_aggregation_key(
self,
@@ -184,15 +190,26 @@ class MessageAggregator:
)
force_flush = False
bypass_aggregation = False
async with self.lock:
buffer = self.buffers.get(aggregation_key)
if buffer is None:
buffer = SessionBuffer(
aggregation_key=aggregation_key,
execution_context=execution_context,
messages=[pending_msg],
workspace_buffer_count = sum(
1
for key in self.buffers
if key[0] == execution_context.instance_uuid
and key[1] == execution_context.workspace_uuid
and key[2] == execution_context.placement_generation
)
self.buffers[aggregation_key] = buffer
if len(self.buffers) >= self.max_buffers or workspace_buffer_count >= self.max_buffers_per_workspace:
bypass_aggregation = True
else:
buffer = SessionBuffer(
aggregation_key=aggregation_key,
execution_context=execution_context,
messages=[pending_msg],
)
self.buffers[aggregation_key] = buffer
else:
if buffer.execution_context != execution_context:
raise ExecutionContextMismatchError('Aggregation buffer ExecutionContext changed for the same key')
@@ -200,14 +217,31 @@ class MessageAggregator:
buffer.timer_task.cancel()
buffer.messages.append(pending_msg)
buffer.last_message_time = time.time()
if len(buffer.messages) >= MAX_BUFFER_MESSAGES:
force_flush = True
else:
buffer.timer_task = create_detached_task(
self._delayed_flush(aggregation_key, delay, execution_context),
after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
)
if not bypass_aggregation:
buffer.last_message_time = time.time()
if len(buffer.messages) >= MAX_BUFFER_MESSAGES:
force_flush = True
else:
buffer.timer_task = create_detached_task(
self._delayed_flush(aggregation_key, delay, execution_context),
after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
workspace_uuid=execution_context.workspace_uuid,
)
if bypass_aggregation:
await self.ap.query_pool.add_query(
bot_uuid=bot_uuid,
launcher_type=launcher_type,
launcher_id=launcher_id,
sender_id=sender_id,
message_event=message_event,
message_chain=message_chain,
adapter=adapter,
pipeline_uuid=pipeline_uuid,
routed_by_rule=routed_by_rule,
execution_context=execution_context,
)
return
if force_flush:
await self._flush_buffer(aggregation_key, execution_context)
src/langbot/pkg/pipeline/cntfilter/filters/baiduexamine.py
+2 -2
View File
@@ -23,7 +23,7 @@ class BaiduCloudExamine(filter_model.ContentFilter):
'client_secret': self.ap.pipeline_cfg.data['baidu-cloud-examine']['api-secret'],
},
) as resp:
return (await resp.json())['access_token']
return (await httpclient.read_json_limited(resp))['access_token']
async def process(self, query: pipeline_query.Query, message: str) -> entities.FilterResult:
session = httpclient.get_session()
@@ -35,7 +35,7 @@ class BaiduCloudExamine(filter_model.ContentFilter):
},
data=f'text={message}'.encode('utf-8'),
) as resp:
result = await resp.json()
result = await httpclient.read_json_limited(resp)
if 'error_code' in result:
return entities.FilterResult(
src/langbot/pkg/pipeline/cntfilter/filters/banwords.py
+15 -17
View File
@@ -1,9 +1,9 @@
from __future__ import annotations
import re
from .. import filter as filter_model
from .. import entities
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from ....utils.safe_regex import SafeRegexError, mask_patterns
@filter_model.filter_class('ban-word-filter')
@@ -14,22 +14,20 @@ class BanWordFilter(filter_model.ContentFilter):
pass
async def process(self, query: pipeline_query.Query, message: str) -> entities.FilterResult:
found = False
for word in self.ap.sensitive_meta.data['words']:
match = re.findall(word, message)
if len(match) > 0:
found = True
for i in range(len(match)):
if self.ap.sensitive_meta.data['mask_word'] == '':
message = message.replace(
match[i],
self.ap.sensitive_meta.data['mask'] * len(match[i]),
)
else:
message = message.replace(match[i], self.ap.sensitive_meta.data['mask_word'])
try:
found, message = await mask_patterns(
self.ap.sensitive_meta.data['words'],
message,
mask=self.ap.sensitive_meta.data['mask'],
mask_word=self.ap.sensitive_meta.data['mask_word'],
)
except SafeRegexError as exc:
return entities.FilterResult(
level=entities.ResultLevel.BLOCK,
replacement='',
user_notice='内容检查规则执行失败，请联系管理员',
console_notice=f'Sensitive-word regex rejected: {exc}',
)
return entities.FilterResult(
level=entities.ResultLevel.MASKED if found else entities.ResultLevel.PASS,
src/langbot/pkg/pipeline/cntfilter/filters/cntignore.py
+20 -9
View File
@@ -1,9 +1,9 @@
from __future__ import annotations
import re
from .. import entities
from .. import filter as filter_model
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from ....utils.safe_regex import SafeRegexError, matches_any
@filter_model.filter_class('content-ignore')
@@ -28,14 +28,25 @@ class ContentIgnore(filter_model.ContentFilter):
)
if 'regexp' in query.pipeline_config['trigger']['ignore-rules']:
for rule in query.pipeline_config['trigger']['ignore-rules']['regexp']:
if re.search(rule, message):
return entities.FilterResult(
level=entities.ResultLevel.BLOCK,
replacement='',
user_notice='',
console_notice='Ignore message according to regexp rule in ignore_rules',
)
try:
matches = await matches_any(
query.pipeline_config['trigger']['ignore-rules']['regexp'],
message,
)
except SafeRegexError as exc:
return entities.FilterResult(
level=entities.ResultLevel.BLOCK,
replacement='',
user_notice='',
console_notice=f'Ignore-rule regex rejected: {exc}',
)
if matches:
return entities.FilterResult(
level=entities.ResultLevel.BLOCK,
replacement='',
user_notice='',
console_notice='Ignore message according to regexp rule in ignore_rules',
)
return entities.FilterResult(
level=entities.ResultLevel.PASS,
src/langbot/pkg/pipeline/controller.py
+97 -48
View File
@@ -38,57 +38,80 @@ class Controller:
raise WorkspaceInvariantError('Queued query instance does not match the active Workspace binding')
return execution_context
async def _process_query(self, selected_query: pipeline_query.Query) -> None:
async def _process_query(
self,
selected_query: pipeline_query.Query,
*,
selected_session=None,
global_slot_reserved: bool = False,
) -> None:
"""Run one selected query and always release its scheduling slot."""
try:
async with self.semaphore:
queued_context = get_query_execution_context(selected_query)
queued_context = get_query_execution_context(selected_query)
async def run_scoped_query() -> None:
execution_context = await self._assert_query_execution_active(selected_query)
pipeline_uuid = selected_query.pipeline_uuid
async def run_scoped_query() -> None:
execution_context = await self._assert_query_execution_active(selected_query)
pipeline_uuid = selected_query.pipeline_uuid
if pipeline_uuid:
pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(
execution_context,
pipeline_uuid,
)
if pipeline:
await pipeline.run(selected_query)
else:
self.ap.logger.warning(
f'Pipeline {pipeline_uuid} not found for query {selected_query.query_id}, query dropped'
)
if pipeline_uuid:
pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(
execution_context,
pipeline_uuid,
)
if pipeline:
await pipeline.run(selected_query)
else:
self.ap.logger.warning(f'No pipeline_uuid for query {selected_query.query_id}, query dropped')
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
cloud_runtime = (
getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
)
if cloud_runtime:
if not callable(tenant_scope):
raise RuntimeError('Cloud query processing requires an explicit tenant scope')
async with tenant_scope(queued_context.workspace_uuid):
await run_scoped_query()
self.ap.logger.warning(
f'Pipeline {pipeline_uuid} not found for query {selected_query.query_id}, query dropped'
)
else:
self.ap.logger.warning(f'No pipeline_uuid for query {selected_query.query_id}, query dropped')
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
if not callable(tenant_scope):
raise RuntimeError('Cloud query processing requires an explicit tenant scope')
async with tenant_scope(queued_context.workspace_uuid):
await run_scoped_query()
else:
await run_scoped_query()
except WorkspaceError as exc:
self.ap.logger.info(
f'Dropped query {selected_query.query_id} because its Workspace execution binding is stale: {exc}'
)
finally:
try:
try:
await self.ap.query_pool.remove_query(selected_query)
finally:
async with self.ap.query_pool:
session = selected_session or await self.ap.sess_mgr.get_session(selected_query)
try:
session._semaphore.release()
finally:
self.ap.query_pool.condition.notify_all()
finally:
if global_slot_reserved:
self.semaphore.release()
async def _drop_selected_query(self, selected_query, selected_session) -> None:
"""Undo scheduler ownership when work cannot be handed to a task."""
try:
await self.ap.query_pool.remove_query(selected_query)
finally:
async with self.ap.query_pool:
(await self.ap.sess_mgr.get_session(selected_query))._semaphore.release()
selected_session._semaphore.release()
self.ap.query_pool.condition.notify_all()
async def consumer(self):
"""事件处理循环"""
try:
while True:
while True:
try:
selected_query: pipeline_query.Query = None
selected_session = None
# 取请求
async with self.ap.query_pool:
@@ -101,7 +124,9 @@ class Controller:
if not session._semaphore.locked():
selected_query = query
selected_session = session
await session._semaphore.acquire()
self.ap.query_pool.mark_query_running_locked(query)
# Only log when actually selecting a query
self.ap.logger.debug(f'Selected query {query.query_id} for processing')
@@ -114,24 +139,48 @@ class Controller:
continue
if selected_query:
execution_context = get_query_execution_context(selected_query)
self.ap.task_mgr.create_task(
self._process_query(selected_query),
kind='query',
name=f'query-{selected_query.query_id}',
scopes=[
core_entities.LifecycleControlScope.APPLICATION,
core_entities.LifecycleControlScope.PLATFORM,
],
instance_uuid=execution_context.instance_uuid,
workspace_uuid=execution_context.workspace_uuid,
placement_generation=execution_context.placement_generation,
)
try:
# Reserve global capacity before creating the task.
# At most one selected query is held by this consumer
# while all pipeline slots are busy.
await self.semaphore.acquire()
except asyncio.CancelledError:
await self._drop_selected_query(selected_query, selected_session)
raise
except Exception as e:
# traceback.print_exc()
self.ap.logger.error(f'控制器循环出错: {e}')
self.ap.logger.error(f'Traceback: {traceback.format_exc()}')
execution_context = get_query_execution_context(selected_query)
process_coro = self._process_query(
selected_query,
selected_session=selected_session,
global_slot_reserved=True,
)
try:
self.ap.task_mgr.create_task(
process_coro,
kind='query',
name=f'query-{selected_query.query_id}',
scopes=[
core_entities.LifecycleControlScope.APPLICATION,
core_entities.LifecycleControlScope.PLATFORM,
],
instance_uuid=execution_context.instance_uuid,
workspace_uuid=execution_context.workspace_uuid,
placement_generation=execution_context.placement_generation,
)
except Exception:
process_coro.close()
self.semaphore.release()
await self._drop_selected_query(selected_query, selected_session)
raise
except asyncio.CancelledError:
raise
except Exception as e:
self.ap.logger.error(f'控制器循环出错: {e}')
self.ap.logger.error(f'Traceback: {traceback.format_exc()}')
# A persistent external failure must not turn this recovery
# loop into a CPU spin.
await asyncio.sleep(1)
async def run(self):
"""运行控制器"""
src/langbot/pkg/pipeline/longtext/strategies/image.py
+60 -63
View File
@@ -1,9 +1,11 @@
from __future__ import annotations
import asyncio
import os
import base64
import time
import re
import uuid
from PIL import Image, ImageDraw, ImageFont
@@ -28,28 +30,34 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
)
async def process(self, message: str, query: pipeline_query.Query) -> list[platform_message.MessageComponent]:
img_path = self.text_to_image(
text_str=message,
save_as='temp/{}.png'.format(int(time.time())),
query=query,
)
def render() -> str:
render_id = f'{int(time.time())}-{uuid.uuid4().hex}'
img_path = f'temp/{render_id}.png'
compressed_path = f'temp/{render_id}-compressed.png'
try:
self.text_to_image(
text_str=message,
save_as=img_path,
query=query,
)
compressed_path, _ = self.compress_image(
img_path,
outfile=compressed_path,
)
with open(compressed_path, 'rb') as f:
return base64.b64encode(f.read()).decode('utf-8')
finally:
for path in {img_path, compressed_path}:
if os.path.exists(path):
os.remove(path)
compressed_path, size = self.compress_image(img_path, outfile='temp/{}_compressed.png'.format(int(time.time())))
with open(compressed_path, 'rb') as f:
img = f.read()
b64 = base64.b64encode(img)
# 删除图片
os.remove(img_path)
if os.path.exists(compressed_path):
os.remove(compressed_path)
# Font measurement, image rendering and compression are CPU-bound PIL
# work and must not block the shared asyncio loop for every tenant.
image_base64 = await asyncio.to_thread(render)
return [
platform_message.Image(
base64=b64.decode('utf-8'),
base64=image_base64,
)
]
@@ -126,6 +134,36 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
o_size = self.get_size(outfile)
return outfile, self.get_size(outfile)
def _split_text_lines(self, text_str: str, text_width: int, font) -> list[str]:
"""Split text while guaranteeing that every loop iteration advances."""
final_lines: list[str] = []
text_width = max(int(text_width), 1)
for line in text_str.replace('\t', ' ').split('\n'):
line_width = font.getlength(line)
if not line or line_width < text_width:
final_lines.append(line)
continue
rest_text = line
while rest_text:
line_width = max(font.getlength(rest_text), 1)
point = int(len(rest_text) * (text_width / line_width))
point = max(1, min(point, len(rest_text)))
for number, number_index in self.indexNumber(rest_text):
if number_index < point < number_index + len(number) and number_index != 0:
point = number_index
break
point = max(1, min(point, len(rest_text)))
final_lines.append(rest_text[:point])
rest_text = rest_text[point:]
if rest_text and font.getlength(rest_text) < text_width:
final_lines.append(rest_text)
break
return final_lines
def text_to_image(
self,
text_str: str,
@@ -133,50 +171,9 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
width=800,
query: pipeline_query.Query = None,
):
text_str = text_str.replace('\t', ' ')
# 分行
lines = text_str.split('\n')
# 计算并分割
final_lines = []
text_width = width - 80
self.ap.logger.debug('lines: {}, text_width: {}'.format(lines, text_width))
for line in lines:
# 如果长了就分割
line_width = self.get_font(query.pipeline_config['output']['long-text-processing']['font-path']).getlength(
line
)
self.ap.logger.debug('line_width: {}'.format(line_width))
if line_width < text_width:
final_lines.append(line)
continue
else:
rest_text = line
while True:
# 分割最前面的一行
point = int(len(rest_text) * (text_width / line_width))
# 检查断点是否在数字中间
numbers = self.indexNumber(rest_text)
for number in numbers:
if number[1] < point < number[1] + len(number[0]) and number[1] != 0:
point = number[1]
break
final_lines.append(rest_text[:point])
rest_text = rest_text[point:]
line_width = self.get_font(
query.pipeline_config['output']['long-text-processing']['font-path']
).getlength(rest_text)
if line_width < text_width:
final_lines.append(rest_text)
break
else:
continue
font = self.get_font(query.pipeline_config['output']['long-text-processing']['font-path'])
text_width = max(width - 80, 1)
final_lines = self._split_text_lines(text_str, text_width, font)
# 准备画布
img = Image.new('RGBA', (width, max(280, len(final_lines) * 35 + 65)), (255, 255, 255, 255))
draw = ImageDraw.Draw(img, mode='RGBA')
@@ -191,7 +188,7 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
(offset_x, offset_y + 35 * line_number),
final_line,
fill=(0, 0, 0),
font=self.get_font(query.pipeline_config['output']['long-text-processing']['font-path']),
font=font,
)
# 遍历此行,检查是否有emoji
idx_in_line = 0
src/langbot/pkg/pipeline/pipelinemgr.py
+107 -18
View File
@@ -496,7 +496,50 @@ class PipelineManager:
def __init__(self, ap: app.Application):
self.ap = ap
self.pipelines = []
self._pipelines_by_key: dict[
tuple[str, str, str],
RuntimePipeline,
] = {}
self._pipeline_keys_by_scope: dict[
tuple[str, str],
set[tuple[str, str, str]],
] = {}
self._scope_generations: dict[tuple[str, str], int] = {}
@property
def pipelines(self) -> list[RuntimePipeline]:
"""Compatibility view over the indexed runtime pipeline registry."""
return list(self._pipelines_by_key.values())
@pipelines.setter
def pipelines(self, pipelines: list[RuntimePipeline]) -> None:
self._pipelines_by_key = {}
self._pipeline_keys_by_scope = {}
for pipeline in pipelines:
context = pipeline.execution_context
pipeline_uuid = (
getattr(getattr(pipeline, 'pipeline_entity', None), 'uuid', None) or context.pipeline_uuid or ''
)
key = (
context.instance_uuid,
pipeline.workspace_uuid,
pipeline_uuid,
)
self._pipelines_by_key[key] = pipeline
self._pipeline_keys_by_scope.setdefault(key[:2], set()).add(key)
def _observe_execution_context(self, context: ExecutionContext) -> None:
scope = (context.instance_uuid, context.workspace_uuid)
previous_generation = self._scope_generations.get(scope)
if previous_generation is not None and context.placement_generation < previous_generation:
raise WorkspaceInvariantError('Pipeline runtime placement generation rolled back')
if previous_generation == context.placement_generation:
return
if previous_generation is not None:
for key in self._pipeline_keys_by_scope.pop(scope, ()):
self._pipelines_by_key.pop(key, None)
self._scope_generations[scope] = context.placement_generation
async def initialize(self):
self.stage_dict = {name: cls for name, cls in stage.preregistered_stages.items()}
@@ -506,7 +549,9 @@ class PipelineManager:
async def load_pipelines_from_db(self):
self.ap.logger.info('Loading pipelines from db...')
self.pipelines = []
self._pipelines_by_key = {}
self._pipeline_keys_by_scope = {}
self._scope_generations = {}
list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
@@ -530,6 +575,7 @@ class PipelineManager:
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
),
pipeline,
_binding_validated=True,
)
return
@@ -575,6 +621,8 @@ class PipelineManager:
pipeline_entity: persistence_pipeline.LegacyPipeline
| sqlalchemy.Row[persistence_pipeline.LegacyPipeline]
| dict,
*,
_binding_validated: bool = False,
):
if isinstance(pipeline_entity, sqlalchemy.Row):
pipeline_entity = persistence_pipeline.LegacyPipeline(**pipeline_entity._mapping)
@@ -584,10 +632,12 @@ class PipelineManager:
execution_context = self._normalize_execution_context(context, pipeline_entity.uuid)
if pipeline_entity.workspace_uuid != execution_context.workspace_uuid:
raise WorkspaceRequiredError('Pipeline entity Workspace does not match its runtime context')
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
if not _binding_validated:
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
self._observe_execution_context(execution_context)
coerce_pipeline_config(
pipeline_entity.config,
@@ -605,13 +655,27 @@ class PipelineManager:
for stage_container in stage_containers:
await stage_container.inst.initialize(pipeline_entity.config)
# Stage initialization can yield while a Workspace is being moved.
# Revalidate before publishing the runtime assembled above.
if not _binding_validated:
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
self._observe_execution_context(execution_context)
runtime_pipeline = RuntimePipeline(
self.ap,
pipeline_entity,
stage_containers,
execution_context,
)
self.pipelines.append(runtime_pipeline)
key = (
execution_context.instance_uuid,
execution_context.workspace_uuid,
pipeline_entity.uuid,
)
self._pipelines_by_key[key] = runtime_pipeline
self._pipeline_keys_by_scope.setdefault(key[:2], set()).add(key)
async def get_pipeline_by_uuid(
self,
@@ -619,13 +683,24 @@ class PipelineManager:
uuid: str,
) -> RuntimePipeline | None:
execution_context = self._normalize_execution_context(context, uuid)
for pipeline in self.pipelines:
if (
pipeline.workspace_uuid == execution_context.workspace_uuid
and pipeline.placement_generation == execution_context.placement_generation
and pipeline.pipeline_entity.uuid == uuid
):
return pipeline
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
self._observe_execution_context(execution_context)
key = (
execution_context.instance_uuid,
execution_context.workspace_uuid,
uuid,
)
pipeline = self._pipelines_by_key.get(key)
if pipeline is not None and pipeline.placement_generation == execution_context.placement_generation:
return pipeline
if not self._pipeline_keys_by_scope.get(key[:2]):
self._scope_generations.pop(
(execution_context.instance_uuid, execution_context.workspace_uuid),
None,
)
return None
async def remove_pipeline(
@@ -634,7 +709,21 @@ class PipelineManager:
uuid: str,
) -> None:
execution_context = self._normalize_execution_context(context, uuid)
for pipeline in self.pipelines:
if pipeline.workspace_uuid == execution_context.workspace_uuid and pipeline.pipeline_entity.uuid == uuid:
self.pipelines.remove(pipeline)
return
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
self._observe_execution_context(execution_context)
key = (
execution_context.instance_uuid,
execution_context.workspace_uuid,
uuid,
)
if self._pipelines_by_key.pop(key, None) is not None:
scope_keys = self._pipeline_keys_by_scope.get(key[:2])
if scope_keys is not None:
scope_keys.discard(key)
if not scope_keys:
self._pipeline_keys_by_scope.pop(key[:2], None)
self._scope_generations.pop(key[:2], None)
return
src/langbot/pkg/pipeline/plugin_diagnostics.py
+11 -3
View File
@@ -159,6 +159,16 @@ def _discard_query_state(query_key: int) -> None:
_QUERY_STATES.pop(query_key, None)
def discard_query_state(query: pipeline_query.Query) -> None:
"""Release all diagnostics retained for a query leaving the runtime pool."""
query_key = id(query)
state = _QUERY_STATES.get(query_key)
if state is not None and state.finalizer is not None:
state.finalizer.detach()
_discard_query_state(query_key)
def _discard_query_state_if_empty(query: pipeline_query.Query) -> None:
query_key = id(query)
state = _QUERY_STATES.get(query_key)
@@ -166,9 +176,7 @@ def _discard_query_state_if_empty(query: pipeline_query.Query) -> None:
return
if state.pending_by_chain_id or state.by_response_index:
return
if state.finalizer is not None:
state.finalizer.detach()
_discard_query_state(query_key)
discard_query_state(query)
def _get_response_sources(
src/langbot/pkg/pipeline/pool.py
+120
View File
@@ -13,6 +13,7 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.provider.session as provider_session
from ..api.http.context import ExecutionContext
from . import plugin_diagnostics
QueryCacheKey = tuple[str, str]
LegacyQueryKey = tuple[str, int]
@@ -35,6 +36,10 @@ class QueryNotFoundError(LookupError):
"""Raised when a query does not exist inside the requested Workspace."""
class QueryPoolCapacityError(RuntimeError):
"""Raised when no queued query can be discarded to admit new work."""
def _validate_execution_context(execution_context: ExecutionContext) -> None:
if not isinstance(execution_context, ExecutionContext):
raise ExecutionContextRequiredError('A trusted ExecutionContext is required')
@@ -120,15 +125,100 @@ class QueryPool:
def __init__(
self,
singleton_context_resolver: SingletonContextResolver | None = None,
*,
max_queries: int = 1000,
max_queries_per_workspace: int = 100,
):
if max_queries < 1:
raise ValueError('max_queries must be positive')
if max_queries_per_workspace < 1:
raise ValueError('max_queries_per_workspace must be positive')
if max_queries_per_workspace > max_queries:
raise ValueError('max_queries_per_workspace cannot exceed max_queries')
self.query_id_counter = 0
self.pool_lock = asyncio.Lock()
self.queries = []
self.cached_queries = {}
self.active_query_count_by_workspace: dict[str, int] = {}
self.legacy_query_index = {}
self.query_count_by_scope = {}
self.dropped_query_count_by_scope: dict[QueryCounterKey, int] = {}
self.condition = asyncio.Condition(self.pool_lock)
self._singleton_context_resolver = singleton_context_resolver
self.max_queries = max_queries
self.max_queries_per_workspace = max_queries_per_workspace
def _discard_queued_query_locked(
self,
*,
workspace_uuid: str | None = None,
) -> pipeline_query.Query | None:
"""Discard the oldest queued query from one scope and all indexes."""
for index, query in enumerate(self.queries):
execution_context = get_query_execution_context(query)
if workspace_uuid is not None and execution_context.workspace_uuid != workspace_uuid:
continue
self.queries.pop(index)
query_uuid = execution_context.query_uuid
if query_uuid is not None:
self.cached_queries.pop((execution_context.workspace_uuid, query_uuid), None)
self.legacy_query_index.pop((execution_context.workspace_uuid, query.query_id), None)
query_workspace_uuid = execution_context.workspace_uuid
remaining = self.active_query_count_by_workspace.get(query_workspace_uuid, 0) - 1
if remaining > 0:
self.active_query_count_by_workspace[query_workspace_uuid] = remaining
else:
self.active_query_count_by_workspace.pop(query_workspace_uuid, None)
plugin_diagnostics.discard_query_state(query)
counter_key = (
execution_context.instance_uuid,
execution_context.workspace_uuid,
execution_context.placement_generation,
)
self.dropped_query_count_by_scope[counter_key] = self.dropped_query_count_by_scope.get(counter_key, 0) + 1
return query
return None
def _admit_query_locked(self, workspace_uuid: str) -> None:
workspace_query_count = self.active_query_count_by_workspace.get(workspace_uuid, 0)
if workspace_query_count >= self.max_queries_per_workspace:
if self._discard_queued_query_locked(workspace_uuid=workspace_uuid) is None:
raise QueryPoolCapacityError(f'Workspace query capacity reached ({self.max_queries_per_workspace})')
if len(self.cached_queries) >= self.max_queries:
if self._discard_queued_query_locked() is None:
raise QueryPoolCapacityError(f'Global query capacity reached ({self.max_queries})')
def mark_query_running_locked(self, query: pipeline_query.Query) -> None:
"""Remove a scheduled query from the overload-discardable queue."""
if not self.pool_lock.locked():
raise RuntimeError('Query pool lock is required to schedule a query')
for index, queued_query in enumerate(self.queries):
if queued_query is query:
self.queries.pop(index)
return
raise QueryNotFoundError('Scheduled query is no longer queued')
def _make_scope_counter_room_locked(self, counter_key: QueryCounterKey) -> None:
"""Retain recent counters without pinning every historical Workspace."""
if counter_key in self.query_count_by_scope:
return
while len(self.query_count_by_scope) >= self.max_queries:
stale_key = next(
(
existing_key
for existing_key in self.query_count_by_scope
if self.active_query_count_by_workspace.get(existing_key[1], 0) <= 0
),
None,
)
if stale_key is None:
raise QueryPoolCapacityError('Query counter capacity reached while every Workspace is active')
self.query_count_by_scope.pop(stale_key, None)
self.dropped_query_count_by_scope.pop(stale_key, None)
async def resolve_execution_context(
self,
@@ -180,6 +270,7 @@ class QueryPool:
)
async with self.condition:
self._admit_query_locked(execution_context.workspace_uuid)
query_id = self.query_id_counter
initial_variables: dict[str, typing.Any] = {'_routed_by_rule': routed_by_rule}
if variables:
@@ -217,6 +308,9 @@ class QueryPool:
self.queries.append(query)
self.cached_queries[(execution_context.workspace_uuid, query_uuid)] = query
self.active_query_count_by_workspace[execution_context.workspace_uuid] = (
self.active_query_count_by_workspace.get(execution_context.workspace_uuid, 0) + 1
)
self.legacy_query_index[(execution_context.workspace_uuid, query_id)] = query_uuid
self.query_id_counter += 1
counter_key = (
@@ -224,6 +318,13 @@ class QueryPool:
execution_context.workspace_uuid,
execution_context.placement_generation,
)
# A Workspace has only one active placement. Drop obsolete
# generation counters so deployment churn cannot grow these maps.
for existing_key in tuple(self.query_count_by_scope):
if existing_key[:2] == counter_key[:2] and existing_key != counter_key:
self.query_count_by_scope.pop(existing_key, None)
self.dropped_query_count_by_scope.pop(existing_key, None)
self._make_scope_counter_room_locked(counter_key)
self.query_count_by_scope[counter_key] = self.query_count_by_scope.get(counter_key, 0) + 1
self.condition.notify_all()
return query
@@ -241,6 +342,19 @@ class QueryPool:
0,
)
def get_dropped_query_count(self, execution_context: ExecutionContext) -> int:
"""Return overload drops for one active placement scope."""
_validate_execution_context(execution_context)
return self.dropped_query_count_by_scope.get(
(
execution_context.instance_uuid,
execution_context.workspace_uuid,
execution_context.placement_generation,
),
0,
)
async def get_query(
self,
workspace_uuid: str,
@@ -290,6 +404,11 @@ class QueryPool:
if cached_query is not query:
return False
del self.cached_queries[cache_key]
remaining = self.active_query_count_by_workspace.get(execution_context.workspace_uuid, 0) - 1
if remaining > 0:
self.active_query_count_by_workspace[execution_context.workspace_uuid] = remaining
else:
self.active_query_count_by_workspace.pop(execution_context.workspace_uuid, None)
self.legacy_query_index.pop(
(execution_context.workspace_uuid, query.query_id),
None,
@@ -298,6 +417,7 @@ class QueryPool:
if queued_query is query:
self.queries.pop(index)
break
plugin_diagnostics.discard_query_state(query)
return True
async def __aenter__(self):
src/langbot/pkg/pipeline/process/handlers/chat.py
+38
View File
@@ -26,6 +26,25 @@ importutil.import_modules_in_pkg(runners)
class ChatMessageHandler(handler.MessageHandler):
def _response_limit(self, name: str, default: int) -> int:
instance_config = getattr(self.ap, 'instance_config', None)
data = getattr(instance_config, 'data', {})
if not isinstance(data, dict):
return default
value = data.get('system', {}).get('response_limits', {}).get(name, default)
try:
return max(int(value), 1)
except (TypeError, ValueError):
return default
def _check_response_size(
self,
result: provider_message.Message | provider_message.MessageChunk,
) -> None:
content = result.content
if isinstance(content, str) and len(content) > self._response_limit('max_generated_chars', 1024 * 1024):
raise RuntimeError('Provider response exceeds the configured limit')
async def handle(
self,
query: pipeline_query.Query,
@@ -87,6 +106,7 @@ class ChatMessageHandler(handler.MessageHandler):
query.user_message.content = [event_ctx.event.user_message_alter]
text_length = 0
runner = None
try:
is_stream = await query.adapter.is_stream_output_supported()
except AttributeError:
@@ -107,6 +127,7 @@ class ChatMessageHandler(handler.MessageHandler):
chunk_count = 0 # Track streaming chunks to reduce excessive logging
async for result in runner.run(query):
self._check_response_size(result)
result.resp_message_id = str(resp_message_id)
if query.resp_messages:
query.resp_messages.pop()
@@ -119,6 +140,11 @@ class ChatMessageHandler(handler.MessageHandler):
query.resp_messages.append(result)
chunk_count += 1
if chunk_count > self._response_limit(
'max_stream_chunks',
100_000,
):
raise RuntimeError('Provider stream exceeds the configured event limit')
# Only log every 10th chunk to reduce excessive logging during streaming
# This prevents memory overflow from thousands of log entries per conversation
# First chunk uses INFO level to confirm connection establishment
@@ -145,6 +171,7 @@ class ChatMessageHandler(handler.MessageHandler):
else:
async for result in runner.run(query):
self._check_response_size(result)
query.resp_messages.append(result)
summary = self.format_result_log(result)
@@ -159,6 +186,10 @@ class ChatMessageHandler(handler.MessageHandler):
query.session.using_conversation.messages.append(query.user_message)
query.session.using_conversation.messages.extend(query.resp_messages)
self.ap.sess_mgr.trim_conversation_messages(
query.session.using_conversation,
max_rounds=query.pipeline_config['ai']['local-agent'].get('max-round', 10),
)
except Exception as e:
error_info = f'{traceback.format_exc()}'
self.ap.logger.error(f'Conversation({query.query_id}) Request Failed: {error_info}')
@@ -181,6 +212,13 @@ class ChatMessageHandler(handler.MessageHandler):
debug_notice=traceback.format_exc(),
)
finally:
if runner is not None:
try:
close_runner = getattr(runner, 'aclose', None)
if close_runner is not None:
await close_runner()
except Exception as ex:
self.ap.logger.warning(f'Failed to close request runner: {ex}')
# Telemetry reporting: collect minimal per-query execution info and send asynchronously
try:
end_ts = time.time()
src/langbot/pkg/pipeline/ratelimit/algos/fixedwin.py
+55 -4
View File
@@ -1,9 +1,17 @@
from __future__ import annotations
import asyncio
from collections import OrderedDict
import time
import typing
from .. import algo
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from ...pool import get_query_execution_context
_MAX_SESSION_CONTAINERS = 10000
_MIN_CONTAINER_TTL_SECONDS = 300
_CLEANUP_INTERVAL_SECONDS = 60
_MAX_EVICTION_PROBES = 64
# 固定窗口算法
@@ -13,9 +21,11 @@ class SessionContainer:
records: dict[int, int]
"""访问记录，key为每窗口长度的起始时间戳，value为访问次数"""
def __init__(self):
def __init__(self, ttl_seconds: int = _MIN_CONTAINER_TTL_SECONDS):
self.wait_lock = asyncio.Lock()
self.records = {}
self.last_accessed = time.monotonic()
self.ttl_seconds = ttl_seconds
@algo.algo_class('fixwin')
@@ -28,7 +38,8 @@ class FixedWindowAlgo(algo.ReteLimitAlgo):
async def initialize(self):
self.containers_lock = asyncio.Lock()
self.containers = {}
self.containers = OrderedDict()
self._last_cleanup = time.monotonic()
async def require_access(
self,
@@ -39,14 +50,53 @@ class FixedWindowAlgo(algo.ReteLimitAlgo):
# 加锁，找容器
container: SessionContainer = None
session_name = f'{launcher_type}_{launcher_id}'
execution_context = get_query_execution_context(query)
session_name = ':'.join(
(
execution_context.instance_uuid,
execution_context.workspace_uuid,
str(execution_context.placement_generation),
str(getattr(query, 'bot_uuid', '')),
str(getattr(query, 'pipeline_uuid', '')),
str(launcher_type),
str(launcher_id),
)
)
async with self.containers_lock:
container = self.containers.get(session_name)
if container is None:
container = SessionContainer()
window_size = query.pipeline_config['safety']['rate-limit']['window-length']
ttl_seconds = max(int(window_size) * 2, _MIN_CONTAINER_TTL_SECONDS)
now_monotonic = time.monotonic()
if now_monotonic - self._last_cleanup >= _CLEANUP_INTERVAL_SECONDS:
self._last_cleanup = now_monotonic
for key, candidate in tuple(self.containers.items()):
if (
not candidate.wait_lock.locked()
and now_monotonic - candidate.last_accessed >= candidate.ttl_seconds
):
self.containers.pop(key, None)
if len(self.containers) >= _MAX_SESSION_CONTAINERS:
for _ in range(min(_MAX_EVICTION_PROBES, len(self.containers))):
oldest_key = next(iter(self.containers))
oldest = self.containers[oldest_key]
if oldest.wait_lock.locked():
self.containers.move_to_end(oldest_key)
continue
self.containers.pop(oldest_key, None)
break
if len(self.containers) >= _MAX_SESSION_CONTAINERS:
# Every retained session is actively waiting. Reject this
# admission instead of growing an attacker-controlled map.
return False
container = SessionContainer(ttl_seconds=ttl_seconds)
self.containers[session_name] = container
else:
self.containers.move_to_end(session_name)
container.last_accessed = time.monotonic()
# 等待锁
async with container.wait_lock:
@@ -87,6 +137,7 @@ class FixedWindowAlgo(algo.ReteLimitAlgo):
container.records[now] = count + 1
# 返回True
container.last_accessed = time.monotonic()
return True
async def release_access(
src/langbot/pkg/pipeline/resprule/rules/regexp.py
+15 -12
View File
@@ -1,10 +1,8 @@
import re
from .. import rule as rule_model
from .. import entities
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from ....utils.safe_regex import SafeRegexError, matches_any
@rule_model.rule_class('regexp')
@@ -16,15 +14,20 @@ class RegExpRule(rule_model.GroupRespondRule):
rule_dict: dict,
query: pipeline_query.Query,
) -> entities.RuleJudgeResult:
regexps = rule_dict['regexp']
try:
matching = await matches_any(
rule_dict['regexp'],
message_text,
mode='match',
)
except SafeRegexError as exc:
self.ap.logger.warning(f'Group response regex rejected: {exc}')
matching = False
for regexp in regexps:
match = re.match(regexp, message_text)
if match:
return entities.RuleJudgeResult(
matching=True,
replacement=message_chain,
)
if matching:
return entities.RuleJudgeResult(
matching=True,
replacement=message_chain,
)
return entities.RuleJudgeResult(matching=False, replacement=message_chain)
src/langbot/pkg/platform/botmgr.py
+373 -106
View File
@@ -1,10 +1,12 @@
from __future__ import annotations
import asyncio
import contextlib
import dataclasses
import functools
import json
import re
import time
import traceback
import uuid
import sqlalchemy
@@ -20,6 +22,7 @@ from ..entity.persistence import workspace as persistence_workspace
from ..entity.errors import platform as platform_errors
from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType, RequestContext
from ..api.http.authz import WorkspaceRequiredError
from ..workspace.errors import WorkspaceInvariantError
from .logger import EventLogger
@@ -40,7 +43,7 @@ class RuntimeBot:
adapter: abstract_platform_adapter.AbstractMessagePlatformAdapter
task_wrapper: taskmgr.TaskWrapper
task_wrapper: taskmgr.TaskWrapper | None
task_context: taskmgr.TaskContext
@@ -80,7 +83,10 @@ class RuntimeBot:
self.enable = bot_entity.enable
self.adapter = adapter
self.task_context = taskmgr.TaskContext()
self.task_wrapper = None
self.logger = logger
self._shutdown_lock = asyncio.Lock()
self._shutdown_complete = False
async def assert_execution_active(self) -> None:
"""Fail closed when this long-lived adapter belongs to a stale placement."""
@@ -490,12 +496,27 @@ class RuntimeBot:
core_entities.LifecycleControlScope.APPLICATION,
core_entities.LifecycleControlScope.PLATFORM,
],
instance_uuid=self.execution_context.instance_uuid,
workspace_uuid=self.execution_context.workspace_uuid,
placement_generation=(self.execution_context.placement_generation),
)
async def shutdown(self):
await self.adapter.kill()
async with self._shutdown_lock:
if self._shutdown_complete:
return
self.ap.task_mgr.cancel_task(self.task_wrapper.id)
wrapper = self.task_wrapper
self.task_wrapper = None
try:
await asyncio.wait_for(self.adapter.kill(), timeout=15)
finally:
if wrapper is not None:
self.ap.task_mgr.cancel_task(wrapper.id)
if wrapper.task is not asyncio.current_task():
with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError):
await asyncio.wait_for(wrapper.task, timeout=5)
self._shutdown_complete = True
# 控制QQ消息输入输出的类
@@ -513,10 +534,140 @@ class PlatformManager:
def __init__(self, ap: app.Application = None):
self.ap = ap
self.bots = []
self._bots_by_key: dict[tuple[str, str, str], RuntimeBot] = {}
self._bot_keys_by_workspace: dict[
str,
set[tuple[str, str, str]],
] = {}
self._bot_keys_by_uuid: dict[str, set[tuple[str, str, str]]] = {}
self.websocket_proxy_bots = {}
self.adapter_components = []
self.adapter_dict = {}
self._scope_generations: dict[tuple[str, str], int] = {}
self._proxy_last_accessed: dict[str, float] = {}
self._runtime_mutation_lock = asyncio.Lock()
@staticmethod
def _runtime_bot_key(bot: RuntimeBot) -> tuple[str, str, str]:
context = getattr(bot, 'execution_context', None)
instance_uuid = str(getattr(context, 'instance_uuid', '__test_instance__'))
workspace_uuid = str(
getattr(bot, 'workspace_uuid', None) or getattr(context, 'workspace_uuid', '__test_workspace__')
)
bot_uuid = str(
getattr(getattr(bot, 'bot_entity', None), 'uuid', None)
or getattr(context, 'bot_uuid', None)
or f'__runtime_{id(bot)}'
)
return instance_uuid, workspace_uuid, bot_uuid
def _register_runtime_bot(self, bot: RuntimeBot) -> RuntimeBot | None:
key = self._runtime_bot_key(bot)
previous = self._bots_by_key.get(key)
self._bots_by_key[key] = bot
self._bot_keys_by_workspace.setdefault(key[1], set()).add(key)
self._bot_keys_by_uuid.setdefault(key[2], set()).add(key)
return previous
def _pop_runtime_bot(
self,
key: tuple[str, str, str],
) -> RuntimeBot | None:
bot = self._bots_by_key.pop(key, None)
if bot is None:
return None
workspace_keys = self._bot_keys_by_workspace.get(key[1])
if workspace_keys is not None:
workspace_keys.discard(key)
if not workspace_keys:
self._bot_keys_by_workspace.pop(key[1], None)
uuid_keys = self._bot_keys_by_uuid.get(key[2])
if uuid_keys is not None:
uuid_keys.discard(key)
if not uuid_keys:
self._bot_keys_by_uuid.pop(key[2], None)
return bot
@property
def bots(self) -> list[RuntimeBot]:
"""Compatibility view over the indexed platform runtime registry."""
return list(self._bots_by_key.values())
@bots.setter
def bots(self, bots: list[RuntimeBot]) -> None:
self._bots_by_key = {}
self._bot_keys_by_workspace = {}
self._bot_keys_by_uuid = {}
for bot in bots:
self._register_runtime_bot(bot)
def _max_workspace_proxies(self) -> int:
instance_data = getattr(
getattr(self.ap, 'instance_config', None),
'data',
{},
)
value = instance_data.get('system', {}).get('websocket_retention', {}).get('max_workspace_proxies', 1024)
try:
return max(int(value), 1)
except (TypeError, ValueError):
return 1024
async def _evict_idle_websocket_proxy_unlocked(self) -> None:
"""Make room without interrupting a live socket or in-flight query."""
if len(self.websocket_proxy_bots) < self._max_workspace_proxies():
return
from .sources.websocket_manager import WebSocketScope, ws_connection_manager
for workspace_uuid in sorted(
self.websocket_proxy_bots,
key=lambda item: self._proxy_last_accessed.get(item, 0.0),
):
proxy_bot = self.websocket_proxy_bots[workspace_uuid]
listener_tasks = getattr(proxy_bot.adapter, 'inbound_listener_tasks', ())
if any(not task.done() for task in tuple(listener_tasks)):
continue
scope = WebSocketScope.from_context(proxy_bot.execution_context)
if ws_connection_manager.get_stats(scope=scope)['total_connections'] > 0:
continue
self.websocket_proxy_bots.pop(workspace_uuid, None)
self._proxy_last_accessed.pop(workspace_uuid, None)
await proxy_bot.shutdown()
if not self._bot_keys_by_workspace.get(workspace_uuid):
self._scope_generations.pop(
(proxy_bot.execution_context.instance_uuid, workspace_uuid),
None,
)
return
raise RuntimeError('WebSocket Workspace proxy capacity reached and every proxy is active')
async def _observe_execution_context(self, context: ExecutionContext) -> None:
"""Shutdown superseded Workspace adapters when placement advances."""
async with self._runtime_mutation_lock:
await self._observe_execution_context_unlocked(context)
async def _observe_execution_context_unlocked(self, context: ExecutionContext) -> None:
scope = (context.instance_uuid, context.workspace_uuid)
previous_generation = self._scope_generations.get(scope)
if previous_generation is not None and context.placement_generation < previous_generation:
raise WorkspaceInvariantError('Platform runtime placement generation rolled back')
if previous_generation == context.placement_generation:
return
if previous_generation is not None:
proxy_bot = self.websocket_proxy_bots.pop(context.workspace_uuid, None)
self._proxy_last_accessed.pop(context.workspace_uuid, None)
if proxy_bot is not None and proxy_bot.enable:
await proxy_bot.shutdown()
for key in tuple(self._bot_keys_by_workspace.get(context.workspace_uuid, ())):
bot = self._pop_runtime_bot(key)
if bot is None:
continue
if bot.enable:
await bot.shutdown()
self._scope_generations[scope] = context.placement_generation
async def initialize(self):
# delete all bot log images
@@ -606,47 +757,60 @@ class PlatformManager:
context: ExecutionContext | RequestContext,
) -> RuntimeBot:
execution_context = self._normalize_execution_context(context)
existing = self.websocket_proxy_bots.get(execution_context.workspace_uuid)
if existing is not None:
if existing.placement_generation != execution_context.placement_generation:
raise WorkspaceRequiredError('WebSocket proxy placement generation is stale')
return existing
binding = await self.ap.workspace_service.get_execution_binding(
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
websocket_adapter_class = self.adapter_dict['websocket']
websocket_logger = EventLogger(
name='websocket-adapter',
ap=self.ap,
execution_context=execution_context,
owner='websocket-proxy-bot',
)
websocket_adapter_inst = websocket_adapter_class({}, websocket_logger, ap=self.ap)
proxy_context = dataclasses.replace(
execution_context,
instance_uuid=binding.instance_uuid,
bot_uuid='websocket-proxy-bot',
)
runtime_bot = RuntimeBot(
ap=self.ap,
bot_entity=persistence_bot.Bot(
uuid='websocket-proxy-bot',
workspace_uuid=binding.workspace_uuid,
name='WebSocket',
description='',
adapter='websocket',
adapter_config={},
enable=True,
),
adapter=websocket_adapter_inst,
logger=websocket_logger,
execution_context=proxy_context,
)
await runtime_bot.initialize()
self.websocket_proxy_bots[binding.workspace_uuid] = runtime_bot
return runtime_bot
async with self._runtime_mutation_lock:
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
await self._observe_execution_context_unlocked(execution_context)
existing = self.websocket_proxy_bots.get(execution_context.workspace_uuid)
if existing is not None:
if existing.placement_generation != execution_context.placement_generation:
raise WorkspaceRequiredError('WebSocket proxy placement generation is stale')
self._proxy_last_accessed[execution_context.workspace_uuid] = time.monotonic()
return existing
await self._evict_idle_websocket_proxy_unlocked()
binding = await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
websocket_adapter_class = self.adapter_dict['websocket']
websocket_logger = EventLogger(
name='websocket-adapter',
ap=self.ap,
execution_context=execution_context,
owner='websocket-proxy-bot',
)
websocket_adapter_inst = websocket_adapter_class({}, websocket_logger, ap=self.ap)
proxy_context = dataclasses.replace(
execution_context,
instance_uuid=binding.instance_uuid,
bot_uuid='websocket-proxy-bot',
)
runtime_bot = RuntimeBot(
ap=self.ap,
bot_entity=persistence_bot.Bot(
uuid='websocket-proxy-bot',
workspace_uuid=binding.workspace_uuid,
name='WebSocket',
description='',
adapter='websocket',
adapter_config={},
enable=True,
),
adapter=websocket_adapter_inst,
logger=websocket_logger,
execution_context=proxy_context,
)
await runtime_bot.initialize()
self.websocket_proxy_bots[binding.workspace_uuid] = runtime_bot
self._proxy_last_accessed[binding.workspace_uuid] = time.monotonic()
return runtime_bot
def get_running_adapters(
self,
@@ -654,13 +818,52 @@ class PlatformManager:
) -> list[abstract_platform_adapter.AbstractMessagePlatformAdapter]:
execution_context = self._normalize_execution_context(context)
return [
bot.adapter for bot in self.bots if bot.enable and bot.workspace_uuid == execution_context.workspace_uuid
bot.adapter
for bot in self.bots
if bot.enable
and bot.workspace_uuid == execution_context.workspace_uuid
and bot.placement_generation == execution_context.placement_generation
]
async def load_bots_from_db(self):
self.ap.logger.info('Loading bots from db...')
self.bots = []
async with self._runtime_mutation_lock:
old_bots = [*self.websocket_proxy_bots.values(), *self.bots]
self.websocket_proxy_bots = {}
self._proxy_last_accessed = {}
self.bots = []
self._scope_generations = {}
for bot in old_bots:
if not bot.enable:
continue
try:
await bot.shutdown()
except Exception as exc:
self.ap.logger.warning(f'Failed to stop old platform runtime during reload: {exc}')
list_bindings = getattr(
self.ap.workspace_service,
'list_active_execution_bindings',
None,
)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
if not callable(list_bindings) or not callable(tenant_uow):
raise RuntimeError('Cloud platform loading requires explicit instance discovery and tenant UoWs')
for binding in await list_bindings():
try:
async with tenant_uow(binding.workspace_uuid):
await self._load_workspace_bots(
binding.workspace_uuid,
_binding=binding,
)
except Exception as exc:
self.ap.logger.error(
f'Failed to load Workspace bots for {binding.workspace_uuid}: {exc}\n{traceback.format_exc()}'
)
return
instance_uow = getattr(self.ap.persistence_mgr, 'instance_discovery_uow', None)
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
@@ -698,15 +901,22 @@ class PlatformManager:
f'Failed to load Workspace bots for {workspace_uuid}: {e}\n{traceback.format_exc()}'
)
async def _load_workspace_bots(self, workspace_uuid: str) -> None:
async def _load_workspace_bots(
self,
workspace_uuid: str,
*,
_binding=None,
) -> None:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_bot.Bot)
.where(persistence_bot.Bot.workspace_uuid == workspace_uuid)
.order_by(persistence_bot.Bot.uuid)
)
binding = _binding
for bot in result.all():
try:
binding = await self.ap.workspace_service.get_execution_binding(workspace_uuid)
if binding is None:
binding = await self.ap.workspace_service.get_execution_binding(workspace_uuid)
execution_context = ExecutionContext(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
@@ -714,7 +924,11 @@ class PlatformManager:
bot_uuid=bot.uuid,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
)
await self.load_bot(execution_context, bot)
await self.load_bot(
execution_context,
bot,
_binding_validated=True,
)
except platform_errors.AdapterNotFoundError as e:
self.ap.logger.warning(f'Adapter {e.adapter_name} not found, skipping bot {bot.uuid}')
except Exception as e:
@@ -724,6 +938,8 @@ class PlatformManager:
self,
context: ExecutionContext | RequestContext,
bot_entity: persistence_bot.Bot | sqlalchemy.Row[persistence_bot.Bot] | dict,
*,
_binding_validated: bool = False,
) -> RuntimeBot:
"""加载机器人"""
if isinstance(bot_entity, sqlalchemy.Row):
@@ -734,45 +950,63 @@ class PlatformManager:
execution_context = self._normalize_execution_context(context, bot_uuid=bot_entity.uuid)
if bot_entity.workspace_uuid != execution_context.workspace_uuid:
raise WorkspaceRequiredError('Bot entity Workspace does not match its runtime context')
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
if not _binding_validated:
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
async with self._runtime_mutation_lock:
if not _binding_validated:
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
await self._observe_execution_context_unlocked(execution_context)
logger = EventLogger(
name=f'platform-adapter-{bot_entity.name}',
ap=self.ap,
execution_context=execution_context,
owner=bot_entity.uuid,
)
logger = EventLogger(
name=f'platform-adapter-{bot_entity.name}',
ap=self.ap,
execution_context=execution_context,
owner=bot_entity.uuid,
)
if bot_entity.adapter not in self.adapter_dict:
raise platform_errors.AdapterNotFoundError(bot_entity.adapter)
if bot_entity.adapter not in self.adapter_dict:
raise platform_errors.AdapterNotFoundError(bot_entity.adapter)
adapter_inst = self.adapter_dict[bot_entity.adapter](
bot_entity.adapter_config,
logger,
)
if hasattr(adapter_inst, 'ap'):
adapter_inst.ap = self.ap
adapter_inst = self.adapter_dict[bot_entity.adapter](
bot_entity.adapter_config,
logger,
)
if hasattr(adapter_inst, 'ap'):
adapter_inst.ap = self.ap
# 如果 adapter 支持 set_bot_uuid 方法，设置 bot_uuid（用于统一 webhook）
if hasattr(adapter_inst, 'set_bot_uuid'):
adapter_inst.set_bot_uuid(bot_entity.uuid)
# 如果 adapter 支持 set_bot_uuid 方法，设置 bot_uuid（用于统一 webhook）
if hasattr(adapter_inst, 'set_bot_uuid'):
adapter_inst.set_bot_uuid(bot_entity.uuid)
runtime_bot = RuntimeBot(
ap=self.ap,
bot_entity=bot_entity,
adapter=adapter_inst,
logger=logger,
execution_context=execution_context,
)
runtime_bot = RuntimeBot(
ap=self.ap,
bot_entity=bot_entity,
adapter=adapter_inst,
logger=logger,
execution_context=execution_context,
)
await runtime_bot.initialize()
await runtime_bot.initialize()
self.bots.append(runtime_bot)
bot_key = self._runtime_bot_key(runtime_bot)
existing_bot = self._bots_by_key.get(bot_key)
if existing_bot is not None and existing_bot is not runtime_bot:
try:
if existing_bot.enable:
await existing_bot.shutdown()
except BaseException:
if runtime_bot.enable:
await runtime_bot.shutdown()
raise
self._register_runtime_bot(runtime_bot)
return runtime_bot
return runtime_bot
async def get_bot_by_uuid(
self,
@@ -780,19 +1014,26 @@ class PlatformManager:
bot_uuid: str,
) -> RuntimeBot | None:
execution_context = self._normalize_execution_context(context, bot_uuid=bot_uuid)
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
await self._observe_execution_context(execution_context)
proxy_bot = self.websocket_proxy_bots.get(execution_context.workspace_uuid)
if proxy_bot and proxy_bot.bot_entity.uuid == bot_uuid:
if proxy_bot.placement_generation != execution_context.placement_generation:
return None
return proxy_bot
for bot in self.bots:
if (
bot.workspace_uuid == execution_context.workspace_uuid
and bot.placement_generation == execution_context.placement_generation
and bot.bot_entity.uuid == bot_uuid
):
return bot
return None
bot = self._bots_by_key.get(
(
execution_context.instance_uuid,
execution_context.workspace_uuid,
bot_uuid,
)
)
if bot is None or bot.placement_generation != execution_context.placement_generation:
return None
return bot
async def resolve_public_bot(self, route_key: str) -> RuntimeBot | None:
"""Resolve an opaque public bot UUID without consulting request headers."""
@@ -801,16 +1042,22 @@ class PlatformManager:
normalized = str(uuid.UUID(route_key))
except (ValueError, AttributeError, TypeError):
return None
for bot in self.bots:
if bot.bot_entity.uuid == normalized:
try:
await self.ap.workspace_service.get_execution_binding(
bot.workspace_uuid,
expected_generation=bot.placement_generation,
)
except Exception:
return None
return bot
keys = tuple(self._bot_keys_by_uuid.get(normalized, ()))
if len(keys) != 1:
return None
key = keys[0]
bot = self._bots_by_key.get(key)
if bot is None:
return None
try:
await self.ap.workspace_service.get_execution_binding(
bot.workspace_uuid,
expected_generation=bot.placement_generation,
)
except Exception:
return None
if self._bots_by_key.get(key) is bot:
return bot
return None
async def remove_bot(
@@ -819,12 +1066,26 @@ class PlatformManager:
bot_uuid: str,
) -> None:
execution_context = self._normalize_execution_context(context, bot_uuid=bot_uuid)
for bot in self.bots[:]:
if bot.workspace_uuid == execution_context.workspace_uuid and bot.bot_entity.uuid == bot_uuid:
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
async with self._runtime_mutation_lock:
await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
await self._observe_execution_context_unlocked(execution_context)
key = (
execution_context.instance_uuid,
execution_context.workspace_uuid,
bot_uuid,
)
bot = self._bots_by_key.get(key)
if bot is not None and bot.placement_generation == execution_context.placement_generation:
if bot.enable:
await bot.shutdown()
self.bots.remove(bot)
return
self._pop_runtime_bot(key)
def get_available_adapters_info(self) -> list[dict]:
return [
@@ -853,10 +1114,16 @@ class PlatformManager:
await bot.run()
async def shutdown(self):
for proxy_bot in self.websocket_proxy_bots.values():
if proxy_bot.enable:
await proxy_bot.shutdown()
for bot in self.bots:
if bot.enable:
await bot.shutdown()
async with self._runtime_mutation_lock:
runtime_bots = [*self.websocket_proxy_bots.values(), *self.bots]
self.websocket_proxy_bots = {}
self._proxy_last_accessed = {}
self.bots = []
for bot in runtime_bots:
if not bot.enable:
continue
try:
await bot.shutdown()
except Exception as exc:
self.ap.logger.warning(f'Failed to stop platform runtime during shutdown: {exc}')
self.ap.task_mgr.cancel_by_scope(core_entities.LifecycleControlScope.PLATFORM)
src/langbot/pkg/platform/logger.py
+6 -1
View File
@@ -55,6 +55,7 @@ class EventLog(pydantic.BaseModel):
MAX_LOG_COUNT = 200
DELETE_COUNT_PER_TIME = 50
MAX_LOG_TEXT_CHARS = 20000
class EventLogger(abstract_platform_event_logger.AbstractEventLogger):
@@ -129,7 +130,7 @@ class EventLogger(abstract_platform_event_logger.AbstractEventLogger):
async def _truncate_logs(self):
if len(self.logs) > MAX_LOG_COUNT:
for i in range(DELETE_COUNT_PER_TIME):
for image_key in self.logs[i].images: # type: ignore
for image_key in self.logs[i].images or []:
await self.ap.storage_mgr.delete_scoped_object_key(
self.execution_context,
image_key,
@@ -147,6 +148,10 @@ class EventLogger(abstract_platform_event_logger.AbstractEventLogger):
):
try:
image_keys = []
text = str(text)
if len(text) > MAX_LOG_TEXT_CHARS:
marker = '\n[log truncated]'
text = text[: MAX_LOG_TEXT_CHARS - len(marker)] + marker
if images is None:
images = []
src/langbot/pkg/platform/sources/aiocqhttp.py
+32
View File
@@ -23,6 +23,7 @@ _GROUP_NAME_LOOKUP_TIMEOUT_SECONDS = 2
_GROUP_MEMBER_INFO_CACHE_TTL_SECONDS = 86400
_GROUP_MEMBER_INFO_NEGATIVE_CACHE_TTL_SECONDS = 600
_GROUP_MEMBER_INFO_LOOKUP_TIMEOUT_SECONDS = 2
_LOOKUP_CACHE_MAX = 4096
def _normalize_base64_payload(value: str) -> str:
@@ -366,6 +367,31 @@ class AiocqhttpEventConverter(abstract_platform_adapter.AbstractEventConverter):
tuple[typing.Union[int, str], typing.Union[int, str]], tuple[dict, float]
] = {}
self._group_member_info_negative_cache: dict[tuple[typing.Union[int, str], typing.Union[int, str]], float] = {}
self._last_cache_cleanup = 0.0
def _prune_caches(self, now: float) -> None:
caches = (
self._group_name_cache,
self._group_name_negative_cache,
self._group_member_info_cache,
self._group_member_info_negative_cache,
)
if now - self._last_cache_cleanup < 60 and all(len(cache) <= _LOOKUP_CACHE_MAX for cache in caches):
return
self._last_cache_cleanup = now
for cache in caches:
for key, value in tuple(cache.items()):
expires_at = value[1] if isinstance(value, tuple) else value
if expires_at <= now:
cache.pop(key, None)
while len(cache) > _LOOKUP_CACHE_MAX:
cache.pop(next(iter(cache)), None)
def clear(self) -> None:
self._group_name_cache.clear()
self._group_name_negative_cache.clear()
self._group_member_info_cache.clear()
self._group_member_info_negative_cache.clear()
@staticmethod
async def yiri2target(event: platform_events.MessageEvent, bot_account_id: int):
@@ -373,6 +399,7 @@ class AiocqhttpEventConverter(abstract_platform_adapter.AbstractEventConverter):
async def _get_group_name(self, group_id: typing.Union[int, str], bot=None) -> str:
now = time.monotonic()
self._prune_caches(now)
if group_id in self._group_name_cache:
group_name, expires_at = self._group_name_cache[group_id]
if expires_at > now:
@@ -408,6 +435,7 @@ class AiocqhttpEventConverter(abstract_platform_adapter.AbstractEventConverter):
bot=None,
) -> dict:
now = time.monotonic()
self._prune_caches(now)
cache_key = (group_id, user_id)
if cache_key in self._group_member_info_cache:
member_info, expires_at = self._group_member_info_cache[cache_key]
@@ -526,6 +554,8 @@ class AiocqhttpAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
return
self.on_websocket_connection_event_cache.append(event)
if len(self.on_websocket_connection_event_cache) > 100:
self.on_websocket_connection_event_cache.pop(0)
await self.logger.info(f'WebSocket connection established, bot id: {event.self_id}')
async def send_message(self, target_type: str, target_id: str, message: platform_message.MessageChain):
@@ -691,4 +721,6 @@ class AiocqhttpAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
async def kill(self) -> bool:
# Current issue: existing connection will not be closed
# self.should_shutdown = True
self.on_websocket_connection_event_cache.clear()
self.event_converter.clear()
return False
src/langbot/pkg/platform/sources/dingtalk.py
+70 -1
View File
@@ -5,6 +5,7 @@ import re
import traceback
import typing
import uuid
import time
from langbot.libs.dingtalk_api.dingtalkevent import DingTalkEvent
import langbot_plugin.api.entities.builtin.platform.message as platform_message
@@ -544,11 +545,59 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
bot=bot,
listeners={},
)
self._background_tasks: set[asyncio.Task] = set()
# Wire the card-action callback after super().__init__ so we can reference
# self.* — the client's handler stores this as a soft reference and reads
# it at fire time.
self.bot.card_action_callback = self._on_card_action
def _start_background_task(self, coro) -> bool:
"""Start one bounded adapter-side auxiliary task."""
background_tasks = getattr(self, '_background_tasks', None)
if background_tasks is None:
background_tasks = set()
object.__setattr__(self, '_background_tasks', background_tasks)
for task in tuple(background_tasks):
if task.done():
background_tasks.discard(task)
if len(background_tasks) >= 100:
coro.close()
return False
task = asyncio.create_task(coro)
background_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
background_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
return True
def _prune_card_state(self) -> None:
now = time.monotonic()
ttl_seconds = 1800
for card_id, state in tuple(self.card_state.items()):
if now - float(state.get('created_at', now)) <= ttl_seconds:
continue
self.card_state.pop(card_id, None)
for session_key, active_card_id in tuple(self.active_turn_card.items()):
if active_card_id == card_id:
self.active_turn_card.pop(session_key, None)
self.active_turn_text.pop(session_key, None)
while len(self.card_state) > 1000:
card_id = next(iter(self.card_state))
self.card_state.pop(card_id, None)
while len(self.active_turn_card) > 1000:
session_key = next(iter(self.active_turn_card))
self.active_turn_card.pop(session_key, None)
self.active_turn_text.pop(session_key, None)
card_instances = getattr(self, 'card_instance_id_dict', None)
if isinstance(card_instances, dict):
while len(card_instances) > 1000:
card_instances.pop(next(iter(card_instances)), None)
async def reply_message(
self,
message_source: platform_events.MessageEvent,
@@ -674,6 +723,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
return is_stream
async def create_message_card(self, message_id, event):
self._prune_card_state()
form_template_id = (self.config.get('human_input_card_template_id') or '').strip()
legacy_template_id = self.config.get('card_template_id', '')
@@ -806,6 +856,20 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
return params
async def kill(self) -> bool:
task_set = getattr(self, '_background_tasks', set())
background_tasks = list(task_set)
for task in background_tasks:
if not task.done():
task.cancel()
if background_tasks:
await asyncio.gather(*background_tasks, return_exceptions=True)
task_set.clear()
card_instances = getattr(self, 'card_instance_id_dict', None)
if isinstance(card_instances, dict):
card_instances.clear()
self.card_state.clear()
self.active_turn_card.clear()
self.active_turn_text.clear()
await self.bot.stop()
return True
@@ -931,6 +995,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
)
# Record form state for the click-handler.
self._prune_card_state()
launcher_type, launcher_id, sender_user_id = self._derive_session_descriptor(message_source)
self.card_state[out_track_id] = {
'session_key': session_key,
@@ -947,6 +1012,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
'current_input_field': str(form_data.get('_current_input_field') or ''),
'input_defs': _dingtalk_form_input_defs(form_data),
'inputs': form_data.get('inputs') or {},
'created_at': time.monotonic(),
}
btns = self._build_btns(actions if should_show_actions else [], out_track_id)
@@ -1040,6 +1106,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
f'options={len(component_params.get("select_options") or [])}'
)
self._prune_card_state()
self.card_state[out_track_id] = {
'session_key': session_key,
'launcher_type': launcher_type.value,
@@ -1057,6 +1124,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
'inputs': form_data.get('inputs') or {},
'open_space_id': open_space_id,
'is_group': is_group,
'created_at': time.monotonic(),
}
parts = []
@@ -1223,6 +1291,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
f'payload_action_id={payload.get("action_id")!r} params={payload.get("params")!r}'
)
out_track_id = payload.get('out_track_id') or ''
self._prune_card_state()
params = payload.get('params') or {}
# ButtonGroup `sendCardRequest` events surface the click id at the
# callback top level as `actionId`; fall back to `params.action_id`
@@ -1359,7 +1428,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# output lives on a separate new card (lazy-created in
# reply_message_chunk on the synthetic event), so the form card
# stays put as a record of the user's selection.
asyncio.create_task(
self._start_background_task(
self._mark_card_resolved(
out_track_id,
action_title,
src/langbot/pkg/platform/sources/discord.py
+70 -66
View File
@@ -28,6 +28,21 @@ import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_
from ..logger import EventLogger
_MAX_DISCORD_MEDIA_BYTES = 10 * 1024 * 1024
def _decode_discord_base64_limited(value: str) -> bytes:
if ',' in value:
value = value.split(',', 1)[1]
max_encoded_bytes = 4 * ((_MAX_DISCORD_MEDIA_BYTES + 2) // 3)
if len(value) > max_encoded_bytes:
raise ValueError('Discord media exceeds the size limit')
decoded = base64.b64decode(value)
if len(decoded) > _MAX_DISCORD_MEDIA_BYTES:
raise ValueError('Discord media exceeds the size limit')
return decoded
# 语音功能相关异常定义
class VoiceConnectionError(Exception):
"""语音连接基础异常"""
@@ -604,7 +619,6 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
for ele in message_chain:
if isinstance(ele, platform_message.Image):
image_bytes = None
filename = f'{uuid.uuid4()}.png' # 默认文件名
if ele.base64:
@@ -618,60 +632,17 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
filename = f'{uuid.uuid4()}.gif'
elif 'webp' in data_header:
filename = f'{uuid.uuid4()}.webp'
# 去掉data:image/xxx;base64,前缀
base64_data = ele.base64.split(',')[1]
else:
base64_data = ele.base64
image_bytes = base64.b64decode(base64_data)
elif ele.url:
# 从URL下载图片
session = httpclient.get_session()
async with session.get(ele.url) as response:
image_bytes = await response.read()
# 从URL或Content-Type推断文件类型
content_type = response.headers.get('Content-Type', '')
if 'jpeg' in content_type or 'jpg' in content_type:
filename = f'{uuid.uuid4()}.jpg'
elif 'gif' in content_type:
filename = f'{uuid.uuid4()}.gif'
elif 'webp' in content_type:
filename = f'{uuid.uuid4()}.webp'
elif ele.url.lower().endswith(('.jpg', '.jpeg')):
filename = f'{uuid.uuid4()}.jpg'
elif ele.url.lower().endswith('.gif'):
filename = f'{uuid.uuid4()}.gif'
elif ele.url.lower().endswith('.webp'):
filename = f'{uuid.uuid4()}.webp'
elif ele.path:
# 从文件路径读取图片
# 确保路径没有空字节
clean_path = ele.path.replace('\x00', '')
clean_path = os.path.abspath(clean_path)
if not os.path.exists(clean_path):
continue # 跳过不存在的文件
try:
with open(clean_path, 'rb') as f:
image_bytes = f.read()
# 从文件路径获取文件名，保持原始扩展名
original_filename = os.path.basename(clean_path)
if original_filename and '.' in original_filename:
# 保持原始文件名的扩展名
ext = original_filename.split('.')[-1].lower()
filename = f'{uuid.uuid4()}.{ext}'
else:
# 如果没有扩展名，尝试从文件内容检测
if image_bytes.startswith(b'\xff\xd8\xff'):
filename = f'{uuid.uuid4()}.jpg'
elif image_bytes.startswith(b'GIF'):
filename = f'{uuid.uuid4()}.gif'
elif image_bytes.startswith(b'RIFF') and b'WEBP' in image_bytes[:20]:
filename = f'{uuid.uuid4()}.webp'
# 默认保持PNG
except Exception as e:
print(f'Error reading image file {clean_path}: {e}')
continue # 跳过读取失败的文件
try:
image_bytes, mime_type = await ele.get_bytes()
except Exception as exc:
print(f'Error reading Discord image: {exc}')
continue
if 'jpeg' in mime_type or 'jpg' in mime_type:
filename = f'{uuid.uuid4()}.jpg'
elif 'gif' in mime_type:
filename = f'{uuid.uuid4()}.gif'
elif 'webp' in mime_type:
filename = f'{uuid.uuid4()}.webp'
if image_bytes:
files.append(discord.File(fp=io.BytesIO(image_bytes), filename=filename))
@@ -702,27 +673,34 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
elif 'webm' in data_header:
filename = f'{uuid.uuid4()}.webm'
file_base64 = ele.base64.split(',')[-1]
file_bytes = base64.b64decode(file_base64)
file_bytes = await asyncio.to_thread(
_decode_discord_base64_limited,
ele.base64,
)
elif ele.url:
session = httpclient.get_session()
async with session.get(ele.url) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_DISCORD_MEDIA_BYTES,
)
if file_bytes:
files.append(discord.File(fp=io.BytesIO(file_bytes), filename=filename))
elif isinstance(ele, platform_message.File):
file_bytes = None
filename = f'{uuid.uuid4()}.{ele.name.split(".")[-1]}'
if ele.base64:
if ele.base64.startswith('data:'):
file_base64 = ele.base64.split(',')[1]
file_bytes = base64.b64decode(file_base64)
else:
file_bytes = base64.b64decode(ele.base64)
file_bytes = await asyncio.to_thread(
_decode_discord_base64_limited,
ele.base64,
)
elif ele.url:
session = httpclient.get_session()
async with session.get(ele.url) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_DISCORD_MEDIA_BYTES,
)
if file_bytes:
files.append(discord.File(fp=io.BytesIO(file_bytes), filename=filename))
elif isinstance(ele, platform_message.Forward):
@@ -780,8 +758,11 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
for attachment in message.attachments:
session = httpclient.get_session(trust_env=True)
async with session.get(attachment.url) as response:
image_data = await response.read()
image_base64 = base64.b64encode(image_data).decode('utf-8')
image_data = await httpclient.read_limited(
response,
max_bytes=_MAX_DISCORD_MEDIA_BYTES,
)
image_base64 = (await asyncio.to_thread(base64.b64encode, image_data)).decode('utf-8')
image_format = response.headers['Content-Type']
element_list.append(platform_message.Image(base64=f'data:{image_format};base64,{image_base64}'))
@@ -968,6 +949,20 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# resume.
self._pending_forms: dict[str, dict] = {}
def _prune_transient_state(self) -> None:
now = time.time()
ttl_seconds = 1800
for message_id, state in tuple(self._stream_buffer.items()):
if now - float(state.get('updated_at', now)) > ttl_seconds:
self._stream_buffer.pop(message_id, None)
for session_key, state in tuple(self._pending_forms.items()):
if now - float(state.get('posted_at', now)) > ttl_seconds:
self._pending_forms.pop(session_key, None)
while len(self._stream_buffer) > 100:
self._stream_buffer.pop(next(iter(self._stream_buffer)), None)
while len(self._pending_forms) > 1000:
self._pending_forms.pop(next(iter(self._pending_forms)), None)
# Voice functionality methods
async def join_voice_channel(self, guild_id: int, channel_id: int, user_id: int = None) -> discord.VoiceClient:
"""
@@ -1246,11 +1241,13 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
source = event.source_platform_object
if not isinstance(source, discord.Message):
return False
self._prune_transient_state()
self._stream_buffer[message_id] = {
'channel': source.channel,
'sent_message': None, # discord.Message set on first send
'last_content': '',
'chunk_count': 0,
'updated_at': time.time(),
}
return True
@@ -1274,6 +1271,8 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
form_data = getattr(bot_message, '_form_data', None) if not isinstance(bot_message, dict) else None
ctx = self._stream_buffer.get(msg_id) if msg_id else None
if ctx is not None:
ctx['updated_at'] = time.time()
# If the stream ctx was not set up (create_message_card wasn't
# called, e.g. synthetic event), or the final chunk carries a
@@ -1342,6 +1341,7 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
callback synthesizes a ``_dify_form_action`` query so the runner's
``_merge_pending_form_action`` resumes the workflow.
"""
self._prune_transient_state()
source = message_source.source_platform_object
actions = form_data.get('actions') or []
@@ -1445,6 +1445,8 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
and disable the View buttons so the choice is visually locked in."""
import langbot_plugin.api.entities.builtin.provider.session as provider_session
self._prune_transient_state()
# ACK first (3-second deadline before Discord shows "interaction failed").
try:
await interaction.response.defer()
@@ -1653,5 +1655,7 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
if self.voice_manager:
await self.voice_manager.disconnect_all()
self._stream_buffer.clear()
self._pending_forms.clear()
await self.bot.close()
return True
src/langbot/pkg/platform/sources/http_bot.py
+89 -6
View File
@@ -54,6 +54,7 @@ _ERR = {
'bad_signature': (401, 40101),
'duplicate': (409, 40901),
'too_large': (413, 41301),
'overloaded': (503, 50301),
'internal': (500, 50001),
}
@@ -63,16 +64,21 @@ _MAX_BODY = 1 * 1024 * 1024
# Idempotency dedup window (seconds) and cap.
_IDEMPOTENCY_TTL = 600
_IDEMPOTENCY_MAX = 4096
_OUTBOUND_QUEUE_MAX = 100
_OUTBOUND_IDLE_SECONDS = 60
_OUTBOUND_STATE_MAX = 4096
_INBOUND_TASK_MAX = 100
class _SessionOutbound:
"""Per-session outbound state: ordered delivery queue + sequence counter."""
def __init__(self) -> None:
self.queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
self.queue: asyncio.Queue = asyncio.Queue(maxsize=_OUTBOUND_QUEUE_MAX)
self.worker: asyncio.Task | None = None
self.sequence: int = 0
self.last_was_final: bool = True # so the first reply of a turn starts at seq 1
self.last_active: float = time.monotonic()
class _SyncCollector:
@@ -99,6 +105,7 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
idempotency_cache: dict[str, float] = pydantic.Field(default_factory=dict, exclude=True)
# session_id -> sync collector (set while a /sync request is awaiting a turn)
sync_waiters: dict[str, '_SyncCollector'] = pydantic.Field(default_factory=dict, exclude=True)
inbound_tasks: set[asyncio.Task] = pydantic.Field(default_factory=set, exclude=True)
model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
@@ -108,6 +115,7 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.outbound_states = {}
self.idempotency_cache = {}
self.sync_waiters = {}
self.inbound_tasks = set()
# -- framework hooks ------------------------------------------------------
@@ -156,10 +164,19 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await asyncio.sleep(3600)
async def kill(self):
# Cancel any outbound workers.
tasks = list(self.inbound_tasks)
for task in tasks:
if not task.done():
task.cancel()
for state in self.outbound_states.values():
if state.worker and not state.worker.done():
state.worker.cancel()
tasks.append(state.worker)
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
self.outbound_states.clear()
self.sync_waiters.clear()
self.inbound_tasks.clear()
return True
# -- inbound --------------------------------------------------------------
@@ -177,6 +194,24 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
for k in expired:
self.idempotency_cache.pop(k, None)
def _start_inbound_task(self, coro: typing.Coroutine) -> asyncio.Task | None:
self.inbound_tasks = {task for task in self.inbound_tasks if not task.done()}
if len(self.inbound_tasks) >= _INBOUND_TASK_MAX:
coro.close()
return None
task = asyncio.create_task(coro)
self.inbound_tasks.add(task)
def task_done(done_task: asyncio.Task) -> None:
self.inbound_tasks.discard(done_task)
if not done_task.cancelled():
# Retrieve failures so fire-and-forget callbacks never emit
# "Task exception was never retrieved" or retain tracebacks.
done_task.exception()
task.add_done_callback(task_done)
return task
async def handle_unified_webhook(self, bot_uuid: str, path: str, request):
"""Handle an inbound POST from the unified webhook dispatcher.
@@ -213,7 +248,7 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
return None, self._err('bad_signature', f'invalid signature: {reason}')
try:
data = json.loads(body)
data = await asyncio.to_thread(json.loads, body)
except (json.JSONDecodeError, ValueError):
return None, self._err('bad_request', 'body is not valid JSON')
if not isinstance(data, dict):
@@ -282,7 +317,8 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
return await self._run_sync(event, listener, session_id, message_id)
# Fire-and-collect: kick the pipeline, return 202 immediately.
asyncio.create_task(listener(event, self))
if self._start_inbound_task(listener(event, self)) is None:
return self._err('overloaded', 'too many inbound messages are already being processed')
return quart.jsonify(
{
'code': 0,
@@ -361,7 +397,9 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
return ''
def _next_sequence(self, session_id: str, is_final: bool) -> int:
self._prune_outbound_states()
state = self.outbound_states.setdefault(session_id, _SessionOutbound())
state.last_active = time.monotonic()
if state.last_was_final:
state.sequence = 1
else:
@@ -369,8 +407,37 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
state.last_was_final = is_final
return state.sequence
def _prune_outbound_states(self) -> None:
now = time.monotonic()
removable = [
(session_id, state)
for session_id, state in self.outbound_states.items()
if (state.worker is None or state.worker.done())
and state.queue.empty()
and now - state.last_active >= _OUTBOUND_IDLE_SECONDS
]
for session_id, state in removable:
if self.outbound_states.get(session_id) is state:
self.outbound_states.pop(session_id, None)
overflow = len(self.outbound_states) - _OUTBOUND_STATE_MAX
if overflow <= 0:
return
idle = sorted(
(
(session_id, state)
for session_id, state in self.outbound_states.items()
if (state.worker is None or state.worker.done()) and state.queue.empty()
),
key=lambda item: item[1].last_active,
)
for session_id, state in idle[:overflow]:
if self.outbound_states.get(session_id) is state:
self.outbound_states.pop(session_id, None)
async def _enqueue_callback(self, session_id: str, payload: dict) -> None:
state = self.outbound_states.setdefault(session_id, _SessionOutbound())
state.last_active = time.monotonic()
if state.worker is None or state.worker.done():
state.worker = asyncio.create_task(self._outbound_worker(session_id, state))
try:
@@ -386,13 +453,23 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async def _outbound_worker(self, session_id: str, state: _SessionOutbound) -> None:
while True:
payload = await state.queue.get()
try:
payload = await asyncio.wait_for(
state.queue.get(),
timeout=_OUTBOUND_IDLE_SECONDS,
)
except asyncio.TimeoutError:
if self.outbound_states.get(session_id) is state and state.queue.empty():
self.outbound_states.pop(session_id, None)
return
continue
try:
await self._deliver_callback(payload)
except Exception as e: # noqa: BLE001
await self.logger.error(f'http_bot callback delivery failed for {session_id}: {e}')
finally:
state.queue.task_done()
state.last_active = time.monotonic()
async def _deliver_callback(self, payload: dict) -> None:
callback_url = self.config.get('callback_url', '')
@@ -508,8 +585,11 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
collector = _SyncCollector()
self.sync_waiters[session_id] = collector
listener_task = self._start_inbound_task(listener(event, self))
if listener_task is None:
self.sync_waiters.pop(session_id, None)
return self._err('overloaded', 'too many inbound messages are already being processed')
try:
asyncio.create_task(listener(event, self))
timeout = int(self.config.get('callback_timeout', 15)) * 4
try:
await asyncio.wait_for(collector.done.wait(), timeout=timeout)
@@ -517,6 +597,9 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await self.logger.warning(f'http_bot sync wait timed out for session {session_id}')
finally:
self.sync_waiters.pop(session_id, None)
state = self.outbound_states.get(session_id)
if state is not None and state.worker is None and state.queue.empty():
self.outbound_states.pop(session_id, None)
return quart.jsonify(
{
src/langbot/pkg/platform/sources/kook.py
+47 -31
View File
@@ -6,7 +6,6 @@ import json
import base64
import zlib
import traceback
import time
import aiohttp
@@ -21,6 +20,39 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
_KOOK_MAX_GATEWAY_MESSAGE_BYTES = 10 * 1024 * 1024
def _bounded_zlib_decompress(payload: bytes) -> bytes:
decompressor = zlib.decompressobj()
decoded = decompressor.decompress(
payload,
_KOOK_MAX_GATEWAY_MESSAGE_BYTES + 1,
)
if len(decoded) > _KOOK_MAX_GATEWAY_MESSAGE_BYTES or decompressor.unconsumed_tail:
raise ValueError('KOOK gateway message exceeds the decompressed size limit')
decoded += decompressor.flush(_KOOK_MAX_GATEWAY_MESSAGE_BYTES + 1 - len(decoded))
if len(decoded) > _KOOK_MAX_GATEWAY_MESSAGE_BYTES or not decompressor.eof:
raise ValueError('KOOK gateway message exceeds the decompressed size limit')
return decoded
def _decode_gateway_message(message: str | bytes) -> dict:
if isinstance(message, bytes):
try:
message_bytes = _bounded_zlib_decompress(message)
except zlib.error:
message_bytes = message
else:
message_bytes = message.encode('utf-8')
if len(message_bytes) > _KOOK_MAX_GATEWAY_MESSAGE_BYTES:
raise ValueError('KOOK gateway message exceeds the size limit')
decoded = json.loads(message_bytes)
if not isinstance(decoded, dict):
raise ValueError('KOOK gateway message must be a JSON object')
return decoded
class KookMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
"""Convert between LangBot MessageChain and KOOK message format"""
@@ -125,8 +157,8 @@ class KookMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
session = httpclient.get_session()
async with session.get(content) as response:
if response.status == 200:
image_bytes = await response.read()
image_base64 = base64.b64encode(image_bytes).decode('utf-8')
image_bytes = await httpclient.read_limited(response)
image_base64 = (await asyncio.to_thread(base64.b64encode, image_bytes)).decode('utf-8')
# Detect image format
content_type = response.headers.get('Content-Type', 'image/png')
components.append(
@@ -270,10 +302,6 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
http_session: typing.Optional[aiohttp.ClientSession] = pydantic.Field(exclude=True, default=None)
def __init__(self, config: dict, logger: abstract_platform_logger.AbstractEventLogger, **kwargs):
# Debug: Track init
with open('/tmp/kook_adapter_init.txt', 'w') as f:
f.write(f'KOOK adapter __init__ called at {time.time()}\n')
# Validate required config
if 'token' not in config:
raise Exception('KOOK adapter requires "token" in config')
@@ -300,7 +328,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
session = httpclient.get_session()
async with session.get(base_url, params=params, headers=headers) as response:
if response.status == 200:
data = await response.json()
data = await httpclient.read_json_limited(response)
if data.get('code') == 0:
gateway_url = data['data']['url']
return gateway_url
@@ -320,7 +348,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
session = httpclient.get_session()
async with session.get(base_url, headers=headers) as response:
if response.status == 200:
data = await response.json()
data = await httpclient.read_json_limited(response)
if data.get('code') == 0:
user_info = data['data']
return user_info
@@ -409,17 +437,10 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# Wait for HELLO within 6 seconds
try:
hello_msg = await asyncio.wait_for(ws.recv(), timeout=6.0)
# Handle compressed messages (same as main message loop)
if isinstance(hello_msg, bytes):
# Decompress if compressed
try:
hello_msg = zlib.decompress(hello_msg).decode('utf-8')
except Exception:
# Not compressed or decompression failed
hello_msg = hello_msg.decode('utf-8')
hello_data = json.loads(hello_msg)
hello_data = await asyncio.to_thread(
_decode_gateway_message,
hello_msg,
)
if hello_data.get('s') == 1: # HELLO signal
await self._handle_hello(hello_data['d'])
@@ -433,16 +454,11 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# Main message loop
async for message in ws:
if isinstance(message, bytes):
# Decompress if compressed
try:
message = zlib.decompress(message).decode('utf-8')
except Exception:
# Not compressed or decompression failed
message = message.decode('utf-8')
try:
msg_data = json.loads(message)
msg_data = await asyncio.to_thread(
_decode_gateway_message,
message,
)
signal = msg_data.get('s')
if signal == 0: # EVENT
@@ -516,7 +532,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async with self.http_session.post(url, json=payload, headers=headers) as response:
if response.status == 200:
result = await response.json()
result = await httpclient.read_json_limited(response)
if result.get('code') == 0:
await self.logger.debug(f'Message sent successfully to {target_id}')
else:
@@ -582,7 +598,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async with self.http_session.post(url, json=payload, headers=headers) as response:
if response.status == 200:
result = await response.json()
result = await httpclient.read_json_limited(response)
if result.get('code') == 0:
await self.logger.debug('Reply sent successfully')
else:
src/langbot/pkg/platform/sources/lark.py
+229 -119
View File
@@ -15,7 +15,7 @@ import hashlib
from Crypto.Cipher import AES
import tempfile
import os
import mimetypes
import threading
from langbot.pkg.utils import httpclient
import lark_oapi.ws.exception
@@ -34,6 +34,53 @@ import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_
import langbot_plugin.api.entities.builtin.provider.session as provider_session
_MAX_LARK_MEDIA_BYTES = 10 * 1024 * 1024
def _decode_lark_base64_limited(value: str) -> bytes:
if ',' in value:
value = value.split(',', 1)[1]
max_encoded_bytes = 4 * ((_MAX_LARK_MEDIA_BYTES + 2) // 3)
if len(value) > max_encoded_bytes:
raise ValueError('Lark media exceeds the size limit')
decoded = base64.b64decode(value)
if len(decoded) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
return decoded
def _read_lark_path_limited(path: str) -> bytes:
if os.path.getsize(path) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
with open(path, 'rb') as file:
body = file.read(_MAX_LARK_MEDIA_BYTES + 1)
if len(body) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
return body
def _write_lark_temp_file(data: bytes) -> str:
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
temp_file.write(data)
temp_file.flush()
return temp_file.name
def _read_lark_response_file_limited(response) -> bytes:
content_length = response.raw.headers.get('content-length')
if content_length is not None:
try:
if int(content_length) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
except (TypeError, ValueError) as exc:
if 'exceeds' in str(exc):
raise
body = response.file.read(_MAX_LARK_MEDIA_BYTES + 1)
if len(body) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
return body
def _lark_form_component_name(prefix: str, field_name: str, index: int) -> str:
safe_name = re.sub(r'[^A-Za-z0-9_]', '_', field_name)[:8] or 'field'
digest = hashlib.sha1(field_name.encode('utf-8')).hexdigest()[:6]
@@ -299,68 +346,33 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
@staticmethod
async def upload_image_to_lark(msg: platform_message.Image, api_client: lark_oapi.Client) -> typing.Optional[str]:
"""Upload an image to Lark and return the image_key, or None if upload fails."""
image_bytes = None
if msg.base64:
try:
# Remove data URL prefix if present
base64_data = msg.base64
if base64_data.startswith('data:'):
base64_data = base64_data.split(',', 1)[1]
image_bytes = base64.b64decode(base64_data)
except Exception as e:
print(f'Failed to decode base64 image: {e}')
traceback.print_exc()
return None
elif msg.url:
try:
session = httpclient.get_session()
async with session.get(msg.url) as response:
if response.status == 200:
image_bytes = await response.read()
else:
print(f'Failed to download image from {msg.url}: HTTP {response.status}')
return None
except Exception as e:
print(f'Failed to download image from {msg.url}: {e}')
traceback.print_exc()
return None
elif msg.path:
try:
with open(msg.path, 'rb') as f:
image_bytes = f.read()
except Exception as e:
print(f'Failed to read image from path {msg.path}: {e}')
traceback.print_exc()
return None
if image_bytes is None:
try:
image_bytes, _mime_type = await msg.get_bytes()
except Exception as exc:
print(f'Failed to load Lark image: {exc}')
traceback.print_exc()
return None
if not image_bytes:
print(
f'No image data available for Image message (url={msg.url}, base64={bool(msg.base64)}, path={msg.path})'
)
return None
try:
# Create a temporary file to store the image bytes
import tempfile
import os
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
temp_file.write(image_bytes)
temp_file.flush()
temp_file_path = temp_file.name
temp_file_path = await asyncio.to_thread(
_write_lark_temp_file,
image_bytes,
)
try:
# Create image request using the temporary file
request = (
CreateImageRequest.builder()
.request_body(
CreateImageRequestBody.builder().image_type('message').image(open(temp_file_path, 'rb')).build()
with open(temp_file_path, 'rb') as upload_file:
request = (
CreateImageRequest.builder()
.request_body(CreateImageRequestBody.builder().image_type('message').image(upload_file).build())
.build()
)
.build()
)
response = await api_client.im.v1.image.acreate(request)
response = await api_client.im.v1.image.acreate(request)
if not response.success():
print(
@@ -395,23 +407,24 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
duration: Duration in milliseconds (for audio files).
"""
try:
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
temp_file.write(file_bytes)
temp_file_path = temp_file.name
if len(file_bytes) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
temp_file_path = await asyncio.to_thread(
_write_lark_temp_file,
file_bytes,
)
try:
body_builder = (
CreateFileRequestBody.builder()
.file_type(file_type)
.file_name(file_name)
.file(open(temp_file_path, 'rb'))
)
if duration is not None:
body_builder = body_builder.duration(duration)
with open(temp_file_path, 'rb') as upload_file:
body_builder = (
CreateFileRequestBody.builder().file_type(file_type).file_name(file_name).file(upload_file)
)
if duration is not None:
body_builder = body_builder.duration(duration)
request = CreateFileRequest.builder().request_body(body_builder.build()).build()
request = CreateFileRequest.builder().request_body(body_builder.build()).build()
response = await api_client.im.v1.file.acreate(request)
response = await api_client.im.v1.file.acreate(request)
if not response.success():
print(
@@ -436,10 +449,10 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
if msg.base64:
try:
base64_str = msg.base64
if ',' in base64_str:
base64_str = base64_str.split(',', 1)[1]
data = base64.b64decode(base64_str)
data = await asyncio.to_thread(
_decode_lark_base64_limited,
msg.base64,
)
except Exception:
pass
elif msg.url:
@@ -447,13 +460,18 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
session = httpclient.get_session()
async with session.get(msg.url) as resp:
if resp.status == 200:
data = await resp.read()
data = await httpclient.read_limited(
resp,
max_bytes=_MAX_LARK_MEDIA_BYTES,
)
except Exception:
pass
elif msg.path:
try:
with open(msg.path, 'rb') as f:
data = f.read()
data = await asyncio.to_thread(
_read_lark_path_limited,
str(msg.path),
)
except Exception:
pass
@@ -694,8 +712,11 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
f'client.im.v1.message_resource.get failed, code: {response.code}, msg: {response.msg}, log_id: {response.get_log_id()}, resp: \n{json.dumps(json.loads(response.raw.content), indent=4, ensure_ascii=False)}'
)
image_bytes = response.file.read()
image_base64 = base64.b64encode(image_bytes).decode()
image_bytes = await asyncio.to_thread(
_read_lark_response_file_limited,
response,
)
image_base64 = (await asyncio.to_thread(base64.b64encode, image_bytes)).decode()
image_format = response.raw.headers['content-type']
@@ -721,27 +742,18 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
lb_msg_list.append(platform_message.Plain(text='[Audio file download failed]'))
return platform_message.MessageChain(lb_msg_list)
# Read audio bytes
audio_bytes = response.file.read()
audio_base64 = base64.b64encode(audio_bytes).decode()
audio_bytes = await asyncio.to_thread(
_read_lark_response_file_limited,
response,
)
audio_base64 = (await asyncio.to_thread(base64.b64encode, audio_bytes)).decode()
# Get content type from response headers
content_type = response.raw.headers.get('content-type', 'audio/mpeg')
mime_main = content_type.split(';')[0].strip()
ext = mimetypes.guess_extension(mime_main) or '.bin'
temp_dir = tempfile.gettempdir()
temp_file_path = os.path.join(temp_dir, f'lark_audio_{file_key}{ext}')
with open(temp_file_path, 'wb') as f:
f.write(audio_bytes)
# Create Voice message: prefer path/url + length, include base64 as optional data URI
lb_msg_list.append(
platform_message.Voice(
voice_id=file_key,
url=f'file://{temp_file_path}',
path=temp_file_path,
base64=f'data:{content_type};base64,{audio_base64}',
length=(duration // 1000) if duration else None,
)
@@ -770,40 +782,22 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
f'client.im.v1.message_resource.get failed, code: {response.code}, msg: {response.msg}, log_id: {response.get_log_id()}, resp: \n{json.dumps(json.loads(response.raw.content), indent=4, ensure_ascii=False)}'
)
file_bytes = response.file.read()
file_base64 = base64.b64encode(file_bytes).decode()
file_bytes = await asyncio.to_thread(
_read_lark_response_file_limited,
response,
)
file_base64 = (await asyncio.to_thread(base64.b64encode, file_bytes)).decode()
file_format = response.raw.headers['content-type']
file_size = len(file_bytes)
# Determine extension from content-type if possible
content_type = response.raw.headers.get('content-type', '')
mime_main = content_type.split(';')[0].strip() if content_type else ''
ext = mimetypes.guess_extension(mime_main) or ''
# Ensure a safe filename (avoid path components)
safe_name = os.path.basename(file_name).replace('/', '_').replace('\\', '_')
if ext and not safe_name.lower().endswith(ext.lower()):
filename_with_ext = f'{safe_name}{ext}'
else:
filename_with_ext = safe_name
temp_dir = tempfile.gettempdir()
temp_file_path = os.path.join(temp_dir, f'lark_{file_key}_{filename_with_ext}')
with open(temp_file_path, 'wb') as f:
f.write(file_bytes)
# Create File message with local path and file:// URL
lb_msg_list.append(
platform_message.File(
id=file_key,
name=file_name,
size=file_size,
url=f'file://{temp_file_path}',
path=temp_file_path,
base64=f'data:{file_format};base64,{file_base64}', # not including base64 by default to save memory; can be added if needed
base64=f'data:{file_format};base64,{file_base64}',
)
)
@@ -1042,16 +1036,26 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# card_id → input_defs / inputs captured for the selected-action notice
card_form_input_defs: dict[str, list[dict]]
card_form_inputs: dict[str, dict]
card_last_accessed: dict[str, float]
card_cleanup_at: float
# set of card_ids that have already transitioned from "buttons visible" to "resume layout"
card_resume_transitioned: set[str]
inbound_event_tasks: set[asyncio.Task]
threadsafe_event_futures: set[typing.Any]
threadsafe_event_lock: typing.Any = pydantic.Field(exclude=True)
_MONITORING_MAPPING_TTL = 600 # 10 minutes
_MAX_INBOUND_EVENTS = 100
_MAX_TENANT_ACCESS_TOKENS = 1024
seq: int # 用于在发送卡片消息中识别消息顺序，直接以seq作为标识
bot_uuid: str = None # 机器人UUID
app_ticket: str = None # 商店应用用到
app_access_token: str = None # 商店应用用到
app_access_token_expire_at: int = None
tenant_access_tokens: dict[str, dict[str, str]] = {} # 租户access_token映射
tenant_access_tokens: dict[str, dict[str, str]] = pydantic.Field(
default_factory=dict,
exclude=True,
)
def __init__(self, config: dict, logger: abstract_platform_logger.AbstractEventLogger, **kwargs):
quart_app = quart.Quart(__name__)
@@ -1062,11 +1066,11 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await self.listeners[type(lb_event)](lb_event, self)
def sync_on_message(event: lark_oapi.im.v1.P2ImMessageReceiveV1):
asyncio.create_task(on_message(event))
self._schedule_inbound_event(on_message(event))
def schedule_on_app_loop(coro):
"""Run a coroutine on the application event loop from sync callbacks."""
return asyncio.run_coroutine_threadsafe(coro, self.ap.event_loop)
return self._schedule_threadsafe_event(coro)
def sync_on_card_action(event):
try:
@@ -1289,7 +1293,13 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
card_form_content={},
card_form_input_defs={},
card_form_inputs={},
card_last_accessed={},
card_cleanup_at=0.0,
card_resume_transitioned=set(),
inbound_event_tasks=set(),
threadsafe_event_futures=set(),
threadsafe_event_lock=threading.Lock(),
tenant_access_tokens={},
seq=1,
listeners={},
quart_app=quart_app,
@@ -1300,6 +1310,45 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
**kwargs,
)
def _schedule_inbound_event(self, coro) -> None:
for task in tuple(self.inbound_event_tasks):
if task.done():
self.inbound_event_tasks.discard(task)
if len(self.inbound_event_tasks) >= self._MAX_INBOUND_EVENTS:
coro.close()
return
task = asyncio.create_task(coro)
self.inbound_event_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
self.inbound_event_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
def _schedule_threadsafe_event(self, coro):
"""Submit one bounded callback from the Lark SDK's sync boundary."""
with self.threadsafe_event_lock:
for future in tuple(self.threadsafe_event_futures):
if future.done():
self.threadsafe_event_futures.discard(future)
if len(self.threadsafe_event_futures) >= self._MAX_INBOUND_EVENTS:
coro.close()
return None
future = asyncio.run_coroutine_threadsafe(coro, self.ap.event_loop)
self.threadsafe_event_futures.add(future)
def done(done_future) -> None:
with self.threadsafe_event_lock:
self.threadsafe_event_futures.discard(done_future)
if not done_future.cancelled():
done_future.exception()
future.add_done_callback(done)
return future
def request_app_ticket(self, api_client, config):
app_id = config['app_id']
app_secret = config['app_secret']
@@ -1376,6 +1425,12 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
'token': tenant_access_token,
'expire_at': int(time.time()) + expire - 300,
}
now = int(time.time())
for cached_key, cached_token in tuple(self.tenant_access_tokens.items()):
if int(cached_token.get('expire_at', 0)) <= now:
self.tenant_access_tokens.pop(cached_key, None)
while len(self.tenant_access_tokens) > self._MAX_TENANT_ACCESS_TOKENS:
self.tenant_access_tokens.pop(next(iter(self.tenant_access_tokens)), None)
def get_tenant_access_token(self, tenant_key: str):
if tenant_key is None or 'isv' != self.config.get('app_type', 'self'):
@@ -1558,6 +1613,8 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
user_msg_id = query.message_event.message_chain.message_id
if user_msg_id:
self.pending_monitoring_msg[user_msg_id] = monitoring_message_id
while len(self.pending_monitoring_msg) > CARD_ID_CACHE_SIZE:
self.pending_monitoring_msg.pop(next(iter(self.pending_monitoring_msg)), None)
except Exception as e:
await self.logger.debug(f'Failed to map message to monitoring message: {e}')
@@ -1570,6 +1627,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
def _next_card_sequence(self, card_id: str, suggested: int = 1) -> int:
"""Return the next strictly increasing sequence for a card update."""
self._touch_card(card_id)
current = self.card_sequence_dict.get(card_id, 0)
next_seq = max(current + 1, suggested)
self.card_sequence_dict[card_id] = next_seq
@@ -1577,6 +1635,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
def _register_card_for_source(self, card_id: str, *source_ids: str) -> None:
"""Register a card_id under one or more source message ids."""
self._touch_card(card_id)
bucket = self.card_id_to_source_ids.setdefault(card_id, set())
for sid in source_ids:
if not sid:
@@ -1596,8 +1655,24 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.card_form_content.pop(card_id, None)
self.card_form_input_defs.pop(card_id, None)
self.card_form_inputs.pop(card_id, None)
self.card_last_accessed.pop(card_id, None)
self.card_resume_transitioned.discard(card_id)
def _touch_card(self, card_id: str) -> None:
now = time.monotonic()
if now - self.card_cleanup_at >= 60 or len(self.card_last_accessed) >= CARD_ID_CACHE_SIZE:
self.card_cleanup_at = now
for stale_card_id, last_accessed in tuple(self.card_last_accessed.items()):
if now - last_accessed >= CARD_ID_CACHE_MAX_LIFETIME:
self._drop_card_state(stale_card_id)
while len(self.card_last_accessed) >= CARD_ID_CACHE_SIZE:
oldest_card_id = min(
self.card_last_accessed,
key=self.card_last_accessed.__getitem__,
)
self._drop_card_state(oldest_card_id)
self.card_last_accessed[card_id] = now
async def create_card_id(self, message_id):
try:
# self.logger.debug('飞书支持stream输出,创建卡片......')
@@ -1793,6 +1868,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.card_id_dict[message_id] = response.data.card_id
card_id = response.data.card_id
self._touch_card(card_id)
self.card_sequence_dict[card_id] = 0
return card_id
@@ -1864,7 +1940,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.reply_to_monitoring_msg[reply_msg_id] = (monitoring_msg_id, time.time())
self._cleanup_monitoring_mapping()
except Exception as e:
asyncio.create_task(self.logger.debug(f'Failed to transfer monitoring mapping in create_message_card: {e}'))
await self.logger.debug(f'Failed to transfer monitoring mapping in create_message_card: {e}')
return True
@@ -2872,8 +2948,8 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
data = await request.json
if 'encrypt' in data:
data = self.cipher.decrypt_string(data['encrypt'])
data = json.loads(data)
encrypted = data['encrypt']
data = await asyncio.to_thread(lambda: json.loads(self.cipher.decrypt_string(encrypted)))
type = self.get_event_type(data)
context = EventContext(data)
if 'url_verification' == type:
@@ -3143,4 +3219,38 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# 所以要设置_auto_reconnect=False,让其不重连。
self.bot._auto_reconnect = False
await self.bot._disconnect()
inbound_tasks = list(self.inbound_event_tasks)
for task in inbound_tasks:
if not task.done():
task.cancel()
if inbound_tasks:
await asyncio.gather(*inbound_tasks, return_exceptions=True)
self.inbound_event_tasks.clear()
with self.threadsafe_event_lock:
threadsafe_futures = list(self.threadsafe_event_futures)
for future in threadsafe_futures:
future.cancel()
if threadsafe_futures:
await asyncio.gather(
*(asyncio.wrap_future(future) for future in threadsafe_futures),
return_exceptions=True,
)
with self.threadsafe_event_lock:
self.threadsafe_event_futures.clear()
self.tenant_access_tokens.clear()
self.card_id_dict.clear()
self.pending_monitoring_msg.clear()
self.reply_to_monitoring_msg.clear()
for card_id in tuple(self.card_last_accessed):
self._drop_card_state(card_id)
self.card_last_accessed.clear()
self.reply_message_card_ids.clear()
self.card_sequence_dict.clear()
self.card_id_to_source_ids.clear()
self.card_streaming_text.clear()
self.card_pre_pause_text.clear()
self.card_form_content.clear()
self.card_form_input_defs.clear()
self.card_form_inputs.clear()
self.card_resume_transitioned.clear()
return False
src/langbot/pkg/platform/sources/legacy/gewechat.py
+13 -8
View File
@@ -6,7 +6,6 @@ import traceback
import time
import re
import copy
import threading
import quart
from langbot.pkg.utils import httpclient
@@ -483,6 +482,7 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.message_converter = GewechatMessageConverter(config)
self.event_converter = GewechatEventConverter(config)
self.listeners = {}
@self.quart_app.route('/gewechat/callback', methods=['POST'])
async def gewechat_callback():
@@ -518,9 +518,13 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
at_targets = at_targets or []
member_info = []
if at_targets:
member_info = self.bot.get_chatroom_member_detail(self.config['app_id'], target_id, at_targets[::-1])[
'data'
]
member_result = await asyncio.to_thread(
self.bot.get_chatroom_member_detail,
self.config['app_id'],
target_id,
at_targets[::-1],
)
member_info = member_result['data']
# 处理消息组件
for msg in content_list:
@@ -596,7 +600,7 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
}
if handler := handler_map.get(msg['type']):
handler(msg)
await asyncio.to_thread(handler, msg)
else:
await self.logger.warning(f'未处理的消息类型: {msg["type"]}')
continue
@@ -645,8 +649,9 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
json={'app_id': self.config['app_id']},
) as response:
if response.status != 200:
raise Exception(f'获取gewechat token失败: {await response.text()}')
self.config['token'] = (await response.json())['data']
error = await httpclient.read_text_limited(response)
raise Exception(f'获取gewechat token失败: {error}')
self.config['token'] = (await httpclient.read_json_limited(response))['data']
self.bot = gewechat_client.GewechatClient(f'{self.config["gewechat_url"]}/v2/api', self.config['token'])
@@ -672,7 +677,7 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
except Exception as e:
raise Exception(f'设置 Gewechat 回调失败， token失效： {e}')
threading.Thread(target=gewechat_login_process).start()
await asyncio.to_thread(gewechat_login_process)
async def shutdown_trigger_placeholder():
while True:
src/langbot/pkg/platform/sources/legacy/nakuru.py
+5 -3
View File
@@ -311,15 +311,17 @@ class NakuruAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
try:
import requests
resp = requests.get(
url='http://{}:{}/get_login_info'.format(self.cfg['host'], self.cfg['http_port']),
resp = await asyncio.to_thread(
requests.get,
'http://{}:{}/get_login_info'.format(self.cfg['host'], self.cfg['http_port']),
headers={'Authorization': 'Bearer ' + self.cfg['token'] if 'token' in self.cfg else ''},
timeout=5,
proxies=None,
)
if resp.status_code == 403:
raise Exception('go-cqhttp拒绝访问，请检查配置文件中nakuru适配器的配置')
self.bot_account_id = int(resp.json()['data']['user_id'])
response_data = await httpclient.parse_json_response(resp)
self.bot_account_id = int(response_data['data']['user_id'])
except Exception:
raise Exception('获取go-cqhttp账号信息失败, 请检查是否已启动go-cqhttp并配置正确')
await self.bot._run()
src/langbot/pkg/platform/sources/legacy/qqbotpy.py
+10 -1
View File
@@ -5,6 +5,7 @@ import typing
import datetime
import re
import traceback
from collections import OrderedDict
import botpy
import botpy.message as botpy_message
@@ -40,7 +41,8 @@ event_handler_mapping = {
}
cached_message_ids = {}
_CACHED_MESSAGE_ID_LIMIT = 10000
cached_message_ids: OrderedDict[str, str] = OrderedDict()
"""由于QQ官方的消息id是字符串，而YiriMirai的消息id是整数，所以需要一个索引来进行转换"""
id_index = 0
@@ -53,6 +55,8 @@ def save_msg_id(message_id: str) -> int:
crt_index = id_index
id_index += 1
cached_message_ids[str(crt_index)] = message_id
while len(cached_message_ids) > _CACHED_MESSAGE_ID_LIMIT:
cached_message_ids.popitem(last=False)
return crt_index
@@ -355,6 +359,7 @@ class OfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.cfg = cfg
self.ap = ap
self.logger = logger
self.cached_official_messages = OrderedDict()
self.group_msg_seq = 1
self.c2c_msg_seq = 1
@@ -490,6 +495,8 @@ class OfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
],
):
self.cached_official_messages[str(message.id)] = message
while len(self.cached_official_messages) > 1000:
self.cached_official_messages.popitem(last=False)
await callback(self.event_converter.target2yiri(message), self)
for event_handler in event_handler_mapping[event_type]:
@@ -519,6 +526,8 @@ class OfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await (await self.bot.start(**self.cfg))
async def kill(self) -> bool:
self.cached_official_messages.clear()
if not self.bot.is_closed():
await self.bot.close()
return True
return True
src/langbot/pkg/platform/sources/line.py
+23 -7
View File
@@ -13,6 +13,7 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
from ..logger import EventLogger
from ...utils import bounded_executor
from linebot.v3 import WebhookHandler
@@ -30,6 +31,14 @@ from linebot.v3.webhooks import (
from linebot.v3.webhook import WebhookParser
from linebot.v3.messaging import MessagingApiBlob
MAX_LINE_MEDIA_BYTES = 10 * 1024 * 1024
def _validate_line_media_content(content: bytes) -> bytes:
if len(content) > MAX_LINE_MEDIA_BYTES:
raise ValueError(f'LINE media exceeds the {MAX_LINE_MEDIA_BYTES}-byte limit')
return content
class LINEMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
@staticmethod
@@ -63,9 +72,13 @@ class LINEMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
elif isinstance(message.message, VideoMessageContent):
pass
elif isinstance(message.message, ImageMessageContent):
message_content = MessagingApiBlob(bot_client).get_message_content(message.message.id)
message_content = await asyncio.to_thread(
MessagingApiBlob(bot_client).get_message_content,
message.message.id,
)
_validate_line_media_content(message_content)
base64_string = base64.b64encode(message_content).decode('utf-8')
base64_string = await asyncio.to_thread(lambda: base64.b64encode(message_content).decode('utf-8'))
# 如果需要Data URI格式（用于直接嵌入HTML等）
# 首先需要知道图片类型，LINE图片通常是JPEG
@@ -173,20 +186,22 @@ class LINEAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
for content in content_list:
if content['type'] == 'text':
self.bot.reply_message_with_http_info(
await asyncio.to_thread(
self.bot.reply_message_with_http_info,
ReplyMessageRequest(
reply_token=message_source.source_platform_object.reply_token,
messages=[TextMessage(text=content['content'])],
)
),
)
elif content['type'] == 'image':
# LINE ImageMessage requires original_content_url and preview_image_url
image_url = content['image']
self.bot.reply_message_with_http_info(
await asyncio.to_thread(
self.bot.reply_message_with_http_info,
ReplyMessageRequest(
reply_token=message_source.source_platform_object.reply_token,
messages=[ImageMessage(original_content_url=image_url, preview_image_url=image_url)],
)
),
)
async def is_muted(self, group_id: int) -> bool:
@@ -266,4 +281,5 @@ class LINEAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await keep_alive()
async def kill(self) -> bool:
pass
await bounded_executor.run_blocking_cleanup(self.api_client.close)
return True
src/langbot/pkg/platform/sources/matrix.py
+98 -27
View File
@@ -5,6 +5,8 @@ import asyncio
import traceback
import base64
import json
import os
from urllib.parse import urlparse
import nio
@@ -16,6 +18,58 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
_MAX_MATRIX_MEDIA_BYTES = 10 * 1024 * 1024
def _decode_matrix_base64_limited(value: str) -> bytes:
if ';base64,' in value:
value = value.split(';base64,', 1)[1]
max_encoded_bytes = 4 * ((_MAX_MATRIX_MEDIA_BYTES + 2) // 3)
if len(value) > max_encoded_bytes:
raise ValueError('Matrix media exceeds the size limit')
decoded = base64.b64decode(value)
if len(decoded) > _MAX_MATRIX_MEDIA_BYTES:
raise ValueError('Matrix media exceeds the size limit')
return decoded
def _read_matrix_file_limited(path: str) -> bytes:
if os.path.getsize(path) > _MAX_MATRIX_MEDIA_BYTES:
raise ValueError('Matrix media exceeds the size limit')
with open(path, 'rb') as file:
body = file.read(_MAX_MATRIX_MEDIA_BYTES + 1)
if len(body) > _MAX_MATRIX_MEDIA_BYTES:
raise ValueError('Matrix media exceeds the size limit')
return body
async def _download_matrix_media_limited(
client: nio.AsyncClient,
mxc_url: str,
) -> tuple[bytes, str]:
parsed = urlparse(mxc_url)
if parsed.scheme != 'mxc' or not parsed.netloc or not parsed.path.strip('/'):
raise ValueError('Invalid Matrix media URL')
method, path = nio.Api.download(
parsed.netloc,
parsed.path.replace('/', ''),
access_token=None,
)
headers = {}
if client.access_token:
headers['Authorization'] = f'Bearer {client.access_token}'
response = await client.send(method, path, headers=headers, timeout=30)
try:
response.raise_for_status()
body = await httpclient.read_limited(
response,
max_bytes=_MAX_MATRIX_MEDIA_BYTES,
)
return body, response.headers.get('Content-Type', 'application/octet-stream')
finally:
response.release()
class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
@staticmethod
async def yiri2target(message_chain: platform_message.MessageChain, client: nio.AsyncClient) -> list[dict]:
@@ -26,17 +80,22 @@ class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
elif isinstance(component, platform_message.Image):
image_bytes = None
if component.base64:
b64_data = component.base64
if ';base64,' in b64_data:
b64_data = b64_data.split(';base64,', 1)[1]
image_bytes = base64.b64decode(b64_data)
image_bytes = await asyncio.to_thread(
_decode_matrix_base64_limited,
component.base64,
)
elif component.url:
session = httpclient.get_session()
async with session.get(component.url) as response:
image_bytes = await response.read()
image_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_MATRIX_MEDIA_BYTES,
)
elif component.path:
with open(component.path, 'rb') as f:
image_bytes = f.read()
image_bytes = await asyncio.to_thread(
_read_matrix_file_limited,
str(component.path),
)
if image_bytes:
resp = await client.upload(image_bytes, content_type='image/png')
if isinstance(resp, nio.UploadResponse):
@@ -44,17 +103,22 @@ class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
elif isinstance(component, platform_message.File):
file_bytes = None
if component.base64:
b64_data = component.base64
if ';base64,' in b64_data:
b64_data = b64_data.split(';base64,', 1)[1]
file_bytes = base64.b64decode(b64_data)
file_bytes = await asyncio.to_thread(
_decode_matrix_base64_limited,
component.base64,
)
elif component.url:
session = httpclient.get_session()
async with session.get(component.url) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_MATRIX_MEDIA_BYTES,
)
elif component.path:
with open(component.path, 'rb') as f:
file_bytes = f.read()
file_bytes = await asyncio.to_thread(
_read_matrix_file_limited,
str(component.path),
)
if file_bytes:
file_name = getattr(component, 'name', None) or 'file'
resp = await client.upload(file_bytes, content_type='application/octet-stream', filename=file_name)
@@ -86,11 +150,12 @@ class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
elif isinstance(event, nio.RoomMessageImage):
mxc_url = event.url
if mxc_url:
resp = await client.download(mxc_url)
if isinstance(resp, nio.DownloadResponse):
b64 = base64.b64encode(resp.body).decode('utf-8')
content_type = resp.content_type or 'image/png'
message_components.append(platform_message.Image(base64=f'data:{content_type};base64,{b64}'))
body, content_type = await _download_matrix_media_limited(
client,
mxc_url,
)
b64 = (await asyncio.to_thread(base64.b64encode, body)).decode('utf-8')
message_components.append(platform_message.Image(base64=f'data:{content_type};base64,{b64}'))
if event.body:
message_components.append(platform_message.Plain(text=event.body))
@@ -431,14 +496,15 @@ class MatrixAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
if not mxc_url:
return
try:
resp = await self.client.download(mxc_url)
if isinstance(resp, nio.DownloadResponse):
b64 = base64.b64encode(resp.body).decode('utf-8')
content_type = resp.content_type or 'image/png'
await self.logger.info(
f'[{_b.user_id}] Bridge 发送了二维码，请扫码登录:',
images=[platform_message.Image(base64=f'data:{content_type};base64,{b64}')],
)
body, content_type = await _download_matrix_media_limited(
self.client,
mxc_url,
)
b64 = (await asyncio.to_thread(base64.b64encode, body)).decode('utf-8')
await self.logger.info(
f'[{_b.user_id}] Bridge 发送了二维码，请扫码登录:',
images=[platform_message.Image(base64=f'data:{content_type};base64,{b64}')],
)
except Exception:
await self.logger.error(
f'[{_b.user_id}] Failed to download bridge QR image: {traceback.format_exc()}'
@@ -672,11 +738,16 @@ class MatrixAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async def kill(self) -> bool:
self._running = False
bridge_tasks = []
for bridge in self._bridges:
if bridge.login_task and not bridge.login_task.done():
bridge.login_task.cancel()
bridge_tasks.append(bridge.login_task)
if bridge.check_task and not bridge.check_task.done():
bridge.check_task.cancel()
bridge_tasks.append(bridge.check_task)
if bridge_tasks:
await asyncio.gather(*bridge_tasks, return_exceptions=True)
if self.client:
await self.client.close()
await self.logger.debug('Matrix adapter stopped')
src/langbot/pkg/platform/sources/officialaccount.py
+1
View File
@@ -164,6 +164,7 @@ class OfficialAccountAdapter(abstract_platform_adapter.AbstractMessagePlatformAd
await keep_alive()
async def kill(self) -> bool:
self.bot.clear()
return False
async def unregister_listener(
src/langbot/pkg/platform/sources/openclaw_weixin.py
+35 -11
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import asyncio
import base64
import os
import traceback
import typing
@@ -26,6 +27,7 @@ from langbot.libs.openclaw_weixin_api.types import (
WeixinMessage,
)
from langbot.pkg.entity.persistence import bot as persistence_bot
from langbot.pkg.utils import httpclient
import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
@@ -35,6 +37,8 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
from langbot.pkg.api.http.context import ExecutionContext
_MAX_OPENCLAW_COMPONENT_BYTES = 10 * 1024 * 1024
class OpenClawWeixinMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
"""Converts between LangBot MessageChain and OpenClaw WeChat message items."""
@@ -114,7 +118,12 @@ class OpenClawWeixinMessageConverter(abstract_platform_adapter.AbstractMessageCo
elif item.type == MessageItem.IMAGE and item.image_item:
if hasattr(item.image_item, '_downloaded_bytes') and item.image_item._downloaded_bytes:
b64 = base64.b64encode(item.image_item._downloaded_bytes).decode('utf-8')
b64 = (
await asyncio.to_thread(
base64.b64encode,
item.image_item._downloaded_bytes,
)
).decode('utf-8')
components.append(platform_message.Image(base64=f'data:image/jpeg;base64,{b64}'))
else:
components.append(platform_message.Unknown(text='[Image]'))
@@ -401,19 +410,30 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
path_val = getattr(component, 'path', None)
if b64_val:
return base64.b64decode(b64_val)
max_encoded_chars = 4 * ((_MAX_OPENCLAW_COMPONENT_BYTES + 2) // 3) + 4
if len(b64_val) > max_encoded_chars:
raise ValueError('OpenClaw media exceeds the size limit')
data = await asyncio.to_thread(base64.b64decode, b64_val)
if len(data) > _MAX_OPENCLAW_COMPONENT_BYTES:
raise ValueError('OpenClaw media exceeds the size limit')
return data
elif url_val and url_val.startswith(('http://', 'https://')):
import aiohttp
async with aiohttp.ClientSession() as session:
async with session.get(url_val) as resp:
if resp.status == 200:
return await resp.read()
session = httpclient.get_session()
async with session.get(url_val) as resp:
if resp.status == 200:
return await httpclient.read_limited(resp)
elif path_val:
import asyncio
if await asyncio.to_thread(os.path.getsize, path_val) > _MAX_OPENCLAW_COMPONENT_BYTES:
raise ValueError('OpenClaw media exceeds the size limit')
with open(path_val, 'rb') as f:
return await asyncio.to_thread(f.read)
def read_file() -> bytes:
with open(path_val, 'rb') as file:
return file.read(_MAX_OPENCLAW_COMPONENT_BYTES + 1)
data = await asyncio.to_thread(read_file)
if len(data) > _MAX_OPENCLAW_COMPONENT_BYTES:
raise ValueError('OpenClaw media exceeds the size limit')
return data
return None
def register_listener(
@@ -544,6 +564,8 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
"""Process a single inbound message from getUpdates."""
if msg.context_token and msg.from_user_id:
self._context_tokens[msg.from_user_id] = msg.context_token
while len(self._context_tokens) > 4096:
self._context_tokens.pop(next(iter(self._context_tokens)), None)
# Download CDN media (files, images) before converting to LangBot events
await self._download_media_items(msg)
@@ -599,6 +621,8 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
await self._poll_task
except asyncio.CancelledError:
pass
self._poll_task = None
self._context_tokens.clear()
await self.client.close()
await self.logger.info('OpenClaw WeChat adapter stopped')
return True
src/langbot/pkg/platform/sources/qqofficial.py
+61 -4
View File
@@ -241,6 +241,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
# per (msg_id|event_id) within 60 min, but each reuse needs a
# fresh ``msg_seq`` — re-sending with msg_seq=1 is silently dedup'd.
self._anchor_msg_seq: dict[str, int] = {}
self._background_tasks: set[asyncio.Task] = set()
# Wire button-click handler so webhook mode catches INTERACTION_CREATE.
# (ws mode is wired separately via on_event in _run_websocket so the
@@ -249,6 +250,30 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
async def _on_interaction(event_data: dict, interaction_id: typing.Optional[str]):
await self._handle_interaction_create(event_data, interaction_id)
def _start_background_task(self, coro) -> bool:
"""Start one bounded adapter-side auxiliary task."""
background_tasks = getattr(self, '_background_tasks', None)
if background_tasks is None:
background_tasks = set()
object.__setattr__(self, '_background_tasks', background_tasks)
for task in tuple(background_tasks):
if task.done():
background_tasks.discard(task)
if len(background_tasks) >= 100:
coro.close()
return False
task = asyncio.create_task(coro)
background_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
background_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
return True
async def reply_message(
self,
message_source: platform_events.MessageEvent,
@@ -449,6 +474,14 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
pass
async def kill(self) -> bool:
task_set = getattr(self, '_background_tasks', set())
background_tasks = list(task_set)
for task in background_tasks:
if not task.done():
task.cancel()
if background_tasks:
await asyncio.gather(*background_tasks, return_exceptions=True)
task_set.clear()
if self._ws_task:
self._ws_task.cancel()
try:
@@ -456,6 +489,14 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
except asyncio.CancelledError:
pass
self._ws_task = None
await self.bot.close()
self._pending_forms.clear()
self._session_event_ids.clear()
self._anchor_msg_seq.clear()
self._stream_ctx.clear()
self._stream_ctx_ts.clear()
self._fallback_text.clear()
self._fallback_text_ts.clear()
return True
# --------------- 流式输出 ---------------
@@ -473,6 +514,14 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
for mid in stale_fb:
self._fallback_text.pop(mid, None)
self._fallback_text_ts.pop(mid, None)
while len(self._stream_ctx) > 1000:
oldest = min(self._stream_ctx_ts, key=self._stream_ctx_ts.__getitem__)
self._stream_ctx.pop(oldest, None)
self._stream_ctx_ts.pop(oldest, None)
while len(self._fallback_text) > 1000:
oldest = min(self._fallback_text_ts, key=self._fallback_text_ts.__getitem__)
self._fallback_text.pop(oldest, None)
self._fallback_text_ts.pop(oldest, None)
if stale_ids or stale_fb:
await self.logger.debug(f'Cleaned up {len(stale_ids)} stream contexts, {len(stale_fb)} fallback texts')
@@ -508,6 +557,8 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
# msg_seq=2 instead of being deduplicated by QQ as another seq=1 send.
if source.d_id:
self._anchor_msg_seq[source.d_id] = max(self._anchor_msg_seq.get(source.d_id, 0), 1)
while len(self._anchor_msg_seq) > 4096:
self._anchor_msg_seq.pop(next(iter(self._anchor_msg_seq)), None)
ctx = {
'user_openid': source.user_openid,
@@ -577,7 +628,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
# 非流式场景（如群聊不支持流式），累积文本后一次性回复
if chunk_text:
# Chunks carry the latest full snapshot, not a text delta.
self._fallback_text[message_id] = chunk_text
self._fallback_text[message_id] = chunk_text[:200000]
self._fallback_text_ts[message_id] = time.time()
if is_final:
full_text = self._fallback_text.pop(message_id, '')
@@ -590,7 +641,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
# 累积文本
if chunk_text:
ctx['accumulated_text'] = chunk_text
ctx['accumulated_text'] = chunk_text[:200000]
# 未启动会话时，等第一个有内容的 chunk 来建立会话
if not ctx['session_started']:
@@ -668,6 +719,8 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
if used >= self._MAX_REPLIES_PER_ANCHOR:
return None
self._anchor_msg_seq[anchor] = used + 1
while len(self._anchor_msg_seq) > 4096:
self._anchor_msg_seq.pop(next(iter(self._anchor_msg_seq)), None)
return used + 1
async def _reply_synthetic(
@@ -791,7 +844,9 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
k for k, v in self._session_event_ids.items() if now - v.get('posted_at', 0) > self._PENDING_FORM_TTL
]
for k in stale_e:
self._session_event_ids.pop(k, None)
stale_event = self._session_event_ids.pop(k, None)
if stale_event:
self._anchor_msg_seq.pop(stale_event.get('event_id'), None)
async def _handle_form_chunk(
self,
@@ -973,7 +1028,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
# ACK uses the interaction id, NOT the ws event id.
interaction_id = event_data.get('id') or ''
if interaction_id:
asyncio.create_task(self.bot.ack_interaction(interaction_id, code=0))
self._start_background_task(self.bot.ack_interaction(interaction_id, code=0))
resolved = (event_data.get('data') or {}).get('resolved') or {}
action_id = str(resolved.get('button_data') or resolved.get('button_id') or '').strip()
@@ -1018,6 +1073,8 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
}
# New anchor → fresh 5-reply budget.
self._anchor_msg_seq[cached_event_id] = 0
while len(self._anchor_msg_seq) > 4096:
self._anchor_msg_seq.pop(next(iter(self._anchor_msg_seq)), None)
if self.ap is not None and not ws_event_id:
self.ap.logger.warning(
'QQ Official: INTERACTION_CREATE lacked ws_event_id; '
src/langbot/pkg/platform/sources/satori.py
+21 -7
View File
@@ -18,6 +18,9 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
from langbot.pkg.utils import httpclient
_MAX_GATEWAY_MESSAGE_BYTES = 1024 * 1024
class SatoriMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
@@ -63,7 +66,15 @@ class SatoriMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
padding = 4 - len(raw_b64) % 4
if padding != 4:
raw_b64 += '=' * padding
image_bytes = base64.b64decode(raw_b64)
max_encoded_chars = 4 * ((10 * 1024 * 1024 + 2) // 3) + 4
if len(raw_b64) > max_encoded_chars:
raise ValueError('Satori image exceeds the 10 MiB limit')
image_bytes = await asyncio.to_thread(
base64.b64decode,
raw_b64,
)
if len(image_bytes) > 10 * 1024 * 1024:
raise ValueError('Satori image exceeds the 10 MiB limit')
uploaded_url = await adapter.upload_image(image_bytes, mime_type)
if uploaded_url:
await adapter.logger.info(f'Satori 图片上传成功: {len(image_bytes)} 字节')
@@ -492,7 +503,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
raise ValueError(f'WebSocket URL必须以ws://或wss://开头: {self.endpoint}')
try:
self.ws = await websockets.connect(self.endpoint)
self.ws = await websockets.connect(self.endpoint, max_size=_MAX_GATEWAY_MESSAGE_BYTES)
await asyncio.sleep(0.1)
await self.send_identify()
@@ -584,7 +595,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async def handle_message(self, message: str):
"""Handle WebSocket message"""
try:
data = json.loads(message)
data = await asyncio.to_thread(json.loads, message)
op = data.get('op')
body = data.get('body', {})
@@ -831,9 +842,10 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
try:
async with self.session.request(method, url, headers=headers, json=data) as response:
if response.status == 200:
return await response.json()
result = await httpclient.read_json_limited(response)
return result if isinstance(result, dict) else None
else:
text = await response.text()
text = await httpclient.read_text_limited(response)
await self.logger.error(f'Satori API 请求失败: {response.status} - {text}')
return None
except Exception as e:
@@ -889,7 +901,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async with self.session.post(url, headers=headers, data=form_data) as response:
if response.status == 200:
result = await response.json()
result = await httpclient.read_json_limited(response)
# The response should contain the URL of the uploaded file
if isinstance(result, dict) and 'url' in result:
return result['url']
@@ -899,7 +911,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await self.logger.warning(f'Satori 图片上传响应格式未知: {result}')
return None
else:
text = await response.text()
text = await httpclient.read_text_limited(response)
await self.logger.error(f'Satori 图片上传失败: {response.status} - {text}')
return None
except Exception as e:
@@ -911,6 +923,8 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.running = False
if self.heartbeat_task:
self.heartbeat_task.cancel()
await asyncio.gather(self.heartbeat_task, return_exceptions=True)
self.heartbeat_task = None
if self.ws:
try:
await self.ws.close()
src/langbot/pkg/platform/sources/telegram.py
+72 -27
View File
@@ -10,6 +10,8 @@ import typing
import traceback
import json
import base64
import asyncio
import os
import time
import uuid
import pydantic
@@ -22,6 +24,31 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
_MAX_TELEGRAM_MEDIA_BYTES = 10 * 1024 * 1024
def _decode_telegram_base64_limited(value: str) -> bytes:
if ';base64,' in value:
value = value.split(';base64,', 1)[1]
max_encoded_bytes = 4 * ((_MAX_TELEGRAM_MEDIA_BYTES + 2) // 3)
if len(value) > max_encoded_bytes:
raise ValueError('Telegram media exceeds the size limit')
decoded = base64.b64decode(value)
if len(decoded) > _MAX_TELEGRAM_MEDIA_BYTES:
raise ValueError('Telegram media exceeds the size limit')
return decoded
def _read_telegram_file_limited(path: str) -> bytes:
if os.path.getsize(path) > _MAX_TELEGRAM_MEDIA_BYTES:
raise ValueError('Telegram media exceeds the size limit')
with open(path, 'rb') as file:
body = file.read(_MAX_TELEGRAM_MEDIA_BYTES + 1)
if len(body) > _MAX_TELEGRAM_MEDIA_BYTES:
raise ValueError('Telegram media exceeds the size limit')
return body
def _telegram_select_field_options(form_data: dict) -> tuple[str, list[str]]:
"""Return the active select field and its option values."""
field_name = str(form_data.get('_current_input_field') or '').strip()
@@ -86,35 +113,29 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
if isinstance(component, platform_message.Plain):
components.append({'type': 'text', 'text': component.text})
elif isinstance(component, platform_message.Image):
photo_bytes = None
if component.base64:
photo_bytes = base64.b64decode(component.base64)
elif component.url:
session = httpclient.get_session()
async with session.get(component.url) as response:
photo_bytes = await response.read()
elif component.path:
with open(component.path, 'rb') as f:
photo_bytes = f.read()
photo_bytes, _mime_type = await component.get_bytes()
components.append({'type': 'photo', 'photo': photo_bytes})
elif isinstance(component, platform_message.File):
file_bytes = None
if component.base64:
# Strip data URI prefix if present (e.g. "data:application/pdf;base64,...")
b64_data = component.base64
if ';base64,' in b64_data:
b64_data = b64_data.split(';base64,', 1)[1]
file_bytes = base64.b64decode(b64_data)
file_bytes = await asyncio.to_thread(
_decode_telegram_base64_limited,
component.base64,
)
elif component.url:
session = httpclient.get_session()
async with session.get(component.url) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
)
elif component.path:
with open(component.path, 'rb') as f:
file_bytes = f.read()
file_bytes = await asyncio.to_thread(
_read_telegram_file_limited,
str(component.path),
)
file_name = getattr(component, 'name', None) or 'file'
components.append({'type': 'document', 'document': file_bytes, 'filename': file_name})
@@ -152,13 +173,15 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
file_format = ''
async with httpclient.get_session(trust_env=True).get(file.file_path) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
)
file_format = 'image/jpeg'
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
message_components.append(
platform_message.Image(
base64=f'data:{file_format};base64,{base64.b64encode(file_bytes).decode("utf-8")}'
)
platform_message.Image(base64=f'data:{file_format};base64,{encoded.decode("utf-8")}')
)
if message.voice:
@@ -171,11 +194,15 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
file_format = message.voice.mime_type or 'audio/ogg'
async with httpclient.get_session(trust_env=True).get(file.file_path) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
)
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
message_components.append(
platform_message.Voice(
base64=f'data:{file_format};base64,{base64.b64encode(file_bytes).decode("utf-8")}',
base64=f'data:{file_format};base64,{encoded.decode("utf-8")}',
length=message.voice.duration,
)
)
@@ -188,16 +215,22 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
file_name = message.document.file_name or 'document'
file_size = message.document.file_size or 0
file_format = message.document.mime_type or 'application/octet-stream'
if file_size > _MAX_TELEGRAM_MEDIA_BYTES:
raise ValueError('Telegram media exceeds the size limit')
file_bytes = None
async with httpclient.get_session(trust_env=True).get(file.file_path) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
)
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
message_components.append(
platform_message.File(
name=file_name,
size=file_size,
base64=f'data:{file_format};base64,{base64.b64encode(file_bytes).decode("utf-8")}',
base64=f'data:{file_format};base64,{encoded.decode("utf-8")}',
)
)
@@ -263,6 +296,8 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
] = {}
_FORM_ACTION_CACHE_TTL = 30 * 60
_MAX_FORM_ACTION_TITLES = 4096
_MAX_STREAM_STATES = 1000
# callback_data -> (display title, pipeline UUID, expiration time, form group id)
_form_action_titles: typing.Dict[str, tuple[str, str, float, str]] = {}
@@ -285,6 +320,8 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self._form_action_titles.update(
{callback_data: (title, pipeline_uuid, expires_at, group_id) for callback_data, title in mappings.items()}
)
while len(self._form_action_titles) > self._MAX_FORM_ACTION_TITLES:
self._form_action_titles.pop(next(iter(self._form_action_titles)), None)
def _take_form_action_context(self, callback_data: str, now: float | None = None) -> tuple[str, str] | None:
"""Consume a callback and invalidate every button from the same form."""
@@ -445,6 +482,11 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
bot_account_id='',
listeners={},
)
self._form_action_titles = {}
def _cap_stream_states(self) -> None:
while len(self.msg_stream_id) > self._MAX_STREAM_STATES:
self.msg_stream_id.pop(next(iter(self.msg_stream_id)), None)
async def send_message(self, target_type: str, target_id: str, message: platform_message.MessageChain):
components = await TelegramMessageConverter.yiri2target(message, self.bot)
@@ -553,6 +595,7 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
args = self._build_message_args(chat_id, 'Thinking...', message_thread_id)
send_msg = await self.bot.send_message(**args)
self.msg_stream_id[message_id] = ('message', send_msg.message_id, False)
self._cap_stream_states()
return True
@@ -844,4 +887,6 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
if self.application.updater:
await self.application.updater.stop()
await self.logger.info('Telegram adapter stopped')
self.msg_stream_id.clear()
self._form_action_titles.clear()
return True
src/langbot/pkg/platform/sources/websocket_adapter.py
+172 -21
View File
@@ -3,6 +3,7 @@
import asyncio
import contextvars
import logging
import time
import typing
from datetime import datetime
@@ -14,6 +15,7 @@ import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
from ...core import app
from ...core import entities as core_entities
from .websocket_manager import WebSocketConnection, WebSocketScope, is_valid_session_id, ws_connection_manager
logger = logging.getLogger(__name__)
@@ -45,21 +47,82 @@ class WebSocketSession:
stream_message_indexes: dict[str, dict[str, int]] = {}
"""流式消息索引 {pipeline_uuid: {resp_message_id: message_index}}"""
def __init__(self, id: str):
def __init__(
self,
id: str = '',
*,
max_conversations: int = 200,
max_messages: int = 100,
idle_ttl_seconds: int = 86400,
):
self.id = id
self.message_lists = {}
self.stream_message_indexes = {}
self.message_counters: dict[str, int] = {}
self.last_accessed: dict[str, float] = {}
self.max_conversations = max(int(max_conversations), 1)
self.max_messages = max(int(max_messages), 1)
self.idle_ttl_seconds = max(int(idle_ttl_seconds), 1)
def _prune(self, now: float) -> None:
expired = [
key for key, last_accessed in self.last_accessed.items() if now - last_accessed >= self.idle_ttl_seconds
]
for key in expired:
self.reset(key)
overflow = len(self.message_lists) - self.max_conversations + 1
if overflow <= 0:
return
oldest = sorted(self.last_accessed, key=self.last_accessed.get)
for key in oldest[:overflow]:
self.reset(key)
def get_message_list(self, pipeline_uuid: str) -> list[WebSocketMessage]:
now = time.monotonic()
self._prune(now)
if pipeline_uuid not in self.message_lists:
self.message_lists[pipeline_uuid] = []
self.last_accessed[pipeline_uuid] = now
return self.message_lists[pipeline_uuid]
def get_stream_message_indexes(self, pipeline_uuid: str) -> dict[str, int]:
if pipeline_uuid not in self.stream_message_indexes:
self.stream_message_indexes[pipeline_uuid] = {}
self.last_accessed[pipeline_uuid] = time.monotonic()
return self.stream_message_indexes[pipeline_uuid]
def next_message_id(self, conversation_key: str) -> int:
next_id = self.message_counters.get(conversation_key, 0) + 1
self.message_counters[conversation_key] = next_id
return next_id
def append_message(self, conversation_key: str, message: WebSocketMessage) -> None:
messages = self.get_message_list(conversation_key)
messages.append(message)
overflow = len(messages) - self.max_messages
if overflow <= 0:
return
del messages[:overflow]
indexes = self.stream_message_indexes.get(conversation_key, {})
adjusted_indexes = {
response_id: index - overflow for response_id, index in indexes.items() if index >= overflow
}
indexes.clear()
indexes.update(adjusted_indexes)
def reset(self, conversation_key: str) -> None:
self.message_lists.pop(conversation_key, None)
self.stream_message_indexes.pop(conversation_key, None)
self.message_counters.pop(conversation_key, None)
self.last_accessed.pop(conversation_key, None)
def clear(self) -> None:
self.message_lists.clear()
self.stream_message_indexes.clear()
self.message_counters.clear()
self.last_accessed.clear()
class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
"""WebSocket适配器 - 支持双向实时通信"""
@@ -75,7 +138,14 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
ap: app.Application = pydantic.Field(exclude=True)
# 主动推送消息的队列
outbound_message_queue: asyncio.Queue = pydantic.Field(default_factory=asyncio.Queue, exclude=True)
outbound_message_queue: asyncio.Queue = pydantic.Field(
default_factory=lambda: asyncio.Queue(maxsize=100),
exclude=True,
)
inbound_listener_tasks: set[asyncio.Task] = pydantic.Field(
default_factory=set,
exclude=True,
)
"""后端主动推送消息的队列"""
# 流式输出开关
@@ -89,11 +159,26 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
**kwargs,
)
self.websocket_person_session = WebSocketSession(id='websocketperson')
self.websocket_group_session = WebSocketSession(id='websocketgroup')
application = kwargs.get('ap')
instance_data = getattr(getattr(application, 'instance_config', None), 'data', {})
retention = (
instance_data.get('system', {}).get('websocket_retention', {}) if isinstance(instance_data, dict) else {}
)
session_options = {
'max_conversations': retention.get('max_conversations_per_workspace', 200),
'max_messages': retention.get('max_messages_per_conversation', 100),
'idle_ttl_seconds': retention.get('conversation_idle_ttl_seconds', 86400),
}
self.websocket_person_session = WebSocketSession(id='websocketperson', **session_options)
self.websocket_group_session = WebSocketSession(id='websocketgroup', **session_options)
self.bot_account_id = 'websocketbot'
self.outbound_message_queue = asyncio.Queue()
try:
outbound_queue_size = max(int(retention.get('send_queue_size', 100)), 1)
except (TypeError, ValueError):
outbound_queue_size = 100
self.outbound_message_queue = asyncio.Queue(maxsize=outbound_queue_size)
self.inbound_listener_tasks = set()
self.stream_enabled = True
@staticmethod
@@ -128,6 +213,25 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
return _current_pipeline_uuid.get()
def _listener_task_done(self, task: asyncio.Task) -> None:
listener_tasks = getattr(self, 'inbound_listener_tasks', None)
if listener_tasks is not None:
listener_tasks.discard(task)
if not task.cancelled():
task.exception()
@staticmethod
def _history_message_chain(message_chain: list[dict]) -> list[dict]:
"""Remove large transient payloads before retaining browser history."""
history = []
for component in message_chain:
copied = dict(component)
if copied.get('base64'):
copied['base64'] = ''
history.append(copied)
return history
async def _get_connection_from_target(self, target_id: str):
"""Resolve a person or group WebSocket launcher to its connection."""
scope = self._scope()
@@ -195,7 +299,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
session = self.websocket_group_session if session_type == 'group' else self.websocket_person_session
msg_id = len(session.get_message_list(conversation_key)) + 1
msg_id = session.next_message_id(conversation_key)
message_data = WebSocketMessage(
id=msg_id,
@@ -206,7 +310,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
is_final=True,
)
session.get_message_list(conversation_key).append(message_data)
session.append_message(conversation_key, message_data)
await ws_connection_manager.broadcast_to_pipeline(
pipeline_uuid,
@@ -241,7 +345,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
session_type = 'group' if isinstance(message_source, platform_events.GroupMessage) else 'person'
conversation_key = self._conversation_key(pipeline_uuid, session_id)
msg_id = len(session.get_message_list(conversation_key)) + 1
msg_id = session.next_message_id(conversation_key)
message_data = WebSocketMessage(
id=msg_id,
@@ -252,7 +356,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
is_final=True,
)
session.get_message_list(conversation_key).append(message_data)
session.append_message(conversation_key, message_data)
await ws_connection_manager.broadcast_to_pipeline(
pipeline_uuid,
@@ -300,7 +404,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
if existing_index is None or existing_index >= len(message_list):
# 创建新消息
msg_id = len(message_list) + 1
msg_id = session.next_message_id(conversation_key)
message_data = WebSocketMessage(
id=msg_id,
role='assistant',
@@ -311,7 +415,8 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
)
# 立即添加到历史记录（即使is_final=False），以便后续块可以更新它
message_list.append(message_data)
session.append_message(conversation_key, message_data)
message_list = session.get_message_list(conversation_key)
if resp_message_id:
stream_message_indexes[resp_message_id] = len(message_list) - 1
else:
@@ -399,7 +504,22 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
async def kill(self):
"""停止适配器"""
pass
await ws_connection_manager.close_scope(self._scope())
listener_tasks = getattr(self, 'inbound_listener_tasks', set())
inbound_tasks = list(listener_tasks)
for task in inbound_tasks:
if not task.done():
task.cancel()
if inbound_tasks:
await asyncio.gather(*inbound_tasks, return_exceptions=True)
listener_tasks.clear()
self.websocket_person_session.clear()
self.websocket_group_session.clear()
while not self.outbound_message_queue.empty():
try:
self.outbound_message_queue.get_nowait()
except asyncio.QueueEmpty:
break
async def _process_image_components(
self,
@@ -445,7 +565,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
try:
file_content = await storage_mgr.storage_provider.load(comp_path)
base64_str = base64.b64encode(file_content).decode('utf-8')
base64_str = (await asyncio.to_thread(base64.b64encode, file_content)).decode('utf-8')
lowered = comp_path.lower()
if comp_type == 'Image':
@@ -507,19 +627,19 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
message_chain = platform_message.MessageChain.model_validate(message_chain_obj)
message_id = len(use_session.get_message_list(conversation_key)) + 1
message_id = use_session.next_message_id(conversation_key)
# 保存用户消息
user_message = WebSocketMessage(
id=message_id,
role='user',
content=str(message_chain),
message_chain=message_chain_obj,
message_chain=self._history_message_chain(message_chain_obj),
timestamp=datetime.now().isoformat(),
connection_id=connection.connection_id,
is_final=True, # 用户消息始终是完整的，非流式
)
use_session.get_message_list(conversation_key).append(user_message)
use_session.append_message(conversation_key, user_message)
await ws_connection_manager.broadcast_to_pipeline(
pipeline_uuid,
@@ -573,9 +693,36 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
owner_bot.adapter.set_ws_adapter(self)
callback_adapter = owner_bot.adapter if (owner_bot and hasattr(owner_bot, 'adapter')) else self
if event.__class__ in listeners:
listener_tasks = getattr(self, 'inbound_listener_tasks', None)
if listener_tasks is None:
listener_tasks = set()
object.__setattr__(self, 'inbound_listener_tasks', listener_tasks)
for task in tuple(listener_tasks):
if task.done():
listener_tasks.discard(task)
if len(listener_tasks) >= 100:
await self.logger.warning('WebSocket inbound listener capacity reached; dropping message')
return
token = _current_pipeline_uuid.set(pipeline_uuid)
try:
asyncio.create_task(listeners[event.__class__](event, callback_adapter))
task_manager = getattr(self.ap, 'task_mgr', None)
if task_manager is None or not isinstance(getattr(task_manager, 'tasks', None), list):
listener_task = asyncio.create_task(listeners[event.__class__](event, callback_adapter))
else:
listener_task = task_manager.create_task(
listeners[event.__class__](event, callback_adapter),
kind='websocket-message',
name=f'websocket-message-{connection.connection_id}',
scopes=[
core_entities.LifecycleControlScope.APPLICATION,
core_entities.LifecycleControlScope.PLATFORM,
],
instance_uuid=connection.instance_uuid,
workspace_uuid=connection.workspace_uuid,
placement_generation=connection.placement_generation,
).task
listener_tasks.add(listener_task)
listener_task.add_done_callback(self._listener_task_done)
finally:
_current_pipeline_uuid.reset(token)
@@ -599,10 +746,14 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
"""Reset one pipeline/client conversation."""
conversation_key = self._conversation_key(pipeline_uuid, session_id)
session = self.websocket_person_session if session_type == 'person' else self.websocket_group_session
if conversation_key in session.message_lists:
session.message_lists[conversation_key] = []
if conversation_key in session.stream_message_indexes:
session.stream_message_indexes[conversation_key] = {}
if isinstance(session, WebSocketSession):
session.reset(conversation_key)
else:
# Compatibility for lightweight adapter doubles.
if conversation_key in session.message_lists:
session.message_lists[conversation_key] = []
if conversation_key in session.stream_message_indexes:
session.stream_message_indexes[conversation_key] = {}
if session_id:
launcher_id = (
src/langbot/pkg/platform/sources/websocket_manager.py
+74 -3
View File
@@ -13,6 +13,7 @@ from ...api.http.context import ExecutionContext
logger = logging.getLogger(__name__)
_SESSION_FILTER_UNSET = object()
_DEFAULT_SEND_QUEUE_SIZE = 100
@dataclasses.dataclass(frozen=True, slots=True)
@@ -82,7 +83,10 @@ class WebSocketConnection(pydantic.BaseModel):
last_active: datetime = pydantic.Field(default_factory=datetime.now)
"""最后活跃时间"""
send_queue: asyncio.Queue = pydantic.Field(default_factory=asyncio.Queue, exclude=True)
send_queue: asyncio.Queue = pydantic.Field(
default_factory=lambda: asyncio.Queue(maxsize=_DEFAULT_SEND_QUEUE_SIZE),
exclude=True,
)
"""发送消息队列"""
is_active: bool = True
@@ -135,9 +139,32 @@ class WebSocketConnectionManager:
session_type: str,
metadata: dict | None = None,
session_id: str | None = None,
send_queue_size: int = _DEFAULT_SEND_QUEUE_SIZE,
max_connections: int = 1024,
max_connections_per_workspace: int = 32,
) -> WebSocketConnection:
"""Register a WebSocket connection and its optional embed session."""
try:
send_queue_size = max(int(send_queue_size), 1)
except (TypeError, ValueError):
send_queue_size = _DEFAULT_SEND_QUEUE_SIZE
max_connections = max(int(max_connections), 1)
max_connections_per_workspace = max(
min(int(max_connections_per_workspace), max_connections),
1,
)
async with self._lock:
if len(self.connections) >= max_connections:
raise RuntimeError(f'WebSocket connection capacity reached ({max_connections})')
workspace_connection_count = sum(
1
for connection in self.connections.values()
if connection.instance_uuid == scope.instance_uuid
and connection.workspace_uuid == scope.workspace_uuid
and connection.placement_generation == scope.placement_generation
)
if workspace_connection_count >= max_connections_per_workspace:
raise RuntimeError(f'Workspace WebSocket connection capacity reached ({max_connections_per_workspace})')
connection = WebSocketConnection(
instance_uuid=scope.instance_uuid,
workspace_uuid=scope.workspace_uuid,
@@ -147,6 +174,7 @@ class WebSocketConnectionManager:
session_id=session_id,
websocket=websocket,
metadata=metadata or {},
send_queue=asyncio.Queue(maxsize=send_queue_size),
)
self.connections[connection.connection_id] = connection
@@ -171,6 +199,31 @@ class WebSocketConnectionManager:
return connection
async def close_scope(self, scope: WebSocketScope) -> None:
"""Close and forget every live connection for one runtime placement."""
async with self._lock:
connection_ids = [
connection_id for connection_id, connection in self.connections.items() if connection.scope == scope
]
for connection_id in connection_ids:
connection = self.connections.get(connection_id)
if connection is None:
continue
close = getattr(connection.websocket, 'close', None)
if close is not None:
try:
result = close()
if asyncio.iscoroutine(result):
await result
except Exception:
logger.debug(
'Failed to close WebSocket connection %s',
connection_id,
exc_info=True,
)
await self.remove_connection(connection_id)
async def remove_connection(self, connection_id: str):
"""移除WebSocket连接"""
async with self._lock:
@@ -237,7 +290,12 @@ class WebSocketConnectionManager:
pipeline_uuid: str | None = None,
) -> WebSocketConnection | None:
"""Get an active embed connection by its stable browser session identifier."""
for connection in self.connections.values():
candidates: typing.Iterable[WebSocketConnection]
if pipeline_uuid is not None:
candidates = await self.get_connections_by_pipeline(pipeline_uuid, scope=scope)
else:
candidates = self.connections.values()
for connection in candidates:
if (
connection.session_id == session_id
and connection.is_active
@@ -307,7 +365,20 @@ class WebSocketConnectionManager:
return
try:
await connection.send_queue.put(message)
try:
connection.send_queue.put_nowait(message)
except asyncio.QueueFull:
# A slow or disconnected browser must not backpressure every
# other connection or retain an unbounded response stream.
try:
connection.send_queue.get_nowait()
except asyncio.QueueEmpty:
pass
connection.send_queue.put_nowait(message)
logger.warning(
'WebSocket send queue full; dropped oldest message for connection %s',
connection_id,
)
connection.last_active = datetime.now()
except Exception as e:
logger.error(f'Failed to send message to connection {connection_id}: {e}')
src/langbot/pkg/platform/sources/wechatpad.py
+142 -65
View File
@@ -1,8 +1,6 @@
import requests
import websocket
import json
import time
import httpx
from langbot.libs.wechatpad_api.client import WeChatPadClient
@@ -17,6 +15,7 @@ import threading
import quart
from langbot.pkg.platform.logger import EventLogger
from langbot.pkg.utils import bounded_executor, httpclient
import xml.etree.ElementTree as ET
from typing import Optional, Tuple
from functools import partial
@@ -27,6 +26,8 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
_MAX_GATEWAY_MESSAGE_CHARS = 1024 * 1024
class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
def __init__(self, config: dict, logger: abstract_platform_logger.AbstractEventLogger):
@@ -53,12 +54,11 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
content_list.append({'type': 'text', 'content': component.text})
elif isinstance(component, platform_message.Image):
if component.url:
async with httpx.AsyncClient() as client:
response = await client.get(component.url)
if response.status_code == 200:
file_bytes = response.content
base64_str = base64.b64encode(file_bytes).decode('utf-8') # 返回字符串格式
session = httpclient.get_session()
async with session.get(component.url) as response:
if response.status == 200:
file_bytes = await httpclient.read_limited(response)
base64_str = (await asyncio.to_thread(base64.b64encode, file_bytes)).decode('utf-8')
else:
raise Exception('获取文件失败')
# pass
@@ -156,9 +156,19 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
cdnthumburl = img_tag.get('cdnthumburl')
# cdnmidimgurl = img_tag.get('cdnmidimgurl')
image_data = self.bot.cdn_download(aeskey=aeskey, file_type=1, file_url=cdnthumburl)
image_data = await asyncio.to_thread(
self.bot.cdn_download,
aeskey=aeskey,
file_type=1,
file_url=cdnthumburl,
)
if image_data['Data']['FileData'] == '':
image_data = self.bot.cdn_download(aeskey=aeskey, file_type=2, file_url=cdnthumburl)
image_data = await asyncio.to_thread(
self.bot.cdn_download,
aeskey=aeskey,
file_type=2,
file_url=cdnthumburl,
)
base64_str = image_data['Data']['FileData']
# self.logger.info(f"data:image/png;base64,{base64_str}")
@@ -186,7 +196,12 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
if voicemsg is not None:
bufid = voicemsg.get('bufid')
length = voicemsg.get('voicelength')
voice_data = self.bot.get_msg_voice(buf_id=str(bufid), length=int(length), msgid=str(new_msg_id))
voice_data = await asyncio.to_thread(
self.bot.get_msg_voice,
buf_id=str(bufid),
length=int(length),
msgid=str(new_msg_id),
)
audio_base64 = voice_data['Data']['Base64']
# 验证语音数据有效性
@@ -319,7 +334,12 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
# print(aeskey,cdnthumburl)
file_data = self.bot.cdn_download(aeskey=aeskey, file_type=5, file_url=cdnthumburl)
file_data = await asyncio.to_thread(
self.bot.cdn_download,
aeskey=aeskey,
file_type=5,
file_url=cdnthumburl,
)
file_base64 = file_data['Data']['FileData']
# print(file_data)
@@ -538,6 +558,7 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
typing.Type[platform_events.Event],
typing.Callable[[platform_events.Event, abstract_platform_adapter.AbstractMessagePlatformAdapter], None],
] = {}
_MAX_CALLBACK_FUTURES = 100
def __init__(self, config: dict, logger: EventLogger):
quart_app = quart.Quart(__name__)
@@ -556,6 +577,12 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
name='WeChatPad',
bot=bot,
)
self._event_loop: asyncio.AbstractEventLoop | None = None
self._ws_app: websocket.WebSocketApp | None = None
self._ws_thread: threading.Thread | None = None
self._stop_event = threading.Event()
self._callback_futures: set = set()
self._callback_futures_lock = threading.Lock()
async def ws_message(self, data):
"""处理接收到的消息"""
@@ -565,7 +592,7 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
except Exception:
await self.logger.error(f'Error in wechatpad callback: {traceback.format_exc()}')
if event.__class__ in self.listeners:
if event is not None and event.__class__ in self.listeners:
await self.listeners[event.__class__](event, self)
return 'ok'
@@ -580,9 +607,8 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
at_targets = at_targets or []
member_info = []
if at_targets:
member_info = self.bot.get_chatroom_member_detail(
target_id,
)['Data']['member_data']['chatroom_member_list']
member_result = await asyncio.to_thread(self.bot.get_chatroom_member_detail, target_id)
member_info = member_result['Data']['member_data']['chatroom_member_list']
# 处理消息组件
for msg in content_list:
@@ -623,11 +649,35 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
}
if handler := handler_map.get(msg['type']):
handler(msg)
await asyncio.to_thread(handler, msg)
else:
self.logger.warning(f'未处理的消息类型: {msg["type"]}')
await self.logger.warning(f'未处理的消息类型: {msg["type"]}')
continue
def _schedule_ws_message(self, data: dict) -> None:
loop = self._event_loop
if loop is None or loop.is_closed() or self._stop_event.is_set():
return
with self._callback_futures_lock:
if len(self._callback_futures) >= self._MAX_CALLBACK_FUTURES:
return
future = asyncio.run_coroutine_threadsafe(self.ws_message(data), loop)
self._callback_futures.add(future)
def done(completed) -> None:
with self._callback_futures_lock:
self._callback_futures.discard(completed)
if completed.cancelled():
return
try:
completed.result()
except asyncio.CancelledError:
pass
except Exception:
logging.getLogger(__name__).exception('WeChatPad callback failed')
future.add_done_callback(done)
async def send_message(self, target_type: str, target_id: str, message: platform_message.MessageChain):
"""主动发送消息"""
return await self._handle_message(message, target_id)
@@ -665,86 +715,113 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
pass
async def run_async(self):
self._event_loop = asyncio.get_running_loop()
self._stop_event.clear()
if not self.config['admin_key'] and not self.config['token']:
raise RuntimeError('无wechatpad管理密匙，请填入配置文件后重启')
else:
if self.config['token']:
self.bot = WeChatPadClient(self.config['wechatpad_url'], self.config['token'])
data = self.bot.get_login_status()
data = await asyncio.to_thread(self.bot.get_login_status)
if data['Code'] == 300 and data['Text'] == '你已退出微信':
response = requests.post(
response = await asyncio.to_thread(
requests.post,
f'{self.config["wechatpad_url"]}/admin/GenAuthKey1?key={self.config["admin_key"]}',
json={'Count': 1, 'Days': 365},
timeout=10,
)
if response.status_code != 200:
raise Exception(f'获取token失败: {response.text}')
self.config['token'] = response.json()['Data'][0]
body = await httpclient.response_text(response)
raise Exception(f'获取token失败: {body}')
response_data = await httpclient.parse_json_response(response)
self.config['token'] = response_data['Data'][0]
elif not self.config['token']:
response = requests.post(
response = await asyncio.to_thread(
requests.post,
f'{self.config["wechatpad_url"]}/admin/GenAuthKey1?key={self.config["admin_key"]}',
json={'Count': 1, 'Days': 365},
timeout=10,
)
if response.status_code != 200:
raise Exception(f'获取token失败: {response.text}')
self.config['token'] = response.json()['Data'][0]
body = await httpclient.response_text(response)
raise Exception(f'获取token失败: {body}')
response_data = await httpclient.parse_json_response(response)
self.config['token'] = response_data['Data'][0]
self.bot = WeChatPadClient(self.config['wechatpad_url'], self.config['token'], logger=self.logger)
await self.logger.info(self.config['token'])
thread_1 = threading.Event()
def wechat_login_process():
# 不登录，这些先注释掉，避免登陆态尝试拉qrcode。
# login_data =self.bot.get_login_qr()
# url = login_data['Data']["QrCodeUrl"]
profile = self.bot.get_profile()
# self.logger.info(profile)
self.bot_account_id = profile['Data']['userInfo']['nickName']['str']
self.config['wxid'] = profile['Data']['userInfo']['userName']['str']
thread_1.set()
# asyncio.create_task(wechat_login_process)
threading.Thread(target=wechat_login_process).start()
profile = await asyncio.to_thread(self.bot.get_profile)
self.bot_account_id = profile['Data']['userInfo']['nickName']['str']
self.config['wxid'] = profile['Data']['userInfo']['userName']['str']
def connect_websocket_sync() -> None:
thread_1.wait()
uri = f'{self.config["wechatpad_ws"]}/GetSyncMsg?key={self.config["token"]}'
print(f'Connecting to WebSocket: {uri}')
def on_message(ws, message):
try:
if len(message) > _MAX_GATEWAY_MESSAGE_CHARS:
logging.getLogger(__name__).warning('WeChatPad WebSocket message exceeds the size limit')
return
data = json.loads(message)
# 这里需要确保ws_message是同步的，或者使用asyncio.run调用异步方法
asyncio.run(self.ws_message(data))
self._schedule_ws_message(data)
except json.JSONDecodeError:
self.logger.error(f'Non-JSON message: {message[:100]}...')
logging.getLogger(__name__).warning('WeChatPad received a non-JSON message')
def on_error(ws, error):
self.logger.error(f'WebSocket error: {str(error)[:200]}')
logging.getLogger(__name__).warning('WeChatPad WebSocket error: %s', str(error)[:200])
def on_close(ws, close_status_code, close_msg):
self.logger.info('WebSocket closed, reconnecting...')
time.sleep(5)
connect_websocket_sync() # 自动重连
logging.getLogger(__name__).info('WeChatPad WebSocket closed')
def on_open(ws):
self.logger.info('WebSocket connected successfully!')
logging.getLogger(__name__).info('WeChatPad WebSocket connected')
ws = websocket.WebSocketApp(
uri, on_message=on_message, on_error=on_error, on_close=on_close, on_open=on_open
)
ws.run_forever(ping_interval=60, ping_timeout=20)
while not self._stop_event.is_set():
ws = websocket.WebSocketApp(
uri,
on_message=on_message,
on_error=on_error,
on_close=on_close,
on_open=on_open,
)
self._ws_app = ws
ws.run_forever(ping_interval=60, ping_timeout=20)
self._ws_app = None
if not self._stop_event.wait(5):
logging.getLogger(__name__).info('Reconnecting WeChatPad WebSocket')
# 直接调用同步版本（会阻塞）
# connect_websocket_sync()
# 这行代码会在WebSocket连接断开后才会执行
thread = threading.Thread(target=connect_websocket_sync, name='WebSocketClientThread', daemon=True)
thread.start()
self.logger.info('WebSocket client thread started')
self._ws_thread = threading.Thread(
target=connect_websocket_sync,
name='WebSocketClientThread',
daemon=True,
)
self._ws_thread.start()
await self.logger.info('WebSocket client thread started')
while not self._stop_event.is_set() and self._ws_thread.is_alive():
await asyncio.sleep(1)
if not self._stop_event.is_set():
raise RuntimeError('WeChatPad WebSocket client thread exited unexpectedly')
async def kill(self) -> bool:
pass
self._stop_event.set()
ws = self._ws_app
if ws is not None:
await bounded_executor.run_blocking_cleanup(ws.close)
with self._callback_futures_lock:
futures = list(self._callback_futures)
for future in futures:
future.cancel()
if futures:
await asyncio.gather(
*(asyncio.wrap_future(future) for future in futures),
return_exceptions=True,
)
thread = self._ws_thread
if thread is not None and thread.is_alive():
await bounded_executor.run_blocking_cleanup(thread.join, 5)
self._ws_thread = None
self._ws_app = None
self._event_loop = None
return True
src/langbot/pkg/platform/sources/wecom.py
+1
View File
@@ -319,6 +319,7 @@ class WecomAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await keep_alive()
async def kill(self) -> bool:
await self.bot.close()
return False
async def unregister_listener(
src/langbot/pkg/platform/sources/wecombot.py
+11
View File
@@ -516,6 +516,8 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# do work. Lazy-create on first call.
object.__setattr__(self, '_synthetic_buffers', {})
buffers: dict[str, str] = self._synthetic_buffers
if buf_key not in buffers and len(buffers) >= 100:
buffers.pop(next(iter(buffers)), None)
if content and not form_data:
previous = buffers.get(buf_key, '')
if previous and content.startswith(previous):
@@ -524,6 +526,8 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
buffers[buf_key] = previous
else:
buffers[buf_key] = previous + content
if len(buffers[buf_key]) > 200000:
buffers[buf_key] = buffers[buf_key][-200000:]
if not is_final:
return {'stream': True, 'synthetic': True, 'buffered': True}
@@ -613,7 +617,11 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
'chat_id': chat_id,
'stream_id': '',
'req_id': '',
'created_at': time.monotonic(),
}
prune = getattr(self.bot, '_prune_pending_forms', None)
if callable(prune):
prune()
return payload
async def send_message(self, target_type, target_id, message):
@@ -745,10 +753,13 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await keep_alive()
async def kill(self) -> bool:
if hasattr(self, '_synthetic_buffers'):
self._synthetic_buffers.clear()
_ws_mode = not self.config.get('enable-webhook', False)
if _ws_mode:
await self.bot.disconnect()
return True
await self.bot.close()
return False
async def unregister_listener(
src/langbot/pkg/platform/sources/wecomcs.py
+2
View File
@@ -254,6 +254,8 @@ class WecomCSAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await keep_alive()
async def kill(self) -> bool:
self.bot.clear()
await self.bot.close()
return False
async def is_muted(self, group_id: int) -> bool:
src/langbot/pkg/platform/webhook_pusher.py
+2 -1
View File
@@ -147,7 +147,8 @@ class WebhookPusher:
else:
self.logger.debug(f'Successfully pushed to webhook {url}')
try:
return await response.json()
result = await httpclient.read_json_limited(response)
return result if isinstance(result, dict) else None
except Exception as json_error:
self.logger.debug(f'Failed to parse JSON response from webhook {url}: {json_error}')
return None
src/langbot/pkg/plugin/archive.py
+98
View File
@@ -0,0 +1,98 @@
from __future__ import annotations
import io
import zipfile
import yaml
_PLUGIN_ARCHIVE_MAX_ENTRIES = 512
_PLUGIN_ARCHIVE_MAX_ENTRY_BYTES = 16 * 1024 * 1024
_PLUGIN_ARCHIVE_MAX_TOTAL_BYTES = 64 * 1024 * 1024
_PLUGIN_ARCHIVE_MAX_COMPRESSION_RATIO = 100
_PLUGIN_METADATA_MAX_BYTES = 1024 * 1024
_PLUGIN_REQUIREMENTS_MAX_ENTRIES = 1000
def _read_plugin_archive_member(
archive: zipfile.ZipFile,
member: zipfile.ZipInfo,
*,
max_bytes: int = _PLUGIN_METADATA_MAX_BYTES,
) -> bytes:
if member.file_size > max_bytes:
raise ValueError(f'Plugin metadata file exceeds the {max_bytes}-byte limit: {member.filename}')
with archive.open(member, 'r') as source:
content = source.read(max_bytes + 1)
if len(content) > max_bytes or len(content) != member.file_size:
raise ValueError(f'Plugin metadata file has an invalid size: {member.filename}')
return content
def inspect_plugin_archive_metadata(
file_bytes: bytes,
*,
require_manifest: bool = True,
) -> tuple[dict, list[str], list[str]]:
"""Validate archive size metadata and read only bounded preview fields."""
with zipfile.ZipFile(io.BytesIO(file_bytes)) as archive:
members = archive.infolist()
if len(members) > _PLUGIN_ARCHIVE_MAX_ENTRIES:
raise ValueError('Plugin archive contains too many entries')
total_uncompressed = 0
files: dict[str, zipfile.ZipInfo] = {}
names: list[str] = []
for member in members:
if member.is_dir():
continue
if member.flag_bits & 0x1:
raise ValueError('Encrypted plugin archives are not supported')
if member.file_size > _PLUGIN_ARCHIVE_MAX_ENTRY_BYTES:
raise ValueError(f'Plugin archive entry exceeds the size limit: {member.filename}')
if (
member.file_size
and member.file_size > max(member.compress_size, 1) * _PLUGIN_ARCHIVE_MAX_COMPRESSION_RATIO
):
raise ValueError(f'Plugin archive entry exceeds the compression-ratio limit: {member.filename}')
total_uncompressed += member.file_size
if total_uncompressed > _PLUGIN_ARCHIVE_MAX_TOTAL_BYTES:
raise ValueError('Plugin archive exceeds the uncompressed size limit')
normalized = member.filename.replace('\\', '/').strip('/')
names.append(member.filename)
files.setdefault(normalized.lower(), member)
manifest_member = files.get('manifest.yaml') or files.get('manifest.yml')
if manifest_member is None:
if require_manifest:
raise ValueError('manifest.yaml is required')
manifest = {}
else:
manifest = yaml.safe_load(_read_plugin_archive_member(archive, manifest_member).decode('utf-8')) or {}
if not isinstance(manifest, dict):
raise ValueError('Plugin manifest must be an object')
requirements: list[str] = []
requirements_member = next(
(
member
for normalized, member in files.items()
if normalized == 'requirements.txt' or normalized.endswith('/requirements.txt')
),
None,
)
if requirements_member is not None:
content = _read_plugin_archive_member(
archive,
requirements_member,
).decode(
'utf-8',
errors='ignore',
)
requirements = [
line.strip()[:1000]
for line in content.splitlines()
if line.strip() and not line.strip().startswith('#')
][:_PLUGIN_REQUIREMENTS_MAX_ENTRIES]
return manifest, requirements, names
src/langbot/pkg/plugin/connector.py
+158 -65
View File
@@ -5,10 +5,9 @@ import asyncio
import contextlib
import contextvars
import hashlib
import io
import json
import time
import uuid
import zipfile
from typing import Any
import typing
import os
@@ -16,17 +15,17 @@ import secrets
import sys
import httpx
import sqlalchemy
import yaml
from urllib.parse import urljoin, urlparse
from langbot_plugin.api.entities.builtin.pipeline.query import provider_session
from ..core import app
from . import handler
from .archive import inspect_plugin_archive_metadata
from .github import (
validate_github_plugin_install_info,
validate_github_release_asset_url,
)
from ..utils import constants, platform
from ..utils import constants, httpclient, platform
from ..utils.managed_runtime import ManagedRuntimeConnector
from langbot_plugin.runtime.io.controllers.stdio import (
client as stdio_client_controller,
@@ -64,6 +63,9 @@ _PLUGIN_ARTIFACT_OWNER_TYPE = 'plugin_artifact'
_PLUGIN_ARTIFACT_KEY = 'package.lbpkg'
_PLUGIN_ARTIFACT_STORAGE_MARKER = 'tenant_binary_storage_v1'
_GITHUB_PLUGIN_DOWNLOAD_MAX_BYTES = 10 * 1024 * 1024
_MARKETPLACE_METADATA_MAX_BYTES = 1024 * 1024
_MARKETPLACE_PLUGIN_DOWNLOAD_MAX_BYTES = 64 * 1024 * 1024
_MARKETPLACE_SKILL_DOWNLOAD_MAX_BYTES = 10 * 1024 * 1024
_GITHUB_PLUGIN_DOWNLOAD_MAX_REDIRECTS = 5
_GITHUB_ASSET_HOSTS = frozenset(
{
@@ -80,6 +82,55 @@ _HEARTBEAT_FAILURE_THRESHOLD = 3
_RECONNECT_MAX_DELAY_SEC = 60.0
async def _read_httpx_response_limited(
response: httpx.Response,
*,
max_bytes: int,
) -> bytes:
content_length = response.headers.get('content-length')
if content_length is not None:
try:
declared_size = int(content_length)
except ValueError:
declared_size = None
if declared_size is not None and declared_size > max_bytes:
raise ValueError(f'Remote response exceeds the {max_bytes}-byte limit')
body = bytearray()
async for chunk in response.aiter_bytes(chunk_size=64 * 1024):
body.extend(chunk)
if len(body) > max_bytes:
raise ValueError(f'Remote response exceeds the {max_bytes}-byte limit')
return bytes(body)
async def _marketplace_get(
client: httpx.AsyncClient,
url: str,
*,
max_bytes: int,
allow_not_found: bool = False,
) -> tuple[int, bytes]:
async with client.stream('GET', url) as response:
if allow_not_found and response.status_code == 404:
return response.status_code, b''
response.raise_for_status()
return response.status_code, await _read_httpx_response_limited(
response,
max_bytes=max_bytes,
)
def _decode_json_object(body: bytes, *, subject: str) -> dict[str, Any]:
try:
payload = json.loads(body)
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
raise ValueError(f'{subject} returned invalid JSON') from exc
if not isinstance(payload, dict):
raise ValueError(f'{subject} returned a non-object response')
return payload
class PluginRuntimeNotConnectedError(RuntimeError):
"""Raised when plugin runtime operations are requested before connection."""
@@ -175,16 +226,23 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
worker = self.ap.instance_config.data.get('plugin', {}).get('worker')
if not isinstance(worker, dict):
raise ValueError('plugin.worker must be configured')
return PluginWorkerPolicy.model_validate(
{
'max_cpus': worker.get('max_cpus'),
'max_memory_mb': worker.get('max_memory_mb'),
'max_pids': worker.get('max_pids'),
'max_open_files': worker.get('max_open_files'),
'max_file_size_mb': worker.get('max_file_size_mb'),
'require_hard_limits': worker.get('require_hard_limits', False),
}
)
policy_data = {
'max_cpus': worker.get('max_cpus'),
'max_memory_mb': worker.get('max_memory_mb'),
'max_pids': worker.get('max_pids'),
'max_open_files': worker.get('max_open_files'),
'max_file_size_mb': worker.get('max_file_size_mb'),
'require_hard_limits': worker.get('require_hard_limits', False),
}
for field_name in (
'max_workers',
'max_total_cpus',
'max_total_memory_mb',
'max_installations',
):
if field_name in PluginWorkerPolicy.model_fields and field_name in worker:
policy_data[field_name] = worker.get(field_name)
return PluginWorkerPolicy.model_validate(policy_data)
def _control_headers(self, *, allow_generate: bool) -> dict[str, str]:
if not self._control_token and allow_generate:
@@ -643,9 +701,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
for context in contexts:
execution_context = await self._validate_execution_context(context)
states = await self._load_workspace_desired_states(execution_context)
workspace_installations[execution_context.workspace_uuid] = {
state.binding.installation_uuid for state in states
}
installation_ids = {state.binding.installation_uuid for state in states}
if installation_ids:
# A newly registered Cloud Workspace normally has no
# plugins. Avoid retaining an empty set for every account.
workspace_installations[execution_context.workspace_uuid] = installation_ids
for state in states:
if state.binding.installation_uuid in all_states:
raise ValueError('Duplicate plugin installation UUID across projected Workspaces')
@@ -706,7 +766,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
# restoring the remaining desired state in this Workspace.
pass
self._known_desired_states[installation_uuid] = desired
self._workspace_installations[execution_context.workspace_uuid] = set(desired_by_uuid)
if desired_by_uuid:
self._workspace_installations[execution_context.workspace_uuid] = set(desired_by_uuid)
else:
self._workspace_installations.pop(
execution_context.workspace_uuid,
None,
)
async def _current_execution_context(self) -> ExecutionContext:
current = self._execution_context.get()
@@ -991,27 +1057,20 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
plugin_name = None
try:
with zipfile.ZipFile(io.BytesIO(file_bytes)) as zf:
try:
manifest = yaml.safe_load(zf.read('manifest.yaml').decode('utf-8', errors='ignore')) or {}
metadata = manifest.get('metadata', {})
plugin_author = metadata.get('author')
plugin_name = metadata.get('name')
except Exception:
pass
if task_context is not None:
for name in zf.namelist():
if name.endswith('requirements.txt'):
content = zf.read(name).decode('utf-8', errors='ignore')
deps = [
line.strip()
for line in content.splitlines()
if line.strip() and not line.strip().startswith('#')
]
task_context.metadata['deps_total'] = len(deps)
task_context.metadata['deps_list'] = deps
break
manifest, dependencies, archive_names = inspect_plugin_archive_metadata(
file_bytes,
require_manifest=False,
)
metadata = manifest.get('metadata', {})
if isinstance(metadata, dict):
plugin_author = metadata.get('author')
plugin_name = metadata.get('name')
has_requirements = any(
name.replace('\\', '/').lower().endswith('requirements.txt') for name in archive_names
)
if task_context is not None and has_requirements:
task_context.metadata['deps_total'] = len(dependencies)
task_context.metadata['deps_list'] = dependencies
except Exception:
pass
@@ -1329,6 +1388,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
trust_env=False,
follow_redirects=False,
timeout=httpx.Timeout(60, connect=10),
event_hooks=httpclient.httpx_response_limit_hooks(_GITHUB_PLUGIN_DOWNLOAD_MAX_BYTES),
) as client:
asset_id: int | None = None
if 'asset_id' in normalized:
@@ -1346,7 +1406,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if response.status_code in _HTTP_REDIRECT_STATUSES:
raise ValueError('GitHub release metadata unexpectedly redirected')
response.raise_for_status()
release = response.json()
release = await httpclient.parse_json_response(response)
if not isinstance(release, dict):
raise ValueError('GitHub release metadata is invalid')
if release.get('id') != release_id or str(release.get('tag_name') or '') != release_tag:
@@ -1500,50 +1560,79 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
"""Return a plugin package, or install an MCP/skill and return none."""
space_url = self.ap.instance_config.data.get('space', {}).get('url', 'https://space.langbot.app').rstrip('/')
async with httpx.AsyncClient(trust_env=True, timeout=15) as client:
mcp_resp = await client.get(f'{space_url}/api/v1/marketplace/mcps/{plugin_author}/{plugin_name}')
if mcp_resp.status_code == 200:
mcp_data = mcp_resp.json().get('data', {}).get('mcp', {})
async with httpx.AsyncClient(
trust_env=True,
timeout=15,
event_hooks=httpclient.httpx_response_limit_hooks(_MARKETPLACE_PLUGIN_DOWNLOAD_MAX_BYTES),
) as client:
mcp_status, mcp_body = await _marketplace_get(
client,
f'{space_url}/api/v1/marketplace/mcps/{plugin_author}/{plugin_name}',
max_bytes=_MARKETPLACE_METADATA_MAX_BYTES,
allow_not_found=True,
)
if mcp_status == 200:
mcp_payload = _decode_json_object(mcp_body, subject='Marketplace MCP metadata')
mcp_data = mcp_payload.get('data', {}).get('mcp', {})
if not isinstance(mcp_data, dict):
raise ValueError(f'MCP {plugin_author}/{plugin_name} metadata is invalid')
if not mcp_data.get('mode'):
raise ValueError(f'MCP {plugin_author}/{plugin_name} has no mode')
await self._install_mcp_from_marketplace(execution_context, mcp_data, task_context)
try:
await client.post(f'{space_url}/api/v1/marketplace/mcps/{plugin_author}/{plugin_name}/install')
async with client.stream(
'POST',
f'{space_url}/api/v1/marketplace/mcps/{plugin_author}/{plugin_name}/install',
):
pass
except Exception as report_err:
self.ap.logger.debug(f'Failed to report MCP install: {report_err}')
return None, None
if mcp_resp.status_code != 404:
mcp_resp.raise_for_status()
skill_resp = await client.get(f'{space_url}/api/v1/marketplace/skills/{plugin_author}/{plugin_name}')
if skill_resp.status_code == 200:
download_resp = await client.get(
f'{space_url}/api/v1/marketplace/skills/download/{plugin_author}/{plugin_name}'
skill_status, _skill_body = await _marketplace_get(
client,
f'{space_url}/api/v1/marketplace/skills/{plugin_author}/{plugin_name}',
max_bytes=_MARKETPLACE_METADATA_MAX_BYTES,
allow_not_found=True,
)
if skill_status == 200:
_download_status, skill_package = await _marketplace_get(
client,
f'{space_url}/api/v1/marketplace/skills/download/{plugin_author}/{plugin_name}',
max_bytes=_MARKETPLACE_SKILL_DOWNLOAD_MAX_BYTES,
)
download_resp.raise_for_status()
await self._install_skill_from_zip(
execution_context,
download_resp.content,
skill_package,
f'{plugin_author}-{plugin_name}',
task_context,
)
return None, None
if skill_resp.status_code != 404:
skill_resp.raise_for_status()
versions_resp = await client.get(
f'{space_url}/api/v1/marketplace/plugins/{plugin_author}/{plugin_name}/versions'
_versions_status, versions_body = await _marketplace_get(
client,
f'{space_url}/api/v1/marketplace/plugins/{plugin_author}/{plugin_name}/versions',
max_bytes=_MARKETPLACE_METADATA_MAX_BYTES,
)
versions_resp.raise_for_status()
versions = versions_resp.json().get('data', {}).get('versions', [])
if not versions or not versions[0].get('version'):
versions_payload = _decode_json_object(
versions_body,
subject='Marketplace plugin versions',
)
versions = versions_payload.get('data', {}).get('versions', [])
if (
not isinstance(versions, list)
or not versions
or not isinstance(versions[0], dict)
or not versions[0].get('version')
):
raise ValueError(f'Plugin {plugin_author}/{plugin_name} has no versions')
latest_version = str(versions[0]['version'])
download_resp = await client.get(
f'{space_url}/api/v1/marketplace/plugins/download/{plugin_author}/{plugin_name}/{latest_version}'
_download_status, plugin_package = await _marketplace_get(
client,
f'{space_url}/api/v1/marketplace/plugins/download/{plugin_author}/{plugin_name}/{latest_version}',
max_bytes=_MARKETPLACE_PLUGIN_DOWNLOAD_MAX_BYTES,
)
download_resp.raise_for_status()
return download_resp.content, latest_version
return plugin_package, latest_version
async def install_plugin(
self,
@@ -1698,7 +1787,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
await delete(self.ap.persistence_mgr.execute_async)
self._known_desired_states.pop(binding.installation_uuid, None)
self._workspace_installations.setdefault(binding.workspace_uuid, set()).discard(binding.installation_uuid)
workspace_installations = self._workspace_installations.get(binding.workspace_uuid)
if workspace_installations is not None:
workspace_installations.discard(binding.installation_uuid)
if not workspace_installations:
self._workspace_installations.pop(binding.workspace_uuid, None)
if task_context is not None:
task_context.set_current_action('plugin removed')
return {}
src/langbot/pkg/plugin/handler.py
+36 -17
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import inspect
import typing
from typing import Any
@@ -43,6 +44,24 @@ from ..entity.persistence import model as persistence_model
from ..core import app
from ..utils import constants
_DEFAULT_BINARY_STORAGE_VALUE_BYTES = 10 * 1024 * 1024
_HARD_MAX_BINARY_STORAGE_VALUE_BYTES = 64 * 1024 * 1024
def _binary_storage_value_limit(ap: Any) -> int:
configured = (
ap.instance_config.data.get('plugin', {})
.get('binary_storage', {})
.get('max_value_bytes', _DEFAULT_BINARY_STORAGE_VALUE_BYTES)
)
try:
configured = int(configured)
except (TypeError, ValueError):
configured = _DEFAULT_BINARY_STORAGE_VALUE_BYTES
if configured < 0:
configured = _DEFAULT_BINARY_STORAGE_VALUE_BYTES
return min(configured, _HARD_MAX_BINARY_STORAGE_VALUE_BYTES)
class _RawAction:
def __init__(self, value: str):
@@ -859,20 +878,15 @@ class RuntimeConnectionHandler(handler.Handler):
return handler.ActionResponse.error(
message=str(e),
)
value = base64.b64decode(data['value_base64'])
max_value_bytes = (
self.ap.instance_config.data.get('plugin', {})
.get('binary_storage', {})
.get(
'max_value_bytes',
10 * 1024 * 1024,
max_value_bytes = _binary_storage_value_limit(self.ap)
encoded_value = data['value_base64']
max_encoded_chars = 4 * ((max_value_bytes + 2) // 3) + 4
if len(encoded_value) > max_encoded_chars:
return handler.ActionResponse.error(
message=f'Binary storage value exceeds the {max_value_bytes}-byte limit',
)
)
try:
max_value_bytes = int(max_value_bytes)
except (TypeError, ValueError):
max_value_bytes = 10 * 1024 * 1024
if max_value_bytes >= 0 and len(value) > max_value_bytes:
value = await asyncio.to_thread(base64.b64decode, encoded_value)
if len(value) > max_value_bytes:
return handler.ActionResponse.error(
message=f'Binary storage value exceeds limit ({len(value)} > {max_value_bytes} bytes)',
)
@@ -936,10 +950,15 @@ class RuntimeConnectionHandler(handler.Handler):
return handler.ActionResponse.error(
message=f'Storage with key {key} not found',
)
max_value_bytes = _binary_storage_value_limit(self.ap)
if len(storage.value) > max_value_bytes:
return handler.ActionResponse.error(
message=f'Binary storage value exceeds the {max_value_bytes}-byte limit',
)
return handler.ActionResponse.success(
data={
'value_base64': base64.b64encode(storage.value).decode('utf-8'),
'value_base64': (await asyncio.to_thread(base64.b64encode, storage.value)).decode('utf-8'),
},
)
@@ -1019,7 +1038,7 @@ class RuntimeConnectionHandler(handler.Handler):
return handler.ActionResponse.success(
data={
'file_base64': base64.b64encode(file_bytes).decode('utf-8'),
'file_base64': (await asyncio.to_thread(base64.b64encode, file_bytes)).decode('utf-8'),
},
)
except Exception as e:
@@ -1744,7 +1763,7 @@ class RuntimeConnectionHandler(handler.Handler):
await self.delete_local_file(plugin_icon_file_key)
return {
'plugin_icon_base64': base64.b64encode(plugin_icon_bytes).decode('utf-8'),
'plugin_icon_base64': (await asyncio.to_thread(base64.b64encode, plugin_icon_bytes)).decode('utf-8'),
'mime_type': mime_type,
}
@@ -1819,7 +1838,7 @@ class RuntimeConnectionHandler(handler.Handler):
asset_bytes = await self.read_local_file(asset_file_key)
await self.delete_local_file(asset_file_key)
return {
'asset_base64': base64.b64encode(asset_bytes).decode('utf-8'),
'asset_base64': (await asyncio.to_thread(base64.b64encode, asset_bytes)).decode('utf-8'),
'mime_type': mime_type,
}
src/langbot/pkg/provider/modelmgr/modelmgr.py
+228 -19
View File
@@ -52,6 +52,128 @@ class ModelManager:
self.rerank_model_dict = {}
self.requester_components = []
self.requester_dict = {}
self._scope_generations: dict[tuple[str, str], int] = {}
self._provider_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
self._llm_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
self._embedding_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
self._rerank_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
def _cache_index(self, cache: dict) -> dict[tuple[str, str], set[_CacheKey]]:
if cache is self.provider_dict:
return self._provider_keys_by_scope
if cache is self.llm_model_dict:
return self._llm_keys_by_scope
if cache is self.embedding_model_dict:
return self._embedding_keys_by_scope
if cache is self.rerank_model_dict:
return self._rerank_keys_by_scope
raise ValueError('Unknown model runtime cache')
def _cache_set(self, cache: dict, key: _CacheKey, value: object) -> None:
cache[key] = value
self._cache_index(cache).setdefault(key[:2], set()).add(key)
def _cache_pop(self, cache: dict, key: _CacheKey) -> object | None:
removed = cache.pop(key, None)
scope = key[:2]
index = self._cache_index(cache)
keys = index.get(scope)
if keys is not None:
keys.discard(key)
if not keys:
index.pop(scope, None)
if not any(
scope in candidate
for candidate in (
self._provider_keys_by_scope,
self._llm_keys_by_scope,
self._embedding_keys_by_scope,
self._rerank_keys_by_scope,
)
):
self._scope_generations.pop(scope, None)
return removed
def _observe_execution_context(
self,
context: ExecutionContext,
) -> tuple[requester.RuntimeProvider, ...]:
"""Prune superseded runtime objects when a Workspace generation advances."""
scope = (context.instance_uuid, context.workspace_uuid)
previous_generation = self._scope_generations.get(scope)
if previous_generation is not None and context.placement_generation < previous_generation:
raise WorkspaceInvariantError('Model runtime placement generation rolled back')
if previous_generation == context.placement_generation:
return ()
retired_providers: list[requester.RuntimeProvider] = []
if previous_generation is not None:
for cache, index in (
(self.provider_dict, self._provider_keys_by_scope),
(self.llm_model_dict, self._llm_keys_by_scope),
(self.embedding_model_dict, self._embedding_keys_by_scope),
(self.rerank_model_dict, self._rerank_keys_by_scope),
):
for key in index.pop(scope, ()):
removed = cache.pop(key, None)
if cache is self.provider_dict and removed is not None:
retired_providers.append(removed)
self._scope_generations[scope] = context.placement_generation
return tuple(retired_providers)
async def _close_runtime_providers(
self,
providers: tuple[requester.RuntimeProvider, ...] | list[requester.RuntimeProvider],
) -> None:
"""Close each retired requester once without blocking other cleanup."""
seen: set[int] = set()
for provider in providers:
provider_id = id(provider)
if provider_id in seen:
continue
seen.add(provider_id)
try:
await provider.requester.aclose()
except Exception as exc:
self.ap.logger.warning(
f'Failed to close model requester for provider {provider.provider_entity.uuid}: {exc}'
)
async def _observe_and_close_execution_context(
self,
context: ExecutionContext,
*,
retain_empty: bool = True,
) -> None:
await self._close_runtime_providers(self._observe_execution_context(context))
if not retain_empty:
scope = (context.instance_uuid, context.workspace_uuid)
if not any(
scope in candidate
for candidate in (
self._provider_keys_by_scope,
self._llm_keys_by_scope,
self._embedding_keys_by_scope,
self._rerank_keys_by_scope,
)
):
self._scope_generations.pop(scope, None)
async def shutdown(self) -> None:
"""Release every requester owned by the model runtime cache."""
providers = list(self.provider_dict.values())
self.provider_dict = {}
self.llm_model_dict = {}
self.embedding_model_dict = {}
self.rerank_model_dict = {}
self._scope_generations = {}
self._provider_keys_by_scope = {}
self._llm_keys_by_scope = {}
self._embedding_keys_by_scope = {}
self._rerank_keys_by_scope = {}
await self._close_runtime_providers(providers)
@staticmethod
def _get_litellm_provider_from_manifest(component: engine.Component | None) -> str | None:
@@ -137,7 +259,17 @@ class ModelManager:
if supplied_instance_uuid is not None and supplied_instance_uuid != binding.instance_uuid:
raise WorkspaceInvariantError('Runtime context belongs to another LangBot instance')
return self._context_from_binding(binding, trigger_principal=trigger_principal)
execution_context = self._context_from_binding(binding, trigger_principal=trigger_principal)
scope = (
execution_context.instance_uuid,
execution_context.workspace_uuid,
)
if scope in self._scope_generations:
await self._observe_and_close_execution_context(
execution_context,
retain_empty=False,
)
return execution_context
async def initialize(self) -> None:
self.requester_components = self.ap.discover.get_components_by_kind('LLMAPIRequester')
@@ -199,10 +331,16 @@ class ModelManager:
"""Load every active projected Workspace into isolated runtime caches."""
self.ap.logger.info('Loading models from db...')
await self._close_runtime_providers(list(self.provider_dict.values()))
self.provider_dict = {}
self.llm_model_dict = {}
self.embedding_model_dict = {}
self.rerank_model_dict = {}
self._scope_generations = {}
self._provider_keys_by_scope = {}
self._llm_keys_by_scope = {}
self._embedding_keys_by_scope = {}
self._rerank_keys_by_scope = {}
list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
@@ -233,6 +371,7 @@ class ModelManager:
binding,
trigger_principal=PrincipalContext(principal_type=PrincipalType.SYSTEM),
)
await self._observe_and_close_execution_context(resolved)
contexts[workspace_uuid] = resolved
return resolved
@@ -243,7 +382,11 @@ class ModelManager:
try:
context = await context_for(provider_entity.workspace_uuid)
runtime_provider = await self._build_provider(context, provider_entity)
self.provider_dict[self._cache_key(context, provider_entity.uuid)] = runtime_provider
self._cache_set(
self.provider_dict,
self._cache_key(context, provider_entity.uuid),
runtime_provider,
)
except provider_errors.RequesterNotFoundError as exc:
self.ap.logger.warning(
f'Requester {exc.requester_name} not found, skipping provider {provider_entity.uuid}'
@@ -282,7 +425,11 @@ class ModelManager:
)
continue
runtime_model = builder(context, model_entity, provider)
cache[self._cache_key(context, model_entity.uuid)] = runtime_model
self._cache_set(
cache,
self._cache_key(context, model_entity.uuid),
runtime_model,
)
except Exception as exc:
self.ap.logger.error(f'Failed to load model {model_entity.uuid}: {exc}\n{traceback.format_exc()}')
@@ -294,10 +441,20 @@ class ModelManager:
persistence_model.ModelProvider.workspace_uuid == context.workspace_uuid
)
)
for provider_entity in providers_result.all():
provider_entities = providers_result.all()
if provider_entities:
# Empty Workspaces are the dominant SaaS registration case. Do
# not retain one generation record per account until the
# Workspace owns an actual runtime model resource.
await self._observe_and_close_execution_context(context)
for provider_entity in provider_entities:
try:
runtime_provider = await self._build_provider(context, provider_entity)
self.provider_dict[self._cache_key(context, provider_entity.uuid)] = runtime_provider
self._cache_set(
self.provider_dict,
self._cache_key(context, provider_entity.uuid),
runtime_provider,
)
except provider_errors.RequesterNotFoundError as exc:
self.ap.logger.warning(
f'Requester {exc.requester_name} not found, skipping provider {provider_entity.uuid}'
@@ -337,7 +494,11 @@ class ModelManager:
)
continue
runtime_model = builder(context, model_entity, provider)
cache[self._cache_key(context, model_entity.uuid)] = runtime_model
self._cache_set(
cache,
self._cache_key(context, model_entity.uuid),
runtime_model,
)
except Exception as exc:
self.ap.logger.error(f'Failed to load model {model_entity.uuid}: {exc}\n{traceback.format_exc()}')
@@ -562,7 +723,12 @@ class ModelManager:
execution_context = await self.resolve_execution_context(context)
self._ensure_same_scope(execution_context, provider.execution_context, resource='Provider')
self._ensure_entity_workspace(provider.provider_entity, execution_context, resource='Provider')
self.provider_dict[self._cache_key(execution_context, provider.provider_entity.uuid)] = provider
self._observe_execution_context(execution_context)
self._cache_set(
self.provider_dict,
self._cache_key(execution_context, provider.provider_entity.uuid),
provider,
)
async def get_provider_by_uuid(
self,
@@ -578,7 +744,12 @@ class ModelManager:
async def remove_provider(self, context: TenantContext, provider_uuid: str) -> None:
execution_context = await self.resolve_execution_context(context)
self.provider_dict.pop(self._cache_key(execution_context, provider_uuid), None)
removed = self._cache_pop(
self.provider_dict,
self._cache_key(execution_context, provider_uuid),
)
if removed is not None:
await self._close_runtime_providers([removed])
async def reload_provider(self, context: TenantContext, provider_uuid: str) -> None:
execution_context = await self.resolve_execution_context(context)
@@ -593,12 +764,26 @@ class ModelManager:
raise provider_errors.ProviderNotFoundError(provider_uuid)
new_provider = await self._build_provider(execution_context, provider_entity)
cache_prefix = self._cache_key(execution_context, '')[:3]
for cache in (self.llm_model_dict, self.embedding_model_dict, self.rerank_model_dict):
for key, model in cache.items():
if key[:3] == cache_prefix and model.provider.provider_entity.uuid == provider_uuid:
scope = (execution_context.instance_uuid, execution_context.workspace_uuid)
for cache, index in (
(self.llm_model_dict, self._llm_keys_by_scope),
(self.embedding_model_dict, self._embedding_keys_by_scope),
(self.rerank_model_dict, self._rerank_keys_by_scope),
):
for key in tuple(index.get(scope, ())):
model = cache.get(key)
if model is not None and model.provider.provider_entity.uuid == provider_uuid:
model.provider = new_provider
self.provider_dict[self._cache_key(execution_context, provider_uuid)] = new_provider
self._observe_execution_context(execution_context)
provider_key = self._cache_key(execution_context, provider_uuid)
old_provider = self.provider_dict.get(provider_key)
self._cache_set(
self.provider_dict,
provider_key,
new_provider,
)
if old_provider is not None and old_provider is not new_provider:
await self._close_runtime_providers([old_provider])
@staticmethod
def _coerce_model(model_info: _ModelEntity | sqlalchemy.Row, entity_type: type[_ModelEntity]) -> _ModelEntity:
@@ -689,7 +874,12 @@ class ModelManager:
async def cache_llm_model(self, context: TenantContext, model: requester.RuntimeLLMModel) -> None:
execution_context = await self.resolve_execution_context(context)
self._ensure_same_scope(execution_context, model.execution_context, resource='LLM model')
self.llm_model_dict[self._cache_key(execution_context, model.model_entity.uuid)] = model
self._observe_execution_context(execution_context)
self._cache_set(
self.llm_model_dict,
self._cache_key(execution_context, model.model_entity.uuid),
model,
)
async def cache_embedding_model(
self,
@@ -698,12 +888,22 @@ class ModelManager:
) -> None:
execution_context = await self.resolve_execution_context(context)
self._ensure_same_scope(execution_context, model.execution_context, resource='Embedding model')
self.embedding_model_dict[self._cache_key(execution_context, model.model_entity.uuid)] = model
self._observe_execution_context(execution_context)
self._cache_set(
self.embedding_model_dict,
self._cache_key(execution_context, model.model_entity.uuid),
model,
)
async def cache_rerank_model(self, context: TenantContext, model: requester.RuntimeRerankModel) -> None:
execution_context = await self.resolve_execution_context(context)
self._ensure_same_scope(execution_context, model.execution_context, resource='Rerank model')
self.rerank_model_dict[self._cache_key(execution_context, model.model_entity.uuid)] = model
self._observe_execution_context(execution_context)
self._cache_set(
self.rerank_model_dict,
self._cache_key(execution_context, model.model_entity.uuid),
model,
)
async def get_model_by_uuid(self, context: TenantContext, model_uuid: str) -> requester.RuntimeLLMModel:
execution_context = await self.resolve_execution_context(context)
@@ -739,15 +939,24 @@ class ModelManager:
async def remove_llm_model(self, context: TenantContext, model_uuid: str) -> None:
execution_context = await self.resolve_execution_context(context)
self.llm_model_dict.pop(self._cache_key(execution_context, model_uuid), None)
self._cache_pop(
self.llm_model_dict,
self._cache_key(execution_context, model_uuid),
)
async def remove_embedding_model(self, context: TenantContext, model_uuid: str) -> None:
execution_context = await self.resolve_execution_context(context)
self.embedding_model_dict.pop(self._cache_key(execution_context, model_uuid), None)
self._cache_pop(
self.embedding_model_dict,
self._cache_key(execution_context, model_uuid),
)
async def remove_rerank_model(self, context: TenantContext, model_uuid: str) -> None:
execution_context = await self.resolve_execution_context(context)
self.rerank_model_dict.pop(self._cache_key(execution_context, model_uuid), None)
self._cache_pop(
self.rerank_model_dict,
self._cache_key(execution_context, model_uuid),
)
def get_available_requesters_info(self, model_type: str) -> list[dict]:
if model_type:
src/langbot/pkg/provider/modelmgr/requester.py
+11
View File
@@ -463,6 +463,17 @@ class ProviderAPIRequester(metaclass=abc.ABCMeta):
async def initialize(self):
pass
async def aclose(self) -> None:
"""Release requester-owned clients when its runtime provider retires.
Most built-in requesters are currently stateless, but provider
extensions may own connection pools or background resources. Keeping
the lifecycle hook on the base class lets Workspace generation changes
and application shutdown retire them deterministically.
"""
return None
async def scan_models(self, api_key: str | None = None) -> dict[str, typing.Any] | list[dict[str, typing.Any]]:
"""Scan models supported by the provider.
src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py
+13 -5
View File
@@ -8,6 +8,7 @@ import litellm
from litellm import acompletion, aembedding, arerank
from .. import errors, requester
from ....utils import httpclient
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
@@ -955,14 +956,17 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
rerank_url = f'{base_url}/rerank'
try:
async with httpx.AsyncClient(timeout=timeout) as client:
async with httpx.AsyncClient(
timeout=timeout,
event_hooks=httpclient.httpx_response_limit_hooks(),
) as client:
resp = await client.post(rerank_url, headers=headers, json=payload)
resp.raise_for_status()
data = resp.json()
data = await httpclient.parse_json_response(resp)
except httpx.HTTPStatusError as e:
body = ''
try:
body = e.response.text
body = await httpclient.response_text(e.response)
except Exception:
pass
raise errors.RequesterError(f'rerank 请求失败 (HTTP {e.response.status_code}): {body or str(e)}')
@@ -998,10 +1002,14 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
models_url = f'{base_url}/models'
try:
async with httpx.AsyncClient(trust_env=True, timeout=timeout) as client:
async with httpx.AsyncClient(
trust_env=True,
timeout=timeout,
event_hooks=httpclient.httpx_response_limit_hooks(),
) as client:
response = await client.get(models_url, headers=headers)
response.raise_for_status()
payload = response.json()
payload = await httpclient.parse_json_response(response)
models = []
for item in payload.get('data', []):
src/langbot/pkg/provider/runner.py
+30
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import abc
import asyncio
import typing
from typing import TYPE_CHECKING
@@ -11,6 +12,32 @@ if TYPE_CHECKING:
preregistered_runners: list[typing.Type[RequestRunner]] = []
_DEFAULT_SYNC_ITERATION_LIMIT = 100_000
_T = typing.TypeVar('_T')
def _next_sync(iterator: typing.Iterator[_T]) -> tuple[bool, _T | None]:
try:
return True, next(iterator)
except StopIteration:
return False, None
async def iterate_sync(
iterable: typing.Iterable[_T],
*,
max_items: int = _DEFAULT_SYNC_ITERATION_LIMIT,
) -> typing.AsyncGenerator[_T, None]:
"""Consume a blocking SDK iterator without stalling the event loop."""
iterator = iter(iterable)
for _ in range(max(max_items, 1)):
has_item, item = await asyncio.to_thread(_next_sync, iterator)
if not has_item:
return
yield typing.cast(_T, item)
raise RuntimeError('Synchronous provider stream exceeded the event limit')
def runner_class(name: str):
@@ -43,3 +70,6 @@ class RequestRunner(abc.ABC):
) -> typing.AsyncGenerator[provider_message.Message | provider_message.MessageChunk, None]:
"""运行请求"""
pass
async def aclose(self) -> None:
"""Release request-scoped resources after one runner invocation."""
src/langbot/pkg/provider/runners/cozeapi.py
+30 -7
View File
@@ -2,7 +2,6 @@ from __future__ import annotations
import typing
import json
import base64
from langbot.pkg.provider import runner
from langbot.pkg.core import app
@@ -11,6 +10,16 @@ from langbot.pkg.utils import image
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from langbot.libs.coze_server_api.client import AsyncCozeAPIClient
_MAX_COZE_GENERATED_CHARS = 1024 * 1024
_MAX_COZE_MEDIA_BYTES = 10 * 1024 * 1024
def _append_bounded(current: str, addition: typing.Any) -> str:
addition = str(addition or '')
if len(current) + len(addition) > _MAX_COZE_GENERATED_CHARS:
raise ValueError('Coze response exceeds the runtime limit')
return current + addition
@runner.runner_class('coze-api')
class CozeAPIRunner(runner.RequestRunner):
@@ -77,7 +86,10 @@ class CozeAPIRunner(runner.RequestRunner):
content_parts.append({'type': 'text', 'text': ce.text})
elif ce.type == 'image_base64':
image_b64, image_format = await image.extract_b64_and_format(ce.image_base64)
file_bytes = base64.b64decode(image_b64)
file_bytes = await image.decode_base64_limited(
image_b64,
max_bytes=_MAX_COZE_MEDIA_BYTES,
)
file_id = await self._get_file_id(file_bytes)
content_parts.append({'type': 'image', 'file_id': file_id})
elif ce.type == 'file':
@@ -144,7 +156,7 @@ class CozeAPIRunner(runner.RequestRunner):
auto_save_history=self.auto_save_history,
stream=True,
):
self.ap.logger.debug(f'coze-chat-stream: {chunk}')
self.ap.logger.debug(f'coze-chat-stream: {str(chunk)[:1000]}')
event_type = chunk.get('event')
data = chunk.get('data', {})
@@ -153,11 +165,17 @@ class CozeAPIRunner(runner.RequestRunner):
if event_type == 'conversation.message.delta':
# 收集内容
if 'content' in data:
full_content += data.get('content', '')
full_content = _append_bounded(
full_content,
data.get('content', ''),
)
# 收集推理内容（如果有）
if 'reasoning_content' in data:
full_reasoning += data.get('reasoning_content', '')
full_reasoning = _append_bounded(
full_reasoning,
data.get('reasoning_content', ''),
)
elif event_type.split('.')[-1] == 'done': # 本地部署coze时，结束event不为done
# 保存会话ID
@@ -179,6 +197,8 @@ class CozeAPIRunner(runner.RequestRunner):
remove_think = self.pipeline_config.get('output', {}).get('misc', {}).get('remove-think', False)
if not remove_think:
content = f'<think>\n{full_reasoning}\n</think>\n{content}'.strip()
if len(content) > _MAX_COZE_GENERATED_CHARS:
raise ValueError('Coze response exceeds the runtime limit')
# 一次性返回完整内容
yield provider_message.Message(
@@ -227,7 +247,7 @@ class CozeAPIRunner(runner.RequestRunner):
auto_save_history=self.auto_save_history,
stream=True,
):
self.ap.logger.debug(f'coze-chat-stream-chunk: {chunk}')
self.ap.logger.debug(f'coze-chat-stream-chunk: {str(chunk)[:1000]}')
event_type = chunk.get('event')
data = chunk.get('data', {})
@@ -263,7 +283,7 @@ class CozeAPIRunner(runner.RequestRunner):
error_msg = f'Coze API错误: {data.get("message", "未知错误")}'
yield provider_message.MessageChunk(role='assistant', content=error_msg, finish_reason='error')
return
full_content += content
full_content = _append_bounded(full_content, content)
if message_idx % 8 == 0 or is_final:
if full_content:
yield provider_message.MessageChunk(role='assistant', content=full_content, is_final=is_final)
@@ -286,3 +306,6 @@ class CozeAPIRunner(runner.RequestRunner):
else:
async for msg in self._chat_messages(query):
yield msg
async def aclose(self) -> None:
await self.coze.close()
src/langbot/pkg/provider/runners/dashscopeapi.py
+87 -42
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import typing
import re
@@ -10,6 +11,9 @@ from ...core import app
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
_MAX_DASHSCOPE_RESPONSE_CHARS = 1024 * 1024
_MAX_DASHSCOPE_REFERENCES = 1024
class DashscopeAPIError(Exception):
"""Dashscope API 请求失败"""
@@ -19,6 +23,13 @@ class DashscopeAPIError(Exception):
super().__init__(self.message)
def _append_bounded(current: str, addition: typing.Any) -> str:
addition = str(addition or '')
if len(current) + len(addition) > _MAX_DASHSCOPE_RESPONSE_CHARS:
raise DashscopeAPIError('Dashscope response exceeds the runtime limit')
return current + addition
@runner.runner_class('dashscope-app-api')
class DashScopeAPIRunner(runner.RequestRunner):
"阿里云百炼DashsscopeAPI对话请求器"
@@ -111,18 +122,16 @@ class DashScopeAPIRunner(runner.RequestRunner):
if remove_think:
has_thoughts = False
# 发送对话请求
response = dashscope.Application.call(
api_key=self.api_key, # 智能体应用的API Key
app_id=self.app_id, # 智能体应用的ID
prompt=plain_text, # 用户输入的文本信息
stream=True, # 流式输出
incremental_output=True, # 增量输出，使用流式输出需要开启增量输出
session_id=query.session.using_conversation.uuid, # 会话ID用于，多轮对话
response = await asyncio.to_thread(
dashscope.Application.call,
api_key=self.api_key,
app_id=self.app_id,
prompt=plain_text,
stream=True,
incremental_output=True,
session_id=query.session.using_conversation.uuid,
enable_thinking=has_thoughts,
has_thoughts=has_thoughts,
# rag_options={ # 主要用于文件交互，暂不支持
# "session_file_ids": ["FILE_ID1"], # FILE_ID1 替换为实际的临时文件ID,逗号隔开多个
# }
)
idx_chunk = 0
try:
@@ -131,7 +140,7 @@ class DashScopeAPIRunner(runner.RequestRunner):
except AttributeError:
is_stream = False
if is_stream:
for chunk in response:
async for chunk in runner.iterate_sync(response):
if chunk.get('status_code') != 200:
raise DashscopeAPIError(
f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '
@@ -145,15 +154,27 @@ class DashScopeAPIRunner(runner.RequestRunner):
if stream_think and stream_think[0].get('thought'):
if not think_start:
think_start = True
pending_content += f'<think>\n{stream_think[0].get("thought")}'
pending_content = _append_bounded(
pending_content,
f'<think>\n{stream_think[0].get("thought")}',
)
else:
# 继续输出 reasoning_content
pending_content += stream_think[0].get('thought')
pending_content = _append_bounded(
pending_content,
stream_think[0].get('thought'),
)
elif think_start and (not stream_think or stream_think[0].get('thought') == '') and not think_end:
think_end = True
pending_content += '\n</think>\n'
pending_content = _append_bounded(
pending_content,
'\n</think>\n',
)
if stream_output.get('text') is not None:
pending_content += stream_output.get('text')
pending_content = _append_bounded(
pending_content,
stream_output.get('text'),
)
# 是否是流式最后一个chunk
is_final = False if stream_output.get('finish_reason', False) == 'null' else True
@@ -162,12 +183,14 @@ class DashScopeAPIRunner(runner.RequestRunner):
# 从模型传出的参考资料信息中提取用于替换的字典
if references_dict_list is not None:
for doc in references_dict_list:
for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]:
if doc.get('index_id') is not None:
references_dict[doc.get('index_id')] = doc.get('doc_name')
# 将参考资料替换到文本中
pending_content = self._replace_references(pending_content, references_dict)
if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS:
raise DashscopeAPIError('Dashscope response exceeds the runtime limit')
if idx_chunk % 8 == 0 or is_final:
yield provider_message.MessageChunk(
@@ -178,7 +201,7 @@ class DashScopeAPIRunner(runner.RequestRunner):
# 保存当前会话的session_id用于下次对话的语境
query.session.using_conversation.uuid = stream_output.get('session_id')
else:
for chunk in response:
async for chunk in runner.iterate_sync(response):
if chunk.get('status_code') != 200:
raise DashscopeAPIError(
f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '
@@ -192,15 +215,27 @@ class DashScopeAPIRunner(runner.RequestRunner):
if stream_think and stream_think[0].get('thought'):
if not think_start:
think_start = True
pending_content += f'<think>\n{stream_think[0].get("thought")}'
pending_content = _append_bounded(
pending_content,
f'<think>\n{stream_think[0].get("thought")}',
)
else:
# 继续输出 reasoning_content
pending_content += stream_think[0].get('thought')
pending_content = _append_bounded(
pending_content,
stream_think[0].get('thought'),
)
elif think_start and (not stream_think or stream_think[0].get('thought') == '') and not think_end:
think_end = True
pending_content += '\n</think>\n'
pending_content = _append_bounded(
pending_content,
'\n</think>\n',
)
if stream_output.get('text') is not None:
pending_content += stream_output.get('text')
pending_content = _append_bounded(
pending_content,
stream_output.get('text'),
)
# 保存当前会话的session_id用于下次对话的语境
query.session.using_conversation.uuid = stream_output.get('session_id')
@@ -210,12 +245,14 @@ class DashScopeAPIRunner(runner.RequestRunner):
# 从模型传出的参考资料信息中提取用于替换的字典
if references_dict_list is not None:
for doc in references_dict_list:
for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]:
if doc.get('index_id') is not None:
references_dict[doc.get('index_id')] = doc.get('doc_name')
# 将参考资料替换到文本中
pending_content = self._replace_references(pending_content, references_dict)
pending_content = self._replace_references(pending_content, references_dict)
if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS:
raise DashscopeAPIError('Dashscope response exceeds the runtime limit')
yield provider_message.Message(
role='assistant',
@@ -240,18 +277,16 @@ class DashScopeAPIRunner(runner.RequestRunner):
biz_params.update(query.variables)
# 发送对话请求
response = dashscope.Application.call(
api_key=self.api_key, # 智能体应用的API Key
app_id=self.app_id, # 智能体应用的ID
prompt=plain_text, # 用户输入的文本信息
stream=True, # 流式输出
incremental_output=True, # 增量输出，使用流式输出需要开启增量输出
session_id=query.session.using_conversation.uuid, # 会话ID用于，多轮对话
biz_params=biz_params, # 工作流应用的自定义输入参数传递
flow_stream_mode='message_format', # 消息模式，输出/结束节点的流式结果
# rag_options={ # 主要用于文件交互，暂不支持
# "session_file_ids": ["FILE_ID1"], # FILE_ID1 替换为实际的临时文件ID,逗号隔开多个
# }
response = await asyncio.to_thread(
dashscope.Application.call,
api_key=self.api_key,
app_id=self.app_id,
prompt=plain_text,
stream=True,
incremental_output=True,
session_id=query.session.using_conversation.uuid,
biz_params=biz_params,
flow_stream_mode='message_format',
)
# 处理API返回的流式输出
@@ -262,7 +297,7 @@ class DashScopeAPIRunner(runner.RequestRunner):
is_stream = False
idx_chunk = 0
if is_stream:
for chunk in response:
async for chunk in runner.iterate_sync(response):
if chunk.get('status_code') != 200:
raise DashscopeAPIError(
f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '
@@ -273,7 +308,10 @@ class DashScopeAPIRunner(runner.RequestRunner):
# 获取流式传输的output
stream_output = chunk.get('output', {})
if stream_output.get('workflow_message') is not None:
pending_content += stream_output.get('workflow_message').get('message').get('content')
pending_content = _append_bounded(
pending_content,
stream_output.get('workflow_message').get('message').get('content'),
)
# if stream_output.get('text') is not None:
# pending_content += stream_output.get('text')
@@ -284,12 +322,14 @@ class DashScopeAPIRunner(runner.RequestRunner):
# 从模型传出的参考资料信息中提取用于替换的字典
if references_dict_list is not None:
for doc in references_dict_list:
for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]:
if doc.get('index_id') is not None:
references_dict[doc.get('index_id')] = doc.get('doc_name')
# 将参考资料替换到文本中
pending_content = self._replace_references(pending_content, references_dict)
if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS:
raise DashscopeAPIError('Dashscope response exceeds the runtime limit')
if idx_chunk % 8 == 0 or is_final:
yield provider_message.MessageChunk(
role='assistant',
@@ -301,7 +341,7 @@ class DashScopeAPIRunner(runner.RequestRunner):
query.session.using_conversation.uuid = stream_output.get('session_id')
else:
for chunk in response:
async for chunk in runner.iterate_sync(response):
if chunk.get('status_code') != 200:
raise DashscopeAPIError(
f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '
@@ -312,7 +352,10 @@ class DashScopeAPIRunner(runner.RequestRunner):
# 获取流式传输的output
stream_output = chunk.get('output', {})
if stream_output.get('text') is not None:
pending_content += stream_output.get('text')
pending_content = _append_bounded(
pending_content,
stream_output.get('text'),
)
is_final = False if stream_output.get('finish_reason', False) == 'null' else True
@@ -324,12 +367,14 @@ class DashScopeAPIRunner(runner.RequestRunner):
# 从模型传出的参考资料信息中提取用于替换的字典
if references_dict_list is not None:
for doc in references_dict_list:
for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]:
if doc.get('index_id') is not None:
references_dict[doc.get('index_id')] = doc.get('doc_name')
# 将参考资料替换到文本中
pending_content = self._replace_references(pending_content, references_dict)
pending_content = self._replace_references(pending_content, references_dict)
if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS:
raise DashscopeAPIError('Dashscope response exceeds the runtime limit')
yield provider_message.Message(
role='assistant',
src/langbot/pkg/provider/runners/difysvapi.py
+166 -27
View File
@@ -1,10 +1,11 @@
from __future__ import annotations
import asyncio
import heapq
import typing
import json
import time
import uuid
import base64
import mimetypes
import os
import re
@@ -16,11 +17,9 @@ from langbot.pkg.provider import runner
from langbot.pkg.core import app
import langbot_plugin.api.entities.builtin.provider.message as provider_message
import langbot_plugin.api.entities.builtin.platform.message as platform_message
from langbot.pkg.utils import image
from langbot.pkg.utils import httpclient, image
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from langbot.libs.dify_service_api.v1 import client, errors
import httpx
# Module-level store for paused-workflow form state. The key includes the full
# execution scope before the bot, pipeline, adapter, and launcher dimensions;
@@ -29,8 +28,29 @@ import httpx
# placement generations.
PendingFormKey = tuple[str, str, int, str, str, str, str, str]
_PENDING_FORMS: dict[PendingFormKey, 'OrderedDict[str, dict[str, typing.Any]]'] = {}
_PENDING_FORM_EXPIRY_HEAP: list[tuple[float, int, PendingFormKey, str]] = []
_PENDING_FORM_ACTIVE_COUNT = 0
_PENDING_FORM_REVISION = 0
_PENDING_FORM_DEFAULT_TTL = 30 * 60 # 30 minutes safety cap
_PENDING_FORM_MAX_SESSIONS = 4096
_PENDING_FORM_MAX_PER_SESSION = 16
_PENDING_FORM_HEAP_COMPACT_FLOOR = 64
_PENDING_FORM_HEAP_MAX_MULTIPLIER = 4
_PENDING_FORM_REVISION_KEY = '_langbot_cache_revision'
_STREAM_FORM_PLACEHOLDER = '\u200b'
_MAX_DIFY_UPLOAD_BYTES = 10 * 1024 * 1024
def _read_local_file_limited(path: str) -> bytes:
"""Read a local platform attachment without allowing an oversized allocation."""
if os.path.getsize(path) > _MAX_DIFY_UPLOAD_BYTES:
raise ValueError('Dify upload file exceeds the size limit')
with open(path, 'rb') as file:
content = file.read(_MAX_DIFY_UPLOAD_BYTES + 1)
if len(content) > _MAX_DIFY_UPLOAD_BYTES:
raise ValueError('Dify upload file exceeds the size limit')
return content
def _merge_stream_text(accumulated: str, incoming: typing.Any) -> str:
@@ -65,19 +85,100 @@ def _session_key_from_query(query: pipeline_query.Query) -> PendingFormKey:
)
def _synchronize_pending_form_cache_if_externally_cleared() -> None:
"""Keep test/debug direct cache clears from retaining stale heap entries."""
global _PENDING_FORM_ACTIVE_COUNT
if _PENDING_FORMS:
return
_PENDING_FORM_EXPIRY_HEAP.clear()
_PENDING_FORM_ACTIVE_COUNT = 0
def _pending_form_entry_is_current(
expires_at: float,
revision: int,
session_key: PendingFormKey,
form_token: str,
) -> bool:
forms = _PENDING_FORMS.get(session_key)
if forms is None:
return False
stored = forms.get(form_token)
if stored is None:
return False
return stored.get(_PENDING_FORM_REVISION_KEY) == revision and stored.get('_expires_at') == expires_at
def _peek_valid_pending_form_expiry(
*,
pop: bool = False,
) -> tuple[float, int, PendingFormKey, str] | None:
while _PENDING_FORM_EXPIRY_HEAP:
entry = _PENDING_FORM_EXPIRY_HEAP[0]
if _pending_form_entry_is_current(*entry):
if pop:
heapq.heappop(_PENDING_FORM_EXPIRY_HEAP)
return entry
heapq.heappop(_PENDING_FORM_EXPIRY_HEAP)
return None
def _drop_pending_form(session_key: PendingFormKey, form_token: str) -> None:
global _PENDING_FORM_ACTIVE_COUNT
forms = _PENDING_FORMS.get(session_key)
if forms is None or forms.pop(form_token, None) is None:
return
_PENDING_FORM_ACTIVE_COUNT = max(_PENDING_FORM_ACTIVE_COUNT - 1, 0)
if not forms:
_PENDING_FORMS.pop(session_key, None)
def _drop_pending_form_session(session_key: PendingFormKey) -> None:
global _PENDING_FORM_ACTIVE_COUNT
forms = _PENDING_FORMS.pop(session_key, None)
if forms is not None:
_PENDING_FORM_ACTIVE_COUNT = max(
_PENDING_FORM_ACTIVE_COUNT - len(forms),
0,
)
def _compact_pending_form_expiry_heap_if_needed() -> None:
max_heap_entries = max(
_PENDING_FORM_HEAP_COMPACT_FLOOR,
_PENDING_FORM_ACTIVE_COUNT * _PENDING_FORM_HEAP_MAX_MULTIPLIER,
)
if len(_PENDING_FORM_EXPIRY_HEAP) <= max_heap_entries:
return
_PENDING_FORM_EXPIRY_HEAP[:] = [
(
float(stored['_expires_at']),
int(stored[_PENDING_FORM_REVISION_KEY]),
session_key,
form_token,
)
for session_key, forms in _PENDING_FORMS.items()
for form_token, stored in forms.items()
]
heapq.heapify(_PENDING_FORM_EXPIRY_HEAP)
def _prune_pending_forms(now: float | None = None) -> None:
_synchronize_pending_form_cache_if_externally_cleared()
if now is None:
now = time.time()
for session_key in list(_PENDING_FORMS.keys()):
forms = _PENDING_FORMS[session_key]
expired_tokens = [token for token, data in forms.items() if data.get('_expires_at', 0) <= now]
for token in expired_tokens:
forms.pop(token, None)
if not forms:
_PENDING_FORMS.pop(session_key, None)
while True:
entry = _peek_valid_pending_form_expiry()
if entry is None or entry[0] > now:
break
_, _, session_key, form_token = _peek_valid_pending_form_expiry(pop=True)
_drop_pending_form(session_key, form_token)
_compact_pending_form_expiry_heap_if_needed()
def _set_pending_form(session_key: PendingFormKey, form_data: dict[str, typing.Any]) -> None:
global _PENDING_FORM_ACTIVE_COUNT, _PENDING_FORM_REVISION
_prune_pending_forms()
if isinstance(session_key, tuple) and len(session_key) == 8:
form_data['pipeline_uuid'] = session_key[4]
@@ -88,11 +189,31 @@ def _set_pending_form(session_key: PendingFormKey, form_data: dict[str, typing.A
except (TypeError, ValueError):
expiration_ts = 0.0
stored['_expires_at'] = expiration_ts or (time.time() + _PENDING_FORM_DEFAULT_TTL)
_PENDING_FORM_REVISION += 1
stored[_PENDING_FORM_REVISION_KEY] = _PENDING_FORM_REVISION
form_token = str(stored.get('form_token') or '')
forms = _PENDING_FORMS.setdefault(session_key, OrderedDict())
# Re-insert at the end so this becomes the "latest" entry
forms.pop(form_token, None)
if forms.pop(form_token, None) is None:
_PENDING_FORM_ACTIVE_COUNT += 1
forms[form_token] = stored
heapq.heappush(
_PENDING_FORM_EXPIRY_HEAP,
(
stored['_expires_at'],
_PENDING_FORM_REVISION,
session_key,
form_token,
),
)
while len(forms) > _PENDING_FORM_MAX_PER_SESSION:
oldest_token = next(iter(forms))
_drop_pending_form(session_key, oldest_token)
if len(_PENDING_FORMS) > _PENDING_FORM_MAX_SESSIONS:
oldest_entry = _peek_valid_pending_form_expiry()
if oldest_entry is not None:
_drop_pending_form_session(oldest_entry[2])
_compact_pending_form_expiry_heap_if_needed()
def _get_pending_form_by_token(session_key: PendingFormKey, form_token: str) -> dict[str, typing.Any] | None:
@@ -144,11 +265,11 @@ def _clear_pending_form(session_key: PendingFormKey, form_token: str | None = No
if not forms:
return
if form_token is None:
_PENDING_FORMS.pop(session_key, None)
_drop_pending_form_session(session_key)
_compact_pending_form_expiry_heap_if_needed()
return
forms.pop(form_token, None)
if not forms:
_PENDING_FORMS.pop(session_key, None)
_drop_pending_form(session_key, form_token)
_compact_pending_form_expiry_heap_if_needed()
def _format_human_input_text(
@@ -721,6 +842,9 @@ class DifyServiceAPIRunner(runner.RequestRunner):
base_url=self.pipeline_config['ai']['dify-service-api']['base-url'],
)
async def aclose(self) -> None:
await self.dify_client.aclose()
def _process_thinking_content(
self,
content: str,
@@ -796,13 +920,16 @@ class DifyServiceAPIRunner(runner.RequestRunner):
async def download_file(file_url: str) -> tuple[bytes, str]:
"""Download file from url (supports data url)."""
async with httpx.AsyncClient() as client_session:
resp = await client_session.get(file_url)
client_session = httpclient.get_session()
async with client_session.get(file_url, timeout=120) as resp:
resp.raise_for_status()
content_type = (
resp.headers.get('content-type') or mimetypes.guess_type(file_url)[0] or 'application/octet-stream'
)
return resp.content, content_type
return (
await httpclient.read_limited(resp, max_bytes=_MAX_DIFY_UPLOAD_BYTES),
content_type,
)
def _detect_file_type(content_type: str) -> str:
"""Map MIME to dify file type."""
@@ -820,7 +947,10 @@ class DifyServiceAPIRunner(runner.RequestRunner):
plain_text += ce.text
elif ce.type == 'image_base64':
image_b64, image_format = await image.extract_b64_and_format(ce.image_base64)
file_bytes = base64.b64decode(image_b64)
file_bytes = await image.decode_base64_limited(
image_b64,
max_bytes=_MAX_DIFY_UPLOAD_BYTES,
)
image_id = await upload_file_bytes(f'img.{image_format}', file_bytes, f'image/{image_format}')
upload_files.append({'type': 'image', 'id': image_id})
elif ce.type == 'file_url':
@@ -840,7 +970,10 @@ class DifyServiceAPIRunner(runner.RequestRunner):
content_type = 'application/octet-stream'
if ';' in header:
content_type = header.split(';')[0][5:] or content_type
file_bytes = base64.b64decode(b64_data)
file_bytes = await image.decode_base64_limited(
b64_data,
max_bytes=_MAX_DIFY_UPLOAD_BYTES,
)
file_id = await upload_file_bytes(file_name, file_bytes, content_type)
file_type = _detect_file_type(content_type)
upload_files.append({'type': file_type, 'id': file_id})
@@ -865,15 +998,19 @@ class DifyServiceAPIRunner(runner.RequestRunner):
}
async def _download_file_for_form(self, file_url: str) -> tuple[bytes, str, str]:
async with httpx.AsyncClient() as client_session:
resp = await client_session.get(file_url)
client_session = httpclient.get_session()
async with client_session.get(file_url, timeout=120) as resp:
resp.raise_for_status()
content_type = (
resp.headers.get('content-type') or mimetypes.guess_type(file_url)[0] or 'application/octet-stream'
)
parsed = urlparse(file_url)
file_name = os.path.basename(parsed.path) or 'file'
return resp.content, content_type, file_name
return (
await httpclient.read_limited(resp, max_bytes=_MAX_DIFY_UPLOAD_BYTES),
content_type,
file_name,
)
async def _platform_file_to_dify(self, item: typing.Any, user: str) -> dict | None:
try:
@@ -890,13 +1027,15 @@ class DifyServiceAPIRunner(runner.RequestRunner):
content_type = header.split(';', 1)[0][5:] or content_type
return await self._upload_file_bytes_for_user(
file_name,
base64.b64decode(b64_data),
await image.decode_base64_limited(
b64_data,
max_bytes=_MAX_DIFY_UPLOAD_BYTES,
),
content_type,
user,
)
if item.path:
with open(item.path, 'rb') as f:
file_bytes = f.read()
file_bytes = await asyncio.to_thread(_read_local_file_limited, str(item.path))
content_type = mimetypes.guess_type(str(item.path))[0] or 'application/octet-stream'
file_name = item.name or os.path.basename(str(item.path)) or 'file'
return await self._upload_file_bytes_for_user(file_name, file_bytes, content_type, user)
src/langbot/pkg/provider/runners/langflowapi.py
+49 -6
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import codecs
import typing
import json
import httpx
@@ -11,6 +12,44 @@ from ...core import app
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
_MAX_LANGFLOW_LINE_CHARS = 1024 * 1024
_MAX_LANGFLOW_TOTAL_BYTES = 16 * 1024 * 1024
_MAX_LANGFLOW_RESPONSE_BYTES = 1024 * 1024
async def _iter_limited_lines(
response: httpx.Response,
) -> typing.AsyncGenerator[str, None]:
decoder = codecs.getincrementaldecoder('utf-8')('replace')
buffer = ''
total_bytes = 0
async for chunk in response.aiter_bytes(chunk_size=8192):
total_bytes += len(chunk)
if total_bytes > _MAX_LANGFLOW_TOTAL_BYTES:
raise ValueError('Langflow stream exceeds the runtime limit')
buffer += decoder.decode(chunk)
while '\n' in buffer:
line, buffer = buffer.split('\n', 1)
if len(line) > _MAX_LANGFLOW_LINE_CHARS:
raise ValueError('Langflow event exceeds the runtime limit')
yield line.rstrip('\r')
if len(buffer) > _MAX_LANGFLOW_LINE_CHARS:
raise ValueError('Langflow event exceeds the runtime limit')
buffer += decoder.decode(b'', final=True)
if buffer:
if len(buffer) > _MAX_LANGFLOW_LINE_CHARS:
raise ValueError('Langflow event exceeds the runtime limit')
yield buffer.rstrip('\r')
async def _read_limited_response(response: httpx.Response) -> bytes:
body = bytearray()
async for chunk in response.aiter_bytes(chunk_size=8192):
body.extend(chunk)
if len(body) > _MAX_LANGFLOW_RESPONSE_BYTES:
raise ValueError('Langflow response exceeds the runtime limit')
return bytes(body)
@runner.runner_class('langflow-api')
class LangflowAPIRunner(runner.RequestRunner):
@@ -99,7 +138,7 @@ class LangflowAPIRunner(runner.RequestRunner):
accumulated_content = ''
message_count = 0
async for line in response.aiter_lines():
async for line in _iter_limited_lines(response):
data_str = line
if data_str.startswith('data: '):
@@ -144,11 +183,15 @@ class LangflowAPIRunner(runner.RequestRunner):
yield provider_message.MessageChunk(role='assistant', content=accumulated_content, is_final=True)
else:
# 非流式请求
response = await client.post(url, json=payload, headers=headers, timeout=120.0)
response.raise_for_status()
# 解析响应
response_data = response.json()
async with client.stream(
'POST',
url,
json=payload,
headers=headers,
timeout=120.0,
) as response:
response.raise_for_status()
response_data = json.loads(await _read_limited_response(response))
# 提取消息内容
# 根据Langflow API文档，响应结构可能在outputs[0].outputs[0].outputs.message.message中
src/langbot/pkg/provider/runners/n8nsvapi.py
+15 -2
View File
@@ -12,6 +12,8 @@ from ...core import app
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
_MAX_N8N_RESPONSE_CHARS = 1024 * 1024
class N8nAPIError(Exception):
"""N8n API 请求失败"""
@@ -94,6 +96,8 @@ class N8nServiceAPIRunner(runner.RequestRunner):
else:
chunk_str = str(raw_chunk)
if len(full_text) + len(chunk_str) > _MAX_N8N_RESPONSE_CHARS:
raise N8nAPIError('n8n response exceeds the runtime limit')
full_text += chunk_str
buffer += chunk_str
@@ -112,7 +116,9 @@ class N8nServiceAPIRunner(runner.RequestRunner):
if obj.get('type') == 'item' and 'content' in obj:
chunk_idx += 1
content = obj['content']
content = str(obj['content'])
if len(full_content) + len(content) > _MAX_N8N_RESPONSE_CHARS:
raise N8nAPIError('n8n response exceeds the runtime limit')
full_content += content
elif obj.get('type') == 'end':
is_final = True
@@ -128,6 +134,8 @@ class N8nServiceAPIRunner(runner.RequestRunner):
except json.JSONDecodeError:
# buffer 末尾可能是一个不完整的 JSON，等待更多数据
break
except N8nAPIError:
raise
except Exception as e:
# 记录解析失败并继续接收后续 chunk
try:
@@ -255,7 +263,12 @@ class N8nServiceAPIRunner(runner.RequestRunner):
self.webhook_url, json=payload, headers=headers, auth=auth, timeout=self.timeout
) as response:
if response.status != 200:
error_text = await response.text()
error_text = (
await httpclient.read_limited(
response,
max_bytes=_MAX_N8N_RESPONSE_CHARS,
)
).decode('utf-8', errors='replace')
self.ap.logger.error(f'n8n webhook call failed: {response.status}, {error_text}')
raise Exception(f'n8n webhook call failed: {response.status}, {error_text}')
src/langbot/pkg/provider/runners/tboxapi.py
+75 -27
View File
@@ -1,8 +1,9 @@
from __future__ import annotations
import asyncio
import typing
import json
import base64
import logging
import tempfile
import os
@@ -11,10 +12,13 @@ from tboxsdk.model.file import File, FileType
from .. import runner
from ...core import app
from ...utils import image
from ...utils import bounded_executor, image
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
_MAX_TBOX_RESPONSE_CHARS = 1024 * 1024
_MAX_TBOX_MEDIA_BYTES = 10 * 1024 * 1024
class TboxAPIError(Exception):
"""TBox API 请求失败"""
@@ -24,6 +28,19 @@ class TboxAPIError(Exception):
super().__init__(self.message)
def _append_bounded(current: str, addition: typing.Any) -> str:
addition = str(addition or '')
if len(current) + len(addition) > _MAX_TBOX_RESPONSE_CHARS:
raise TboxAPIError('Tbox response exceeds the runtime limit')
return current + addition
def _write_temp_media(file_bytes: bytes, suffix: str) -> str:
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp_file:
tmp_file.write(file_bytes)
return tmp_file.name
@runner.runner_class('tbox-app-api')
class TboxAPIRunner(runner.RequestRunner):
"蚂蚁百宝箱API对话请求器"
@@ -42,6 +59,7 @@ class TboxAPIRunner(runner.RequestRunner):
self.api_key = self.pipeline_config['ai']['tbox-app-api']['api-key']
# 初始化Tbox client
logging.getLogger('tbox.client').setLevel(logging.WARNING)
self.tbox_client = TboxClient(authorization=self.api_key)
async def _preprocess_user_message(self, query: pipeline_query.Query) -> tuple[str, list[str]]:
@@ -59,19 +77,29 @@ class TboxAPIRunner(runner.RequestRunner):
plain_text += ce.text
elif ce.type == 'image_base64':
image_b64, image_format = await image.extract_b64_and_format(ce.image_base64)
# 创建临时文件
file_bytes = base64.b64decode(image_b64)
file_bytes = await image.decode_base64_limited(
image_b64,
max_bytes=_MAX_TBOX_MEDIA_BYTES,
)
tmp_file_path: str | None = None
try:
with tempfile.NamedTemporaryFile(suffix=f'.{image_format}', delete=False) as tmp_file:
tmp_file.write(file_bytes)
tmp_file_path = tmp_file.name
file_upload_resp = self.tbox_client.upload_file(tmp_file_path)
tmp_file_path = await asyncio.to_thread(
_write_temp_media,
file_bytes,
f'.{image_format}',
)
file_upload_resp = await asyncio.to_thread(
self.tbox_client.upload_file,
tmp_file_path,
)
image_id = file_upload_resp.get('data', '')
image_ids.append(image_id)
finally:
# 清理临时文件
if os.path.exists(tmp_file_path):
os.unlink(tmp_file_path)
if tmp_file_path and os.path.exists(tmp_file_path):
await bounded_executor.run_blocking_cleanup(
os.unlink,
tmp_file_path,
)
elif isinstance(query.user_message.content, str):
plain_text = query.user_message.content
@@ -98,18 +126,23 @@ class TboxAPIRunner(runner.RequestRunner):
files = [File(file_id=image_id, type=FileType.IMAGE) for image_id in image_ids]
# 发送对话请求
response = self.tbox_client.chat(
app_id=self.app_id, # Tbox中智能体应用的ID
user_id=query.bot_uuid, # 用户ID
query=plain_text, # 用户输入的文本信息
stream=is_stream, # 是否流式输出
conversation_id=conversation_id, # 会话ID，为None时Tbox会自动创建一个新会话
files=files, # 图片内容
response = await asyncio.to_thread(
self.tbox_client.chat,
app_id=self.app_id,
user_id=query.bot_uuid,
query=plain_text,
stream=is_stream,
conversation_id=conversation_id,
files=files,
)
if is_stream:
# 解析Tbox流式输出内容，并发送给上游
for chunk in self._process_stream_message(response, query, remove_think):
async for chunk in self._process_stream_message(
response,
query,
remove_think,
):
yield chunk
else:
message = self._process_non_stream_message(response, query, remove_think)
@@ -127,13 +160,16 @@ class TboxAPIRunner(runner.RequestRunner):
thinking_content = payload.get('reasoningContent', [])
result = ''
if thinking_content and not remove_think:
result += f'<think>\n{thinking_content[0].get("text", "")}\n</think>\n'
result = _append_bounded(
result,
f'<think>\n{thinking_content[0].get("text", "")}\n</think>\n',
)
content = payload.get('result', [])
if content:
result += content[0].get('chunk', '')
result = _append_bounded(result, content[0].get('chunk', ''))
return result
def _process_stream_message(
async def _process_stream_message(
self, response: typing.Generator[dict], query: pipeline_query.Query, remove_think: bool
):
idx_msg = 0
@@ -141,7 +177,7 @@ class TboxAPIRunner(runner.RequestRunner):
conversation_id = None
think_start = False
think_end = False
for chunk in response:
async for chunk in runner.iterate_sync(response):
if chunk.get('type', '') == 'chunk':
"""
Tbox返回的消息内容chunk结构
@@ -149,7 +185,10 @@ class TboxAPIRunner(runner.RequestRunner):
"""
# 如果包含思考过程，拼接</think>
if think_start and not think_end:
pending_content += '\n</think>\n'
pending_content = _append_bounded(
pending_content,
'\n</think>\n',
)
think_end = True
payload = chunk.get('payload', {})
@@ -158,7 +197,10 @@ class TboxAPIRunner(runner.RequestRunner):
query.session.using_conversation.uuid = conversation_id
if payload.get('text'):
idx_msg += 1
pending_content += payload.get('text')
pending_content = _append_bounded(
pending_content,
payload.get('text'),
)
elif chunk.get('type', '') == 'thinking' and not remove_think:
"""
Tbox返回的思考过程chunk结构
@@ -170,9 +212,15 @@ class TboxAPIRunner(runner.RequestRunner):
content = payload.get('ext_data', {}).get('text')
if not think_start:
think_start = True
pending_content += f'<think>\n{content}'
pending_content = _append_bounded(
pending_content,
f'<think>\n{content}',
)
else:
pending_content += content
pending_content = _append_bounded(
pending_content,
content,
)
elif chunk.get('type', '') == 'error':
raise TboxAPIError(
f'Tbox API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '
src/langbot/pkg/provider/runners/weknoraapi.py
+17 -8
View File
@@ -10,6 +10,15 @@ import langbot_plugin.api.entities.builtin.provider.message as provider_message
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from langbot.libs.weknora_api import client, errors
_MAX_WEKNORA_GENERATED_CHARS = 1024 * 1024
def _append_bounded(current: str, addition: typing.Any) -> str:
addition = str(addition or '')
if len(current) + len(addition) > _MAX_WEKNORA_GENERATED_CHARS:
raise errors.WeKnoraAPIError('WeKnora response exceeds the runtime limit')
return current + addition
@runner.runner_class('weknora-api')
class WeKnoraAPIRunner(runner.RequestRunner):
@@ -94,7 +103,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
web_search_enabled=web_search_enabled,
timeout=timeout,
):
self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk))
self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk)[:1000])
response_type = chunk.get('response_type', '')
content = chunk.get('content', '')
@@ -120,7 +129,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
elif response_type == 'answer':
if content:
full_answer += content
full_answer = _append_bounded(full_answer, content)
elif response_type == 'error':
raise errors.WeKnoraAPIError(f'WeKnora 服务错误: {content}')
@@ -158,14 +167,14 @@ class WeKnoraAPIRunner(runner.RequestRunner):
knowledge_base_ids=knowledge_base_ids,
timeout=timeout,
):
self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk))
self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk)[:1000])
response_type = chunk.get('response_type', '')
content = chunk.get('content', '')
if response_type == 'answer':
if content:
full_answer += content
full_answer = _append_bounded(full_answer, content)
elif response_type == 'error':
raise errors.WeKnoraAPIError(f'WeKnora 服务错误: {content}')
@@ -207,7 +216,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
web_search_enabled=web_search_enabled,
timeout=timeout,
):
self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk))
self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk)[:1000])
response_type = chunk.get('response_type', '')
content = chunk.get('content', '')
@@ -235,7 +244,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
elif response_type == 'answer':
message_idx += 1
if content:
pending_answer += content
pending_answer = _append_bounded(pending_answer, content)
if done:
is_final = True
@@ -288,7 +297,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
knowledge_base_ids=knowledge_base_ids,
timeout=timeout,
):
self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk))
self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk)[:1000])
response_type = chunk.get('response_type', '')
content = chunk.get('content', '')
@@ -297,7 +306,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
if response_type == 'answer':
message_idx += 1
if content:
pending_answer += content
pending_answer = _append_bounded(pending_answer, content)
if done:
is_final = True
src/langbot/pkg/provider/session/sessionmgr.py
+280 -7
View File
@@ -2,6 +2,8 @@ from __future__ import annotations
import asyncio
import dataclasses
import heapq
import time
from langbot_plugin.api.entities.builtin.provider import message as provider_message, prompt as provider_prompt
import langbot_plugin.api.entities.builtin.provider.session as provider_session
@@ -24,6 +26,10 @@ SessionKey = tuple[
str,
int | str,
]
SessionExpiryEntry = tuple[float, int, SessionKey]
_SESSION_EXPIRY_HEAP_MIN_LIMIT = 64
_SESSION_EXPIRY_HEAP_ACTIVE_MULTIPLIER = 4
def _query_session_key(query: pipeline_query.Query) -> tuple[SessionKey, ExecutionContext]:
@@ -49,11 +55,215 @@ class SessionManager:
ap: app.Application
session_list: list[provider_session.Session]
def __init__(self, ap: app.Application):
self.ap = ap
self.session_list = []
self._legacy_sessions: list[provider_session.Session] = []
self._session_index: dict[SessionKey, provider_session.Session] = {}
self._session_keys_by_workspace: dict[str, set[SessionKey]] = {}
self._session_expiry_heap: list[SessionExpiryEntry] = []
self._next_access_revision = 0
@property
def session_list(self) -> list[provider_session.Session]:
"""Compatibility view for API services that enumerate sessions."""
return [
*self._legacy_sessions,
*self._session_index.values(),
]
@session_list.setter
def session_list(self, sessions: list[provider_session.Session]) -> None:
"""Replace the cache while keeping the O(1) index consistent."""
session_values = list(sessions)
self._legacy_sessions = []
self._session_index = {}
self._session_keys_by_workspace = {}
self._session_expiry_heap = []
self._next_access_revision = 0
now = time.monotonic()
for session in session_values:
key = getattr(session, '_langbot_session_key', None)
if isinstance(key, tuple) and len(key) == 6:
self._session_index[key] = session
self._session_keys_by_workspace.setdefault(key[1], set()).add(key)
last_accessed = getattr(session, '_langbot_last_accessed', None)
if last_accessed is None:
last_accessed = now
self._touch_session(
session,
key,
float(last_accessed),
compact=False,
)
else:
self._legacy_sessions.append(session)
self._compact_session_expiry_heap(force=True)
def _retention_config(self) -> dict:
instance_config = getattr(getattr(self.ap, 'instance_config', None), 'data', {})
if not isinstance(instance_config, dict):
return {}
config = instance_config.get('system', {}).get('session_retention', {})
return config if isinstance(config, dict) else {}
def _positive_config_int(self, name: str, default: int) -> int:
try:
value = int(self._retention_config().get(name, default))
except (TypeError, ValueError):
value = default
return max(value, 1)
@staticmethod
def _session_is_idle(session: provider_session.Session) -> bool:
semaphore = getattr(session, '_semaphore', None)
concurrency = getattr(session, '_langbot_session_concurrency', None)
if semaphore is None or not isinstance(concurrency, int):
return True
return getattr(semaphore, '_value', -1) == concurrency
def _remove_session(self, session: provider_session.Session) -> None:
key = getattr(session, '_langbot_session_key', None)
if isinstance(key, tuple) and len(key) == 6 and self._session_index.get(key) is session:
self._session_index.pop(key, None)
workspace_keys = self._session_keys_by_workspace.get(key[1])
if workspace_keys is not None:
workspace_keys.discard(key)
if not workspace_keys:
self._session_keys_by_workspace.pop(key[1], None)
else:
try:
self._legacy_sessions.remove(session)
except ValueError:
pass
def _touch_session(
self,
session: provider_session.Session,
key: SessionKey,
now: float,
*,
compact: bool = True,
) -> None:
self._next_access_revision += 1
revision = self._next_access_revision
object.__setattr__(session, '_langbot_last_accessed', now)
object.__setattr__(session, '_langbot_access_revision', revision)
heapq.heappush(
self._session_expiry_heap,
(now, revision, key),
)
if compact:
self._compact_session_expiry_heap()
def _compact_session_expiry_heap(self, *, force: bool = False) -> None:
limit = max(
len(self._session_index) * _SESSION_EXPIRY_HEAP_ACTIVE_MULTIPLIER,
_SESSION_EXPIRY_HEAP_MIN_LIMIT,
)
if not force and len(self._session_expiry_heap) <= limit:
return
self._session_expiry_heap = [
(
float(getattr(session, '_langbot_last_accessed', 0.0)),
int(getattr(session, '_langbot_access_revision', 0)),
key,
)
for key, session in self._session_index.items()
]
heapq.heapify(self._session_expiry_heap)
def _pop_current_expiry_entry(
self,
) -> tuple[float, int, SessionKey, provider_session.Session] | None:
while self._session_expiry_heap:
last_accessed, revision, key = heapq.heappop(self._session_expiry_heap)
session = self._session_index.get(key)
if session is None:
continue
if getattr(session, '_langbot_access_revision', None) != revision:
continue
return last_accessed, revision, key, session
return None
def _prune_expired_sessions(self, now: float) -> None:
idle_ttl = self._positive_config_int('idle_ttl_seconds', 86400)
cutoff = now - idle_ttl
while self._session_expiry_heap:
last_accessed, _, _ = self._session_expiry_heap[0]
if last_accessed > cutoff:
break
current = self._pop_current_expiry_entry()
if current is None:
break
last_accessed, revision, key, session = current
if last_accessed > cutoff:
heapq.heappush(
self._session_expiry_heap,
(last_accessed, revision, key),
)
break
if self._session_is_idle(session):
self._remove_session(session)
continue
# The session became active without another cache lookup. Give it
# a fresh TTL instead of repeatedly examining the same expired
# entry or losing its future expiry record.
self._touch_session(session, key, now)
def _prune_workspace_capacity(
self,
workspace_uuid: str,
max_entries_per_workspace: int,
) -> None:
workspace_keys = self._session_keys_by_workspace.get(workspace_uuid, set())
overflow = len(workspace_keys) - max_entries_per_workspace + 1
if overflow <= 0:
return
idle_workspace_sessions = sorted(
(
session
for key in tuple(workspace_keys)
if (session := self._session_index.get(key)) is not None and self._session_is_idle(session)
),
key=lambda session: float(getattr(session, '_langbot_last_accessed', 0.0)),
)
for session in idle_workspace_sessions[:overflow]:
self._remove_session(session)
def _evict_oldest_idle_session(self, now: float) -> bool:
# At most one current entry per active session is examined. Stale heap
# revisions do not count and are discarded in O(log N).
current_probes = 0
max_probes = len(self._session_index)
while current_probes < max_probes:
current = self._pop_current_expiry_entry()
if current is None:
return False
_, _, key, session = current
current_probes += 1
if self._session_is_idle(session):
self._remove_session(session)
return True
self._touch_session(session, key, now)
return False
def _prune_sessions(self, now: float, workspace_uuid: str) -> None:
self._prune_expired_sessions(now)
max_entries_per_workspace = self._positive_config_int('max_entries_per_workspace', 200)
self._prune_workspace_capacity(
workspace_uuid,
max_entries_per_workspace,
)
max_entries = self._positive_config_int('max_entries', 2000)
overflow = len(self._session_index) - max_entries + 1
if overflow <= 0:
return
for _ in range(overflow):
if not self._evict_oldest_idle_session(now):
break
async def initialize(self):
pass
@@ -61,9 +271,25 @@ class SessionManager:
async def get_session(self, query: pipeline_query.Query) -> provider_session.Session:
"""获取会话"""
session_key, execution_context = _query_session_key(query)
for session in self.session_list:
if getattr(session, '_langbot_session_key', None) == session_key:
return session
now = time.monotonic()
session = self._session_index.get(session_key)
if session is not None:
self._touch_session(session, session_key, now)
return session
self._prune_sessions(now, execution_context.workspace_uuid)
max_entries_per_workspace = self._positive_config_int('max_entries_per_workspace', 200)
workspace_entries = len(
self._session_keys_by_workspace.get(
execution_context.workspace_uuid,
(),
)
)
if workspace_entries >= max_entries_per_workspace:
raise RuntimeError(f'Workspace session cache capacity reached ({max_entries_per_workspace})')
max_entries = self._positive_config_int('max_entries', 2000)
if len(self._session_index) >= max_entries:
raise RuntimeError(f'Session cache capacity reached ({max_entries})')
session_concurrency = self.ap.instance_config.data['concurrency']['session']
@@ -93,8 +319,14 @@ class SessionManager:
object.__setattr__(session, 'bot_uuid', query.bot_uuid)
object.__setattr__(session, '_execution_context', session_context)
object.__setattr__(session, '_langbot_session_key', session_key)
object.__setattr__(session, '_langbot_session_concurrency', session_concurrency)
session._semaphore = asyncio.Semaphore(session_concurrency)
self.session_list.append(session)
self._session_index[session_key] = session
self._session_keys_by_workspace.setdefault(
execution_context.workspace_uuid,
set(),
).add(session_key)
self._touch_session(session, session_key, now)
return session
async def get_conversation(
@@ -144,6 +376,47 @@ class SessionManager:
bot_uuid=bot_uuid,
)
session.conversations.append(conversation)
max_conversations = self._positive_config_int('max_conversations_per_session', 20)
if len(session.conversations) > max_conversations:
del session.conversations[:-max_conversations]
session.using_conversation = conversation
return session.using_conversation
def trim_conversation_messages(
self,
conversation: provider_session.Conversation,
*,
max_rounds: int,
) -> None:
"""Bound retained process-local history after a completed turn."""
try:
max_rounds = int(max_rounds)
except (TypeError, ValueError):
max_rounds = 10
max_rounds = max(max_rounds, 1)
max_messages = self._positive_config_int('max_messages_per_conversation', 100)
kept_reversed = []
user_rounds = 0
for message in reversed(conversation.messages):
if user_rounds >= max_rounds:
break
kept_reversed.append(message)
if getattr(message, 'role', None) == 'user':
user_rounds += 1
retained = list(reversed(kept_reversed))[-max_messages:]
# Binary payloads are needed for the current model call, but retaining
# them in process-local history makes a few image/file turns consume
# hundreds of MB. Historical URL and text references remain intact.
for message in retained:
content = getattr(message, 'content', None)
if not isinstance(content, list):
continue
for element in content:
if getattr(element, 'image_base64', None) is not None:
element.image_base64 = None
if getattr(element, 'file_base64', None) is not None:
element.file_base64 = None
conversation.messages = retained
src/langbot/pkg/provider/tools/loaders/mcp.py
+348 -72
View File
@@ -1,12 +1,13 @@
from __future__ import annotations
import base64
import enum
import json
import math
import re
import time
import typing
import ipaddress
from urllib.parse import urlparse
from contextlib import AsyncExitStack, asynccontextmanager
from datetime import timedelta
import traceback
@@ -50,6 +51,7 @@ MCP_TOOL_READ_RESOURCE = 'langbot_mcp_read_resource'
MCP_RESOURCE_DISCOVERY_MAX_PAGES = 20
MCP_RESOURCE_CACHE_TTL_SECONDS = 30
MCP_RESOURCE_CACHE_MAX_ENTRIES = 32
MCP_RESOURCE_PREVIEW_MAX_BYTES = 64 * 1024
MCP_RESOURCE_AGENT_READ_MAX_BYTES = 64 * 1024
MCP_RESOURCE_AGENT_READ_MAX_TOKENS = 12000
@@ -134,10 +136,13 @@ def _truncate_text(text: str, max_bytes: int, max_tokens: int | None = None) ->
def _blob_size(blob: str) -> int:
try:
return len(base64.b64decode(blob, validate=False))
except Exception:
# MCP BlobResourceContents is schema-validated base64 without whitespace.
# Compute decoded size in O(1) without allocating a second binary copy.
encoded_chars = len(blob)
if encoded_chars % 4:
return len(blob.encode('utf-8', errors='ignore'))
padding = 2 if blob.endswith('==') else 1 if blob.endswith('=') else 0
return max((encoded_chars // 4) * 3 - padding, 0)
def _resource_to_dict(resource: mcp_types.Resource | mcp_types.ResourceLink) -> dict:
@@ -434,12 +439,24 @@ class RuntimeMCPSession:
await self._box_stdio_runtime.initialize()
async def _init_sse_server(self):
trust_env = self._remote_http_trust_env()
def httpx_client_factory(headers=None, timeout=None, auth=None):
return httpx.AsyncClient(
headers=headers,
timeout=timeout,
auth=auth,
follow_redirects=True,
trust_env=trust_env,
)
sse_transport = await self.exit_stack.enter_async_context(
sse_client(
self.server_config['url'],
headers=self.server_config.get('headers', {}),
timeout=self.server_config.get('timeout', 10),
sse_read_timeout=self.server_config.get('ssereadtimeout', 30),
httpx_client_factory=httpx_client_factory,
)
)
@@ -449,6 +466,18 @@ class RuntimeMCPSession:
await self.session.initialize()
def _remote_http_trust_env(self) -> bool:
configured = self.server_config.get('trust_env')
if isinstance(configured, bool):
return configured
hostname = (urlparse(str(self.server_config.get('url', ''))).hostname or '').lower()
if hostname == 'localhost':
return False
try:
return not ipaddress.ip_address(hostname).is_loopback
except ValueError:
return True
@asynccontextmanager
async def _streamable_http_session(self) -> typing.AsyncIterator[ClientSession]:
"""Enter a fully initialized Streamable HTTP session as one context.
@@ -465,6 +494,7 @@ class RuntimeMCPSession:
headers=self.server_config.get('headers', {}),
timeout=self.server_config.get('timeout', 10),
follow_redirects=True,
trust_env=self._remote_http_trust_env(),
) as http_client:
async with streamable_http_client(
self.server_config['url'],
@@ -1215,6 +1245,9 @@ class RuntimeMCPSession:
cache_key = (uri, max_bytes, max_tokens, include_blob)
now = time.time()
for expired_key, entry in tuple(self._resource_cache.items()):
if now - entry.get('cached_at', 0) > MCP_RESOURCE_CACHE_TTL_SECONDS:
self._resource_cache.pop(expired_key, None)
cached = self._resource_cache.get(cache_key)
if cached and now - cached.get('cached_at', 0) <= MCP_RESOURCE_CACHE_TTL_SECONDS:
envelope = {
@@ -1314,6 +1347,12 @@ class RuntimeMCPSession:
'warnings': warnings,
}
await self._assert_execution_active()
if cache_key not in self._resource_cache and len(self._resource_cache) >= MCP_RESOURCE_CACHE_MAX_ENTRIES:
oldest_key = min(
self._resource_cache,
key=lambda key: self._resource_cache[key].get('cached_at', 0),
)
self._resource_cache.pop(oldest_key, None)
self._resource_cache[cache_key] = {'cached_at': now, 'envelope': envelope}
self._record_resource_read_trace(query, envelope)
return envelope
@@ -1498,17 +1537,249 @@ class MCPLoader(loader.ToolLoader):
在此加载器中管理所有与 MCP Server 的连接。
"""
sessions: dict[tuple[str, str, int, str], RuntimeMCPSession]
_last_listed_functions: list[resource_tool.LLMTool]
_sessions: dict[tuple[str, str, int, str], RuntimeMCPSession]
_hosted_mcp_tasks: list[asyncio.Task]
def __init__(self, ap: app.Application):
super().__init__(ap)
self.sessions = {}
self._last_listed_functions = []
self._hosted_mcp_tasks = []
self._hosted_mcp_tasks_by_scope: dict[
tuple[str, str, int],
set[asyncio.Task],
] = {}
self._host_dispatch_tasks: set[asyncio.Task] = set()
config = getattr(getattr(ap, 'instance_config', None), 'data', {})
mcp_config = config.get('mcp', {}) if isinstance(config, dict) else {}
raw_lifecycle_concurrency = mcp_config.get('lifecycle_concurrency', 16) if isinstance(mcp_config, dict) else 16
if (
isinstance(raw_lifecycle_concurrency, bool)
or not isinstance(raw_lifecycle_concurrency, int)
or raw_lifecycle_concurrency < 1
):
raw_lifecycle_concurrency = 16
self._lifecycle_concurrency = min(
raw_lifecycle_concurrency,
128,
)
self._lifecycle_semaphore = asyncio.Semaphore(self._lifecycle_concurrency)
@property
def sessions(
self,
) -> dict[tuple[str, str, int, str], RuntimeMCPSession]:
return self._sessions
@sessions.setter
def sessions(self, sessions: dict) -> None:
"""Compatibility setter that rebuilds the per-scope session index."""
self._sessions = sessions
self._session_keys_by_scope: dict[
tuple[str, str, int],
set[tuple[str, str, int, str]],
] = {}
self._scope_generations: dict[tuple[str, str], int] = {}
for key in sessions:
if not isinstance(key, tuple) or len(key) != 4:
continue
scope_key = key[:3]
self._session_keys_by_scope.setdefault(scope_key, set()).add(key)
self._scope_generations[scope_key[:2]] = scope_key[2]
def _register_session(
self,
context: TenantContext,
server_name: str,
session: RuntimeMCPSession,
) -> None:
scope_key = self._scope_key(context)
workspace_scope = scope_key[:2]
previous_generation = self._scope_generations.get(workspace_scope)
if previous_generation is not None and previous_generation != scope_key[2]:
raise WorkspaceInvariantError('MCP session registration crossed a Workspace generation')
key = (*scope_key, server_name)
self._sessions[key] = session
self._session_keys_by_scope.setdefault(scope_key, set()).add(key)
self._scope_generations[workspace_scope] = scope_key[2]
def _pop_session(
self,
context: TenantContext,
server_name: str,
) -> RuntimeMCPSession | None:
scope_key = self._scope_key(context)
key = (*scope_key, server_name)
session = self._sessions.pop(key, None)
keys = self._session_keys_by_scope.get(scope_key)
if keys is not None:
keys.discard(key)
if not keys:
self._session_keys_by_scope.pop(scope_key, None)
self._drop_empty_scope(scope_key)
return session
def _drop_empty_scope(self, scope_key: tuple[str, str, int]) -> None:
if (
scope_key not in self._session_keys_by_scope
and scope_key not in self._hosted_mcp_tasks_by_scope
and self._scope_generations.get(scope_key[:2]) == scope_key[2]
):
self._scope_generations.pop(scope_key[:2], None)
def track_hosted_task(
self,
task: asyncio.Task,
context: TenantContext,
) -> asyncio.Task:
"""Track a host task without retaining it after completion."""
scope_key = self._scope_key(context)
workspace_scope = scope_key[:2]
previous_generation = self._scope_generations.get(workspace_scope)
if previous_generation is not None and previous_generation != scope_key[2]:
task.cancel()
raise WorkspaceInvariantError('MCP host task crossed a Workspace generation')
self._scope_generations[workspace_scope] = scope_key[2]
self._hosted_mcp_tasks.append(task)
self._hosted_mcp_tasks_by_scope.setdefault(scope_key, set()).add(task)
def discard(completed: asyncio.Task) -> None:
try:
self._hosted_mcp_tasks.remove(completed)
except ValueError:
pass
tasks = self._hosted_mcp_tasks_by_scope.get(scope_key)
if tasks is not None:
tasks.discard(completed)
if not tasks:
self._hosted_mcp_tasks_by_scope.pop(scope_key, None)
self._drop_empty_scope(scope_key)
task.add_done_callback(discard)
return task
def _track_host_dispatch_task(self, task: asyncio.Task) -> None:
"""Track the bounded startup dispatcher without retaining it."""
self._host_dispatch_tasks.add(task)
def discard(completed: asyncio.Task) -> None:
self._host_dispatch_tasks.discard(completed)
if completed.cancelled():
return
exception = completed.exception()
if exception is not None:
self.ap.logger.error(
f'MCP startup dispatcher failed: {exception}',
)
task.add_done_callback(discard)
async def _retire_runtime_scope(
self,
scope_key: tuple[str, str, int],
) -> None:
tasks = tuple(self._hosted_mcp_tasks_by_scope.pop(scope_key, ()))
for task in tasks:
if not task.done():
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
keys = tuple(self._session_keys_by_scope.pop(scope_key, ()))
sessions = [session for key in keys if (session := self._sessions.pop(key, None)) is not None]
await self._shutdown_sessions(sessions)
self._scope_generations.pop(scope_key[:2], None)
async def _observe_execution_context(
self,
context: ExecutionContext,
) -> None:
workspace_scope = (
context.instance_uuid,
context.workspace_uuid,
)
previous_generation = self._scope_generations.get(workspace_scope)
if previous_generation is None:
return
if context.placement_generation < previous_generation:
raise WorkspaceInvariantError('MCP runtime placement generation rolled back')
if context.placement_generation == previous_generation:
return
await self._retire_runtime_scope((*workspace_scope, previous_generation))
async def _reset_runtime_state(self) -> None:
"""Cancel host tasks and close sessions before reload or shutdown."""
dispatch_tasks = tuple(self._host_dispatch_tasks)
self._host_dispatch_tasks.clear()
for task in dispatch_tasks:
if not task.done():
task.cancel()
if dispatch_tasks:
await asyncio.gather(*dispatch_tasks, return_exceptions=True)
tasks = tuple(self._hosted_mcp_tasks)
self._hosted_mcp_tasks.clear()
self._hosted_mcp_tasks_by_scope.clear()
for task in tasks:
if not task.done():
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
sessions = tuple(self._sessions.values())
self.sessions = {}
await self._shutdown_sessions(sessions)
async def _shutdown_sessions(
self,
sessions: typing.Iterable[RuntimeMCPSession],
) -> None:
"""Close MCP sessions in bounded batches to avoid shutdown storms."""
session_list = list(sessions)
for offset in range(0, len(session_list), self._lifecycle_concurrency):
batch = session_list[offset : offset + self._lifecycle_concurrency]
results = await asyncio.gather(
*(session.shutdown() for session in batch),
return_exceptions=True,
)
for session, result in zip(batch, results, strict=True):
if isinstance(result, BaseException):
self.ap.logger.error(f'Error shutting down MCP session {session.server_name}: {result}')
async def _host_server_configs_bounded(
self,
server_configs: typing.Sequence[tuple[ExecutionContext, dict],],
) -> None:
"""Create at most one lifecycle batch of MCP host tasks at a time."""
for offset in range(0, len(server_configs), self._lifecycle_concurrency):
batch = server_configs[offset : offset + self._lifecycle_concurrency]
tasks: list[asyncio.Task] = []
for execution_context, config in batch:
task = create_detached_task(
self.host_mcp_server(execution_context, config),
after_commit_manager=getattr(
self.ap,
'persistence_mgr',
None,
),
workspace_uuid=execution_context.workspace_uuid,
)
tasks.append(task)
try:
self.track_hosted_task(task, execution_context)
except WorkspaceInvariantError as exc:
self.ap.logger.warning(
f'Skipping stale MCP startup task for {execution_context.workspace_uuid}: {exc}'
)
continue
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
async def _assert_execution_active(
self,
@@ -1523,6 +1794,7 @@ class MCPLoader(loader.ToolLoader):
)
if binding.instance_uuid != execution_context.instance_uuid:
raise WorkspaceInvariantError('MCP caller instance does not match the active Workspace binding')
await self._observe_execution_context(execution_context)
return execution_context
async def initialize(self):
@@ -1531,9 +1803,36 @@ class MCPLoader(loader.ToolLoader):
async def load_mcp_servers_from_db(self):
self.ap.logger.info('Loading MCP servers from db...')
self.sessions = {}
await self._reset_runtime_state()
pending_hosts: list[tuple[ExecutionContext, dict]] = []
async def queue_server(binding, server) -> None:
config = self.ap.persistence_mgr.serialize_model(
persistence_mcp.MCPServer,
server,
)
if config.get('mode') == 'stdio' and not stdio_mcp_enabled(self.ap):
self.ap.logger.info(
f'Skipping disabled stdio MCP server {server.uuid}; '
'the persisted configuration is retained but no process is launched'
)
return
try:
if binding is None:
binding = await self.ap.workspace_service.get_execution_binding(server.workspace_uuid)
execution_context = ExecutionContext(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
placement_generation=binding.placement_generation,
)
except Exception as exc:
self.ap.logger.warning(
f'Skipping MCP server {server.uuid}: Workspace execution binding is unavailable: {exc}'
)
return
pending_hosts.append((execution_context, config))
server_configs: list[tuple[typing.Any, typing.Any, dict]] = []
list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
@@ -1548,51 +1847,19 @@ class MCPLoader(loader.ToolLoader):
.order_by(persistence_mcp.MCPServer.uuid)
)
for server in result.all():
server_configs.append(
(
binding,
server,
self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server),
)
)
await queue_server(binding, server)
else:
# Compatibility path for isolated loader tests and older embedders.
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_mcp.MCPServer))
for server in result.all():
server_configs.append(
(
None,
server,
self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server),
)
)
await queue_server(None, server)
for binding, server, config in server_configs:
if config.get('mode') == 'stdio' and not stdio_mcp_enabled(self.ap):
self.ap.logger.info(
f'Skipping disabled stdio MCP server {server.uuid}; '
'the persisted configuration is retained but no process is launched'
)
continue
try:
if binding is None:
binding = await self.ap.workspace_service.get_execution_binding(server.workspace_uuid)
execution_context = ExecutionContext(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
placement_generation=binding.placement_generation,
)
except Exception as exc:
self.ap.logger.warning(
f'Skipping MCP server {server.uuid}: Workspace execution binding is unavailable: {exc}'
)
continue
task = create_detached_task(
self.host_mcp_server(execution_context, config),
if pending_hosts:
dispatch_task = create_detached_task(
self._host_server_configs_bounded(pending_hosts),
after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
)
self._hosted_mcp_tasks.append(task)
self._track_host_dispatch_task(dispatch_task)
@staticmethod
def _scope_key(context: TenantContext) -> tuple[str, str, int]:
@@ -1609,9 +1876,25 @@ class MCPLoader(loader.ToolLoader):
def _sessions_for_context(self, context: TenantContext) -> list[RuntimeMCPSession]:
scope_key = self._scope_key(context)
return [session for key, session in self.sessions.items() if key[:3] == scope_key]
return [
session
for key in self._session_keys_by_scope.get(scope_key, ())
if (session := self._sessions.get(key)) is not None
]
async def host_mcp_server(self, context: TenantContext, server_config: dict):
async def host_mcp_server(
self,
context: TenantContext,
server_config: dict,
) -> None:
async with self._lifecycle_semaphore:
await self._host_mcp_server(context, server_config)
async def _host_mcp_server(
self,
context: TenantContext,
server_config: dict,
) -> None:
requested_context = _execution_context_from_tenant(context)
execution_context = await run_in_workspace_uow(
self.ap,
@@ -1627,7 +1910,17 @@ class MCPLoader(loader.ToolLoader):
try:
session = await self.load_mcp_server(execution_context, server_config)
await self._assert_execution_active(execution_context)
self.sessions[self._session_key(execution_context, server_config['name'])] = session
old_session = self._pop_session(
execution_context,
server_config['name'],
)
if old_session is not None:
await old_session.shutdown()
self._register_session(
execution_context,
server_config['name'],
session,
)
except Exception as e:
self.ap.logger.error(
f'Failed to load MCP server from db: {server_config["name"]}({server_config["uuid"]}): {e}\n{traceback.format_exc()}'
@@ -1876,8 +2169,6 @@ class MCPLoader(loader.ToolLoader):
if include_resource_tools and self._eligible_resource_sessions_for_bound(context, bound_mcp_servers):
all_functions.extend(self._mcp_synthetic_resource_tools())
self._last_listed_functions = all_functions
return all_functions
async def get_tool_catalog(
@@ -2140,7 +2431,9 @@ class MCPLoader(loader.ToolLoader):
self.ap.logger.warning(f'MCP server {server_name} not found in sessions, skipping removal')
return
session = self.sessions.pop(key)
session = self._pop_session(context, server_name)
if session is None:
return
await session.shutdown()
self.ap.logger.info(f'Removed MCP server: {server_name}')
@@ -2180,22 +2473,5 @@ class MCPLoader(loader.ToolLoader):
"""关闭所有工具"""
self.ap.logger.info('Shutting down all MCP sessions...')
hosted_tasks = [task for task in self._hosted_mcp_tasks if not task.done()]
for task in hosted_tasks:
task.cancel()
if hosted_tasks:
await asyncio.gather(*hosted_tasks, return_exceptions=True)
self._hosted_mcp_tasks.clear()
async def shutdown_session(session: RuntimeMCPSession) -> None:
try:
await session.shutdown()
self.ap.logger.debug(f'Shutdown MCP session: {session.server_name}')
except Exception as e:
self.ap.logger.error(
f'Error shutting down MCP session {session.server_name}: {e}\n{traceback.format_exc()}'
)
await asyncio.gather(*(shutdown_session(session) for session in list(self.sessions.values())))
self.sessions.clear()
await self._reset_runtime_state()
self.ap.logger.info('All MCP sessions shutdown complete')
src/langbot/pkg/provider/tools/loaders/mcp_stdio.py
+8 -2
View File
@@ -6,10 +6,12 @@ import os
import shutil
import shlex
import threading
import weakref
from contextlib import suppress, AsyncExitStack, asynccontextmanager
from typing import TYPE_CHECKING, Any
import pydantic
from ....utils import bounded_executor
from mcp import ClientSession
from mcp.client.websocket import websocket_client
from ....box.workspace import (
@@ -27,7 +29,7 @@ if TYPE_CHECKING:
from .mcp import RuntimeMCPSession
_WORKSPACE_COPY_LOCKS: dict[str, threading.Lock] = {}
_WORKSPACE_COPY_LOCKS: weakref.WeakValueDictionary[str, threading.Lock] = weakref.WeakValueDictionary()
_WORKSPACE_COPY_LOCKS_GUARD = threading.Lock()
@@ -536,7 +538,11 @@ class BoxStdioSessionRuntime:
return
try:
process_host_root = os.path.join(self._shared_workspace_host_path(), '.mcp', self.process_id)
await asyncio.to_thread(shutil.rmtree, process_host_root, True)
await bounded_executor.run_blocking_cleanup(
shutil.rmtree,
process_host_root,
True,
)
except Exception as exc:
self.ap.logger.warning(
f'MCP server {self.server_name}: failed to clean staged workspace '

Some files were not shown because too many files have changed in this diff Show More