mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
fix(runtime): bound tenant resource amplification
This commit is contained in:
@@ -8,12 +8,31 @@ from .secrets import SECRET_MASK, mask_secret_value, restore_secret_placeholders
|
||||
from .tenant import TenantContext, require_workspace_uuid, scope_statement
|
||||
|
||||
|
||||
_DEFAULT_MAX_WEBHOOKS_PER_WORKSPACE = 16
|
||||
_HARD_MAX_WEBHOOKS_PER_WORKSPACE = 64
|
||||
|
||||
|
||||
class WebhookService:
|
||||
ap: app.Application
|
||||
|
||||
def __init__(self, ap: app.Application) -> None:
|
||||
self.ap = ap
|
||||
|
||||
def max_per_workspace(self) -> int:
|
||||
"""Return the configured webhook cap within the process hard limit."""
|
||||
|
||||
config = getattr(getattr(self.ap, 'instance_config', None), 'data', {})
|
||||
try:
|
||||
value = int(
|
||||
config.get('webhooks', {}).get(
|
||||
'max_per_workspace',
|
||||
_DEFAULT_MAX_WEBHOOKS_PER_WORKSPACE,
|
||||
)
|
||||
)
|
||||
except (AttributeError, TypeError, ValueError):
|
||||
value = _DEFAULT_MAX_WEBHOOKS_PER_WORKSPACE
|
||||
return min(max(value, 1), _HARD_MAX_WEBHOOKS_PER_WORKSPACE)
|
||||
|
||||
def _serialize_webhook(self, entity, *, include_secret: bool) -> dict:
|
||||
serialized = self.ap.persistence_mgr.serialize_model(webhook.Webhook, entity)
|
||||
if not include_secret:
|
||||
@@ -24,7 +43,11 @@ class WebhookService:
|
||||
async def get_webhooks(self, context: TenantContext, *, include_secret: bool = False) -> list[dict]:
|
||||
"""Get all webhooks"""
|
||||
result = await self.ap.persistence_mgr.execute_async(
|
||||
scope_statement(sqlalchemy.select(webhook.Webhook), webhook.Webhook, context)
|
||||
scope_statement(
|
||||
sqlalchemy.select(webhook.Webhook).order_by(webhook.Webhook.id).limit(_HARD_MAX_WEBHOOKS_PER_WORKSPACE),
|
||||
webhook.Webhook,
|
||||
context,
|
||||
)
|
||||
)
|
||||
|
||||
webhooks = result.all()
|
||||
@@ -40,6 +63,15 @@ class WebhookService:
|
||||
) -> dict:
|
||||
"""Create a new webhook"""
|
||||
workspace_uuid = require_workspace_uuid(context)
|
||||
max_webhooks = self.max_per_workspace()
|
||||
count_result = await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.select(sqlalchemy.func.count())
|
||||
.select_from(webhook.Webhook)
|
||||
.where(webhook.Webhook.workspace_uuid == workspace_uuid)
|
||||
)
|
||||
if (count_result.scalar() or 0) >= max_webhooks:
|
||||
raise ValueError(f'Maximum number of webhooks ({max_webhooks}) reached')
|
||||
|
||||
url = restore_secret_placeholders(url, sensitive=True)
|
||||
webhook_data = {
|
||||
'workspace_uuid': workspace_uuid,
|
||||
@@ -143,6 +175,8 @@ class WebhookService:
|
||||
webhook.Webhook,
|
||||
context,
|
||||
)
|
||||
.order_by(webhook.Webhook.id)
|
||||
.limit(self.max_per_workspace())
|
||||
)
|
||||
|
||||
webhooks = result.all()
|
||||
|
||||
@@ -38,6 +38,8 @@ _INT_ADAPTER = pydantic.TypeAdapter(int)
|
||||
_UTC = _dt.timezone.utc
|
||||
_MAX_RECENT_ERRORS = 50
|
||||
_MIB = 1024 * 1024
|
||||
_DEFAULT_MAX_WORKSPACE_ENTRIES = 100_000
|
||||
_HARD_MAX_WORKSPACE_ENTRIES = 1_000_000
|
||||
|
||||
|
||||
def _create_shared_workspace_probe(root: str, marker_name: str, payload: bytes) -> None:
|
||||
@@ -1214,25 +1216,53 @@ class BoxService:
|
||||
import json as _json
|
||||
|
||||
target_dir = f'{self.OUTBOX_MOUNT_DIR}/{self._attachment_query_key(query)}'
|
||||
max_bytes = self._EXEC_FALLBACK_MAX_BYTES
|
||||
max_file_bytes = self._EXEC_FALLBACK_MAX_BYTES
|
||||
max_files = self._ATTACHMENT_MAX_FILES
|
||||
max_total_bytes = max_file_bytes * max_files
|
||||
max_scan_entries = 1000
|
||||
script = (
|
||||
'import base64, json, os\n'
|
||||
f'target = {target_dir!r}\n'
|
||||
f'max_bytes = {max_bytes}\n'
|
||||
f'max_file_bytes = {max_file_bytes}\n'
|
||||
f'max_files = {max_files}\n'
|
||||
f'max_total_bytes = {max_total_bytes}\n'
|
||||
f'max_scan_entries = {max_scan_entries}\n'
|
||||
'out = []\n'
|
||||
'total_bytes = 0\n'
|
||||
'scanned_entries = 0\n'
|
||||
'stack = [target]\n'
|
||||
'if os.path.isdir(target):\n'
|
||||
' for root, _dirs, names in os.walk(target):\n'
|
||||
' for n in sorted(names):\n'
|
||||
' p = os.path.join(root, n)\n'
|
||||
' while stack and len(out) < max_files and scanned_entries < max_scan_entries:\n'
|
||||
' current = stack.pop()\n'
|
||||
' try:\n'
|
||||
' with os.scandir(current) as iterator:\n'
|
||||
' entries = sorted(iterator, key=lambda item: item.name, reverse=True)\n'
|
||||
' except OSError:\n'
|
||||
' continue\n'
|
||||
' for entry in entries:\n'
|
||||
' scanned_entries += 1\n'
|
||||
' if scanned_entries > max_scan_entries:\n'
|
||||
' break\n'
|
||||
' try:\n'
|
||||
' if os.path.getsize(p) > max_bytes:\n'
|
||||
' if entry.is_dir(follow_symlinks=False):\n'
|
||||
' stack.append(entry.path)\n'
|
||||
' continue\n'
|
||||
' if not entry.is_file(follow_symlinks=False):\n'
|
||||
' continue\n'
|
||||
' size = entry.stat(follow_symlinks=False).st_size\n'
|
||||
' if size > max_file_bytes or total_bytes + size > max_total_bytes:\n'
|
||||
' continue\n'
|
||||
" with open(entry.path, 'rb') as f:\n"
|
||||
' data = f.read(max_file_bytes + 1)\n'
|
||||
' if len(data) > max_file_bytes or total_bytes + len(data) > max_total_bytes:\n'
|
||||
' continue\n'
|
||||
" with open(p, 'rb') as f:\n"
|
||||
' data = f.read()\n'
|
||||
' except OSError:\n'
|
||||
' continue\n'
|
||||
' rel = os.path.relpath(p, target)\n'
|
||||
' rel = os.path.relpath(entry.path, target)\n'
|
||||
" out.append({'name': rel, 'b64': base64.b64encode(data).decode('ascii')})\n"
|
||||
' total_bytes += len(data)\n'
|
||||
' if len(out) >= max_files:\n'
|
||||
' break\n'
|
||||
'print(json.dumps(out))\n'
|
||||
)
|
||||
result = await self.execute_tool(
|
||||
@@ -1888,29 +1918,50 @@ class BoxService:
|
||||
if normalized_timeout > profile.max_timeout_sec:
|
||||
params['timeout_sec'] = profile.max_timeout_sec
|
||||
|
||||
def _get_workspace_size_bytes(self, root: str) -> int:
|
||||
total = 0
|
||||
def _max_workspace_entries(self) -> int:
|
||||
data = getattr(getattr(self.ap, 'instance_config', None), 'data', {})
|
||||
try:
|
||||
configured = int(
|
||||
data.get('box', {}).get('limits', {}).get('max_workspace_entries', _DEFAULT_MAX_WORKSPACE_ENTRIES)
|
||||
)
|
||||
except (AttributeError, TypeError, ValueError):
|
||||
configured = _DEFAULT_MAX_WORKSPACE_ENTRIES
|
||||
return min(max(configured, 1), _HARD_MAX_WORKSPACE_ENTRIES)
|
||||
|
||||
def _walk(path: str):
|
||||
nonlocal total
|
||||
@staticmethod
|
||||
def _get_workspace_usage(
|
||||
root: str,
|
||||
*,
|
||||
stop_after_bytes: int,
|
||||
max_entries: int,
|
||||
) -> tuple[int, int, bool]:
|
||||
"""Scan depth-first without recursion and stop at either hard bound."""
|
||||
|
||||
total = 0
|
||||
entries_seen = 0
|
||||
directories = [root]
|
||||
while directories:
|
||||
path = directories.pop()
|
||||
try:
|
||||
with os.scandir(path) as entries:
|
||||
for entry in entries:
|
||||
entries_seen += 1
|
||||
if entries_seen > max_entries:
|
||||
return total, entries_seen, True
|
||||
try:
|
||||
if entry.is_symlink():
|
||||
total += entry.stat(follow_symlinks=False).st_size
|
||||
continue
|
||||
if entry.is_dir(follow_symlinks=False):
|
||||
_walk(entry.path)
|
||||
continue
|
||||
total += entry.stat(follow_symlinks=False).st_size
|
||||
elif entry.is_dir(follow_symlinks=False):
|
||||
directories.append(entry.path)
|
||||
else:
|
||||
total += entry.stat(follow_symlinks=False).st_size
|
||||
except FileNotFoundError:
|
||||
continue
|
||||
if total > stop_after_bytes:
|
||||
return total, entries_seen, False
|
||||
except FileNotFoundError:
|
||||
return
|
||||
|
||||
_walk(root)
|
||||
return total
|
||||
continue
|
||||
return total, entries_seen, False
|
||||
|
||||
async def _enforce_workspace_quota(self, spec: BoxSpec, *, phase: str) -> None:
|
||||
if spec.host_path is None or spec.workspace_quota_mb <= 0:
|
||||
@@ -1923,15 +1974,26 @@ class BoxService:
|
||||
# Walk the workspace off the event loop — this runs on every
|
||||
# quota-enforced exec, and a large tree would otherwise block the whole
|
||||
# asyncio runtime (all bots/pipelines) for the duration of the scan.
|
||||
used_bytes = await asyncio.to_thread(self._get_workspace_size_bytes, host_path)
|
||||
limit_bytes = spec.workspace_quota_mb * _MIB
|
||||
max_entries = self._max_workspace_entries()
|
||||
used_bytes, entries_seen, entry_limit_exceeded = await asyncio.to_thread(
|
||||
self._get_workspace_usage,
|
||||
host_path,
|
||||
stop_after_bytes=limit_bytes,
|
||||
max_entries=max_entries,
|
||||
)
|
||||
if entry_limit_exceeded:
|
||||
raise BoxValidationError(
|
||||
f'workspace entry limit exceeded {phase}: '
|
||||
f'entries>{max_entries} host_path={host_path} session_id={spec.session_id}'
|
||||
)
|
||||
if used_bytes <= limit_bytes:
|
||||
return
|
||||
|
||||
raise BoxValidationError(
|
||||
f'workspace quota exceeded {phase}: '
|
||||
f'used={used_bytes} bytes limit={limit_bytes} bytes '
|
||||
f'host_path={host_path} session_id={spec.session_id}'
|
||||
f'entries={entries_seen} host_path={host_path} session_id={spec.session_id}'
|
||||
)
|
||||
|
||||
async def _cleanup_exceeded_session(self, context: TenantContext, spec: BoxSpec) -> None:
|
||||
|
||||
@@ -171,17 +171,23 @@ def wrap_python_command_with_env(
|
||||
import sys
|
||||
|
||||
root = "{mount_path}"
|
||||
max_manifest_bytes = 10 * 1024 * 1024
|
||||
digest = hashlib.sha256()
|
||||
manifest_files = []
|
||||
for rel in ("requirements.txt", "pyproject.toml", "setup.py", "setup.cfg"):
|
||||
path = os.path.join(root, rel)
|
||||
if not os.path.isfile(path):
|
||||
continue
|
||||
if os.path.getsize(path) > max_manifest_bytes:
|
||||
raise RuntimeError(
|
||||
f"Python project manifest exceeds {{max_manifest_bytes}} bytes: {{rel}}"
|
||||
)
|
||||
manifest_files.append(rel)
|
||||
with open(path, "rb") as handle:
|
||||
digest.update(rel.encode("utf-8"))
|
||||
digest.update(b"\\0")
|
||||
digest.update(handle.read())
|
||||
while chunk := handle.read(1024 * 1024):
|
||||
digest.update(chunk)
|
||||
digest.update(b"\\0")
|
||||
|
||||
print(
|
||||
|
||||
@@ -69,7 +69,19 @@ _RUNTIME_POLICY_DEFAULTS = {
|
||||
},
|
||||
'auto_cleanup': {'max_batches_per_table_per_run': 4},
|
||||
},
|
||||
'storage': {'cleanup': {'max_files_per_run': 1000}},
|
||||
'storage': {
|
||||
'max_object_read_bytes': 10485760,
|
||||
'cleanup': {'max_files_per_run': 1000},
|
||||
},
|
||||
'webhooks': {
|
||||
'max_per_workspace': 16,
|
||||
'max_inflight_requests': 16,
|
||||
},
|
||||
'box': {
|
||||
'limits': {
|
||||
'max_workspace_entries': 100000,
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -12,10 +12,21 @@ from PIL import Image, ImageDraw, ImageFont
|
||||
import functools
|
||||
|
||||
from .. import strategy as strategy_model
|
||||
from .forward import ForwardComponentStrategy
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
|
||||
|
||||
_MAX_TEXT_TO_IMAGE_CHARS = 100000
|
||||
_MAX_TEXT_TO_IMAGE_LINES = 256
|
||||
_MAX_TEXT_TO_IMAGE_PIXELS = 8_000_000
|
||||
_MAX_RENDERED_IMAGE_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
class _TextToImageCapacityError(ValueError):
|
||||
"""The requested image would exceed a deterministic resource boundary."""
|
||||
|
||||
|
||||
@strategy_model.strategy_class('image')
|
||||
class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
async def initialize(self):
|
||||
@@ -30,6 +41,12 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
)
|
||||
|
||||
async def process(self, message: str, query: pipeline_query.Query) -> list[platform_message.MessageComponent]:
|
||||
if len(message) > _MAX_TEXT_TO_IMAGE_CHARS:
|
||||
self.ap.logger.warning(
|
||||
f'Text-to-image input exceeds {_MAX_TEXT_TO_IMAGE_CHARS} characters; using forward message'
|
||||
)
|
||||
return await ForwardComponentStrategy(self.ap).process(message, query)
|
||||
|
||||
def render() -> str:
|
||||
render_id = f'{int(time.time())}-{uuid.uuid4().hex}'
|
||||
img_path = f'temp/{render_id}.png'
|
||||
@@ -45,7 +62,12 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
outfile=compressed_path,
|
||||
)
|
||||
with open(compressed_path, 'rb') as f:
|
||||
return base64.b64encode(f.read()).decode('utf-8')
|
||||
image_bytes = f.read(_MAX_RENDERED_IMAGE_BYTES + 1)
|
||||
if len(image_bytes) > _MAX_RENDERED_IMAGE_BYTES:
|
||||
raise _TextToImageCapacityError(
|
||||
f'Rendered image exceeds the {_MAX_RENDERED_IMAGE_BYTES}-byte limit'
|
||||
)
|
||||
return base64.b64encode(image_bytes).decode('utf-8')
|
||||
finally:
|
||||
for path in {img_path, compressed_path}:
|
||||
if os.path.exists(path):
|
||||
@@ -53,7 +75,11 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
|
||||
# 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)
|
||||
try:
|
||||
image_base64 = await asyncio.to_thread(render)
|
||||
except _TextToImageCapacityError as exc:
|
||||
self.ap.logger.warning(f'{exc}; using forward message')
|
||||
return await ForwardComponentStrategy(self.ap).process(message, query)
|
||||
|
||||
return [
|
||||
platform_message.Image(
|
||||
@@ -67,38 +93,7 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
:param path:目标字符串
|
||||
:return:<class 'list'>: <class 'list'>: [['1', 16], ['2', 35], ['1', 51]]
|
||||
"""
|
||||
kv = []
|
||||
nums = []
|
||||
beforeDatas = re.findall('[\\d]+', path)
|
||||
for num in beforeDatas:
|
||||
indexV = []
|
||||
times = path.count(num)
|
||||
if times > 1:
|
||||
if num not in nums:
|
||||
indexs = re.finditer(num, path)
|
||||
for index in indexs:
|
||||
iV = []
|
||||
i = index.span()[0]
|
||||
iV.append(num)
|
||||
iV.append(i)
|
||||
kv.append(iV)
|
||||
nums.append(num)
|
||||
else:
|
||||
index = path.find(num)
|
||||
indexV.append(num)
|
||||
indexV.append(index)
|
||||
kv.append(indexV)
|
||||
# 根据数字位置排序
|
||||
indexSort = []
|
||||
resultIndex = []
|
||||
for vi in kv:
|
||||
indexSort.append(vi[1])
|
||||
indexSort.sort()
|
||||
for i in indexSort:
|
||||
for v in kv:
|
||||
if i == v[1]:
|
||||
resultIndex.append(v)
|
||||
return resultIndex
|
||||
return [[match.group(0), match.start()] for match in re.finditer(r'\d+', path)]
|
||||
|
||||
def get_size(self, file):
|
||||
# 获取文件大小:KB
|
||||
@@ -126,9 +121,9 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
return infile, o_size
|
||||
outfile = self.get_outfile(infile, outfile)
|
||||
while o_size > kb:
|
||||
im = Image.open(infile)
|
||||
im.save(outfile, quality=quality)
|
||||
if quality - step < 0:
|
||||
with Image.open(infile) as im:
|
||||
im.save(outfile, quality=quality)
|
||||
if step <= 0 or quality - step < 0:
|
||||
break
|
||||
quality -= step
|
||||
o_size = self.get_size(outfile)
|
||||
@@ -137,12 +132,21 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
def _split_text_lines(self, text_str: str, text_width: int, font) -> list[str]:
|
||||
"""Split text while guaranteeing that every loop iteration advances."""
|
||||
|
||||
if len(text_str) > _MAX_TEXT_TO_IMAGE_CHARS:
|
||||
raise _TextToImageCapacityError(f'Text-to-image input exceeds {_MAX_TEXT_TO_IMAGE_CHARS} characters')
|
||||
|
||||
final_lines: list[str] = []
|
||||
|
||||
def append_line(value: str) -> None:
|
||||
if len(final_lines) >= _MAX_TEXT_TO_IMAGE_LINES:
|
||||
raise _TextToImageCapacityError(f'Text-to-image output exceeds {_MAX_TEXT_TO_IMAGE_LINES} lines')
|
||||
final_lines.append(value)
|
||||
|
||||
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)
|
||||
append_line(line)
|
||||
continue
|
||||
|
||||
rest_text = line
|
||||
@@ -151,16 +155,18 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
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
|
||||
if 0 < point < len(rest_text) and rest_text[point - 1].isdigit() and rest_text[point].isdigit():
|
||||
number_start = point - 1
|
||||
while number_start > 0 and rest_text[number_start - 1].isdigit():
|
||||
number_start -= 1
|
||||
if number_start > 0:
|
||||
point = number_start
|
||||
|
||||
point = max(1, min(point, len(rest_text)))
|
||||
final_lines.append(rest_text[:point])
|
||||
append_line(rest_text[:point])
|
||||
rest_text = rest_text[point:]
|
||||
if rest_text and font.getlength(rest_text) < text_width:
|
||||
final_lines.append(rest_text)
|
||||
append_line(rest_text)
|
||||
break
|
||||
return final_lines
|
||||
|
||||
@@ -171,38 +177,34 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
width=800,
|
||||
query: pipeline_query.Query = None,
|
||||
):
|
||||
width = int(width)
|
||||
if width < 1:
|
||||
raise _TextToImageCapacityError('Text-to-image width must be positive')
|
||||
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)
|
||||
image_height = max(280, len(final_lines) * 35 + 65)
|
||||
if width * image_height > _MAX_TEXT_TO_IMAGE_PIXELS:
|
||||
raise _TextToImageCapacityError(f'Text-to-image canvas exceeds the {_MAX_TEXT_TO_IMAGE_PIXELS}-pixel limit')
|
||||
# 准备画布
|
||||
img = Image.new('RGBA', (width, max(280, len(final_lines) * 35 + 65)), (255, 255, 255, 255))
|
||||
draw = ImageDraw.Draw(img, mode='RGBA')
|
||||
img = Image.new('RGBA', (width, image_height), (255, 255, 255, 255))
|
||||
try:
|
||||
draw = ImageDraw.Draw(img, mode='RGBA')
|
||||
|
||||
self.ap.logger.debug('正在绘制图片...')
|
||||
# 绘制正文
|
||||
line_number = 0
|
||||
offset_x = 20
|
||||
offset_y = 30
|
||||
for final_line in final_lines:
|
||||
draw.text(
|
||||
(offset_x, offset_y + 35 * line_number),
|
||||
final_line,
|
||||
fill=(0, 0, 0),
|
||||
font=font,
|
||||
)
|
||||
# 遍历此行,检查是否有emoji
|
||||
idx_in_line = 0
|
||||
for ch in final_line:
|
||||
# 检查字符占位宽
|
||||
char_code = ord(ch)
|
||||
if char_code >= 127:
|
||||
idx_in_line += 1
|
||||
else:
|
||||
idx_in_line += 0.5
|
||||
self.ap.logger.debug('正在绘制图片...')
|
||||
offset_x = 20
|
||||
offset_y = 30
|
||||
for line_number, final_line in enumerate(final_lines):
|
||||
draw.text(
|
||||
(offset_x, offset_y + 35 * line_number),
|
||||
final_line,
|
||||
fill=(0, 0, 0),
|
||||
font=font,
|
||||
)
|
||||
|
||||
line_number += 1
|
||||
|
||||
self.ap.logger.debug('正在保存图片...')
|
||||
img.save(save_as)
|
||||
self.ap.logger.debug('正在保存图片...')
|
||||
img.save(save_as)
|
||||
finally:
|
||||
img.close()
|
||||
|
||||
return save_as
|
||||
|
||||
@@ -564,7 +564,11 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
raise ValueError('Attachment key does not belong to this WebSocket connection')
|
||||
|
||||
try:
|
||||
file_content = await storage_mgr.storage_provider.load(comp_path)
|
||||
file_content = await storage_mgr.load_scoped_object_key(
|
||||
execution_context,
|
||||
comp_path,
|
||||
expected_owner_type='upload_image',
|
||||
)
|
||||
base64_str = (await asyncio.to_thread(base64.b64encode, file_content)).decode('utf-8')
|
||||
|
||||
lowered = comp_path.lower()
|
||||
|
||||
@@ -15,6 +15,10 @@ if TYPE_CHECKING:
|
||||
import langbot_plugin.api.entities.builtin.platform.events as platform_events
|
||||
|
||||
|
||||
_DEFAULT_MAX_INFLIGHT_WEBHOOK_REQUESTS = 16
|
||||
_HARD_MAX_INFLIGHT_WEBHOOK_REQUESTS = 128
|
||||
|
||||
|
||||
class WebhookPusher:
|
||||
"""Push bot events to configured webhooks"""
|
||||
|
||||
@@ -24,6 +28,56 @@ class WebhookPusher:
|
||||
def __init__(self, ap: app.Application):
|
||||
self.ap = ap
|
||||
self.logger = self.ap.logger
|
||||
self._delivery_lock = asyncio.Lock()
|
||||
self._inflight_requests = 0
|
||||
|
||||
def _max_inflight_requests(self) -> int:
|
||||
config = getattr(getattr(self.ap, 'instance_config', None), 'data', {})
|
||||
try:
|
||||
value = int(
|
||||
config.get('webhooks', {}).get(
|
||||
'max_inflight_requests',
|
||||
_DEFAULT_MAX_INFLIGHT_WEBHOOK_REQUESTS,
|
||||
)
|
||||
)
|
||||
except (AttributeError, TypeError, ValueError):
|
||||
value = _DEFAULT_MAX_INFLIGHT_WEBHOOK_REQUESTS
|
||||
return min(max(value, 1), _HARD_MAX_INFLIGHT_WEBHOOK_REQUESTS)
|
||||
|
||||
async def _reserve_delivery_slots(self, requested: int) -> int:
|
||||
async with self._delivery_lock:
|
||||
available = max(self._max_inflight_requests() - self._inflight_requests, 0)
|
||||
admitted = min(max(requested, 0), available)
|
||||
self._inflight_requests += admitted
|
||||
return admitted
|
||||
|
||||
async def _release_delivery_slots(self, released: int) -> None:
|
||||
async with self._delivery_lock:
|
||||
self._inflight_requests = max(self._inflight_requests - released, 0)
|
||||
|
||||
async def _push_to_webhooks(self, webhooks: list[dict], payload: dict) -> list[object]:
|
||||
"""Dispatch only requests admitted by the instance-wide hard bound."""
|
||||
|
||||
admitted = await self._reserve_delivery_slots(len(webhooks))
|
||||
if admitted < len(webhooks):
|
||||
self.logger.warning(
|
||||
'Webhook delivery capacity reached; skipped %d of %d destinations',
|
||||
len(webhooks) - admitted,
|
||||
len(webhooks),
|
||||
)
|
||||
if admitted == 0:
|
||||
return []
|
||||
|
||||
tasks = [asyncio.create_task(self._push_to_webhook(webhook['url'], payload)) for webhook in webhooks[:admitted]]
|
||||
try:
|
||||
return await asyncio.gather(*tasks, return_exceptions=True)
|
||||
except asyncio.CancelledError:
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
raise
|
||||
finally:
|
||||
await self._release_delivery_slots(admitted)
|
||||
|
||||
async def push_person_message(
|
||||
self,
|
||||
@@ -58,9 +112,7 @@ class WebhookPusher:
|
||||
},
|
||||
}
|
||||
|
||||
# Push to all webhooks asynchronously
|
||||
tasks = [self._push_to_webhook(webhook['url'], payload) for webhook in webhooks]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
results = await self._push_to_webhooks(webhooks, payload)
|
||||
|
||||
# Check if any webhook responded with skip_pipeline=true
|
||||
for result in results:
|
||||
@@ -111,9 +163,7 @@ class WebhookPusher:
|
||||
},
|
||||
}
|
||||
|
||||
# Push to all webhooks asynchronously
|
||||
tasks = [self._push_to_webhook(webhook['url'], payload) for webhook in webhooks]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
results = await self._push_to_webhooks(webhooks, payload)
|
||||
|
||||
# Check if any webhook responded with skip_pipeline=true
|
||||
for result in results:
|
||||
|
||||
@@ -905,17 +905,27 @@ if not path.startswith('/workspace'):
|
||||
print(json.dumps({{'ok': False, 'error': 'Path must be under /workspace.'}}))
|
||||
elif not os.path.isfile(path):
|
||||
print(json.dumps({{'ok': False, 'error': f'File not found: {{path}}'}}))
|
||||
elif os.path.getsize(path) > {_MAX_HOST_EDIT_FILE_BYTES}:
|
||||
print(json.dumps({{'ok': False, 'error': 'File exceeds the {_MAX_HOST_EDIT_FILE_BYTES}-byte edit limit.'}}))
|
||||
else:
|
||||
with open(path, 'r', encoding='utf-8', errors='replace') as f:
|
||||
content = f.read()
|
||||
with open(path, 'rb') as f:
|
||||
raw_content = f.read({_MAX_HOST_EDIT_FILE_BYTES + 1})
|
||||
if len(raw_content) > {_MAX_HOST_EDIT_FILE_BYTES}:
|
||||
print(json.dumps({{'ok': False, 'error': 'File exceeds the {_MAX_HOST_EDIT_FILE_BYTES}-byte edit limit.'}}))
|
||||
raise SystemExit(0)
|
||||
content = raw_content.decode('utf-8', errors='replace')
|
||||
count = content.count(old_string)
|
||||
if count == 0:
|
||||
print(json.dumps({{'ok': False, 'error': 'old_string not found in file.'}}))
|
||||
elif count > 1:
|
||||
print(json.dumps({{'ok': False, 'error': f'old_string matches {{count}} locations; provide a more unique string.'}}))
|
||||
else:
|
||||
new_content = content.replace(old_string, new_string, 1)
|
||||
if len(new_content.encode('utf-8')) > {_MAX_HOST_EDIT_FILE_BYTES}:
|
||||
print(json.dumps({{'ok': False, 'error': 'Edited file exceeds the {_MAX_HOST_EDIT_FILE_BYTES}-byte limit.'}}))
|
||||
raise SystemExit(0)
|
||||
with open(path, 'w', encoding='utf-8') as f:
|
||||
f.write(content.replace(old_string, new_string, 1))
|
||||
f.write(new_content)
|
||||
print(json.dumps({{'ok': True, 'path': path}}))
|
||||
""".strip()
|
||||
return await self._run_workspace_file_script(script, query)
|
||||
|
||||
@@ -14,6 +14,7 @@ from .providers import localstorage
|
||||
|
||||
|
||||
_SAFE_OWNER_TYPE = re.compile(r'^[a-z][a-z0-9_-]{0,63}$')
|
||||
_DEFAULT_OBJECT_READ_BYTES = 10 * 1024 * 1024
|
||||
_SCOPED_KEY = re.compile(
|
||||
r'^v1/(?P<instance>[a-f0-9]{24})/'
|
||||
r'(?P<workspace>[0-9a-fA-F-]{36})/'
|
||||
@@ -34,6 +35,35 @@ class StorageMgr:
|
||||
def __init__(self, ap: app.Application):
|
||||
self.ap = ap
|
||||
|
||||
def _object_read_limit(self) -> int:
|
||||
config = getattr(getattr(self.ap, 'instance_config', None), 'data', {})
|
||||
try:
|
||||
configured = int(
|
||||
config.get('storage', {}).get(
|
||||
'max_object_read_bytes',
|
||||
_DEFAULT_OBJECT_READ_BYTES,
|
||||
)
|
||||
)
|
||||
except (AttributeError, TypeError, ValueError):
|
||||
configured = _DEFAULT_OBJECT_READ_BYTES
|
||||
return min(max(configured, 1), provider.HARD_MAX_STORAGE_OBJECT_BYTES)
|
||||
|
||||
async def _load_object_bounded(self, object_key: str) -> bytes:
|
||||
max_bytes = self._object_read_limit()
|
||||
bounded_loader = getattr(self.storage_provider, 'load_bounded', None)
|
||||
if callable(bounded_loader):
|
||||
return await bounded_loader(object_key, max_bytes=max_bytes)
|
||||
|
||||
# Compatibility for lightweight and third-party providers. Built-in
|
||||
# providers enforce the same bound in the actual read operation.
|
||||
object_size = await self.storage_provider.size(object_key)
|
||||
if object_size > max_bytes:
|
||||
raise ValueError(f'Storage object exceeds the {max_bytes}-byte read limit')
|
||||
value = await self.storage_provider.load(object_key)
|
||||
if len(value) > max_bytes:
|
||||
raise ValueError(f'Storage object exceeds the {max_bytes}-byte read limit')
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _require_execution_scope(
|
||||
context: ExecutionContext | RequestContext,
|
||||
@@ -150,6 +180,9 @@ class StorageMgr:
|
||||
preserve_suffix: bool = True,
|
||||
) -> str:
|
||||
await self._require_active_execution_scope(context)
|
||||
max_bytes = self._object_read_limit()
|
||||
if len(value) > max_bytes:
|
||||
raise ValueError(f'Storage object exceeds the {max_bytes}-byte write limit')
|
||||
object_key = self.scoped_object_key(
|
||||
context,
|
||||
owner_type=owner_type,
|
||||
@@ -177,7 +210,7 @@ class StorageMgr:
|
||||
key=key,
|
||||
preserve_suffix=preserve_suffix,
|
||||
)
|
||||
return await self.storage_provider.load(object_key)
|
||||
return await self._load_object_bounded(object_key)
|
||||
|
||||
async def delete_scoped(
|
||||
self,
|
||||
@@ -223,7 +256,7 @@ class StorageMgr:
|
||||
return None
|
||||
if not await self.storage_provider.exists(object_key):
|
||||
return None
|
||||
return await self.storage_provider.load(object_key)
|
||||
return await self._load_object_bounded(object_key)
|
||||
|
||||
@classmethod
|
||||
def require_scoped_object_key(
|
||||
@@ -274,7 +307,7 @@ class StorageMgr:
|
||||
object_key,
|
||||
expected_owner_type=expected_owner_type,
|
||||
)
|
||||
return await self.storage_provider.load(object_key)
|
||||
return await self._load_object_bounded(object_key)
|
||||
|
||||
async def size_scoped_object_key(
|
||||
self,
|
||||
|
||||
@@ -5,6 +5,19 @@ import abc
|
||||
from ..core import app
|
||||
|
||||
|
||||
HARD_MAX_STORAGE_OBJECT_BYTES = 64 * 1024 * 1024
|
||||
|
||||
|
||||
def normalize_read_limit(max_bytes: int) -> int:
|
||||
"""Validate a provider read limit without allowing callers to bypass the hard cap."""
|
||||
|
||||
try:
|
||||
normalized = int(max_bytes)
|
||||
except (TypeError, ValueError):
|
||||
normalized = HARD_MAX_STORAGE_OBJECT_BYTES
|
||||
return min(max(normalized, 1), HARD_MAX_STORAGE_OBJECT_BYTES)
|
||||
|
||||
|
||||
class StorageProvider(abc.ABC):
|
||||
ap: app.Application
|
||||
|
||||
@@ -34,6 +47,23 @@ class StorageProvider(abc.ABC):
|
||||
) -> bytes:
|
||||
pass
|
||||
|
||||
async def load_bounded(self, key: str, *, max_bytes: int) -> bytes:
|
||||
"""Fallback for third-party providers that have not implemented streaming bounds.
|
||||
|
||||
Built-in providers override this method so the byte limit is enforced by
|
||||
the actual read. The size check still protects compatible providers from
|
||||
downloading a known oversized object.
|
||||
"""
|
||||
|
||||
max_bytes = normalize_read_limit(max_bytes)
|
||||
object_size = await self.size(key)
|
||||
if object_size > max_bytes:
|
||||
raise ValueError(f'Storage object exceeds the {max_bytes}-byte read limit')
|
||||
value = await self.load(key)
|
||||
if len(value) > max_bytes:
|
||||
raise ValueError(f'Storage object exceeds the {max_bytes}-byte read limit')
|
||||
return value
|
||||
|
||||
@abc.abstractmethod
|
||||
async def exists(
|
||||
self,
|
||||
|
||||
@@ -51,9 +51,21 @@ class LocalStorageProvider(provider.StorageProvider):
|
||||
self,
|
||||
key: str,
|
||||
) -> bytes:
|
||||
return await self.load_bounded(key, max_bytes=provider.HARD_MAX_STORAGE_OBJECT_BYTES)
|
||||
|
||||
async def load_bounded(
|
||||
self,
|
||||
key: str,
|
||||
*,
|
||||
max_bytes: int,
|
||||
) -> bytes:
|
||||
max_bytes = provider.normalize_read_limit(max_bytes)
|
||||
resolved = await asyncio.to_thread(_safe_resolve, LOCAL_STORAGE_PATH, key)
|
||||
async with aiofiles.open(resolved, 'rb') as f:
|
||||
return await f.read()
|
||||
value = await f.read(max_bytes + 1)
|
||||
if len(value) > max_bytes:
|
||||
raise ValueError(f'Storage object exceeds the {max_bytes}-byte read limit')
|
||||
return value
|
||||
|
||||
async def exists(
|
||||
self,
|
||||
|
||||
@@ -101,22 +101,37 @@ class S3StorageProvider(provider.StorageProvider):
|
||||
async def load(
|
||||
self,
|
||||
key: str,
|
||||
) -> bytes:
|
||||
return await self.load_bounded(key, max_bytes=provider.HARD_MAX_STORAGE_OBJECT_BYTES)
|
||||
|
||||
async def load_bounded(
|
||||
self,
|
||||
key: str,
|
||||
*,
|
||||
max_bytes: int,
|
||||
) -> bytes:
|
||||
"""Load bytes from S3"""
|
||||
max_bytes = provider.normalize_read_limit(max_bytes)
|
||||
try:
|
||||
return await self._run_io(self._load_sync, key)
|
||||
return await self._run_io(self._load_sync, key, max_bytes)
|
||||
except Exception as e:
|
||||
self.ap.logger.error(f'Failed to load from S3: {e}')
|
||||
raise
|
||||
|
||||
def _load_sync(self, key: str) -> bytes:
|
||||
def _load_sync(self, key: str, max_bytes: int) -> bytes:
|
||||
response = self.s3_client.get_object(
|
||||
Bucket=self.bucket_name,
|
||||
Key=key,
|
||||
)
|
||||
body = response['Body']
|
||||
try:
|
||||
return body.read()
|
||||
declared_size = response.get('ContentLength')
|
||||
if declared_size is not None and declared_size > max_bytes:
|
||||
raise ValueError(f'Storage object exceeds the {max_bytes}-byte read limit')
|
||||
value = body.read(max_bytes + 1)
|
||||
if len(value) > max_bytes:
|
||||
raise ValueError(f'Storage object exceeds the {max_bytes}-byte read limit')
|
||||
return value
|
||||
finally:
|
||||
body.close()
|
||||
|
||||
|
||||
@@ -82,6 +82,7 @@ _MATCH_ALL = '-@file_id:{__langbot_match_all_sentinel__}'
|
||||
# files/filters matching more than one page of chunks are fully removed
|
||||
# (no silent truncation / orphaned vectors).
|
||||
_DELETE_SCAN_BATCH = 10000
|
||||
_MAX_DELETE_SCAN_ROUNDS = 1000
|
||||
|
||||
# Characters Valkey Search's TAG query parser cannot handle even when
|
||||
# backslash-escaped (the brace delimiters and the wildcard). file_id TAG
|
||||
@@ -652,11 +653,9 @@ class ValkeySearchVectorDatabase(VectorDatabase):
|
||||
return
|
||||
|
||||
query = f'@{_FIELD_FILE_ID}:{{{self._encode_and_escape_tag(file_id)}}}'
|
||||
keys = await self._search_keys(client, index, query)
|
||||
if keys:
|
||||
await client.delete(keys)
|
||||
deleted = await self._delete_search_results(client, index, query)
|
||||
self.ap.logger.info(
|
||||
f"Deleted {len(keys)} embeddings from Valkey Search collection '{collection}' with file_id: {file_id}"
|
||||
f"Deleted {deleted} embeddings from Valkey Search collection '{collection}' with file_id: {file_id}"
|
||||
)
|
||||
|
||||
async def delete_by_filter(self, collection: str, filter: dict[str, Any]) -> int:
|
||||
@@ -676,11 +675,9 @@ class ValkeySearchVectorDatabase(VectorDatabase):
|
||||
collection,
|
||||
)
|
||||
return 0
|
||||
keys = await self._search_keys(client, index, query)
|
||||
if keys:
|
||||
await client.delete(keys)
|
||||
self.ap.logger.info(f"Deleted {len(keys)} embeddings from Valkey Search collection '{collection}' by filter")
|
||||
return len(keys)
|
||||
deleted = await self._delete_search_results(client, index, query)
|
||||
self.ap.logger.info(f"Deleted {deleted} embeddings from Valkey Search collection '{collection}' by filter")
|
||||
return deleted
|
||||
|
||||
async def list_by_filter(
|
||||
self,
|
||||
@@ -783,38 +780,30 @@ class ValkeySearchVectorDatabase(VectorDatabase):
|
||||
except RequestError:
|
||||
return False
|
||||
|
||||
async def _search_keys(self, client: GlideClient, index: str, query: str) -> list[str]:
|
||||
"""Return all matching document keys for a query (NOCONTENT).
|
||||
async def _delete_search_results(self, client: GlideClient, index: str, query: str) -> int:
|
||||
"""Delete matching hashes in fixed batches without retaining every key.
|
||||
|
||||
Paginates through the full result set in pages of ``_DELETE_SCAN_BATCH``
|
||||
so that queries matching more than one page of chunks are fully
|
||||
enumerated (avoids silently truncating deletes and leaving orphaned
|
||||
vectors).
|
||||
Each deletion shrinks the result set, so every search starts at offset
|
||||
zero. Advancing an offset after deleting the preceding page would skip
|
||||
records as the remaining results shift left.
|
||||
"""
|
||||
keys: list[str] = []
|
||||
offset = 0
|
||||
while True:
|
||||
|
||||
deleted = 0
|
||||
for _round in range(_MAX_DELETE_SCAN_ROUNDS):
|
||||
options = FtSearchOptions(
|
||||
nocontent=True,
|
||||
limit=FtSearchLimit(offset, _DELETE_SCAN_BATCH),
|
||||
limit=FtSearchLimit(0, _DELETE_SCAN_BATCH),
|
||||
dialect=2,
|
||||
)
|
||||
try:
|
||||
reply = await ft.search(client, index, query, options)
|
||||
except Exception as exc:
|
||||
if self._is_missing_index_error(exc):
|
||||
return keys
|
||||
return deleted
|
||||
raise
|
||||
|
||||
if not reply or len(reply) < 2:
|
||||
break
|
||||
|
||||
# reply[0] is the total match count; reply[1] holds this page.
|
||||
total = 0
|
||||
try:
|
||||
total = int(reply[0])
|
||||
except (TypeError, ValueError):
|
||||
total = 0
|
||||
return deleted
|
||||
|
||||
docs = reply[1]
|
||||
if isinstance(docs, dict):
|
||||
@@ -825,11 +814,17 @@ class ValkeySearchVectorDatabase(VectorDatabase):
|
||||
page = []
|
||||
|
||||
if not page:
|
||||
break
|
||||
keys.extend(page)
|
||||
return deleted
|
||||
await client.delete(page)
|
||||
deleted += len(page)
|
||||
|
||||
offset += len(page)
|
||||
if offset >= total or len(page) < _DELETE_SCAN_BATCH:
|
||||
break
|
||||
try:
|
||||
total = int(reply[0])
|
||||
except (TypeError, ValueError):
|
||||
total = len(page)
|
||||
if total <= len(page) or len(page) < _DELETE_SCAN_BATCH:
|
||||
return deleted
|
||||
|
||||
return keys
|
||||
raise RuntimeError(
|
||||
f'Valkey deletion exceeded {_MAX_DELETE_SCAN_ROUNDS} batches ({_DELETE_SCAN_BATCH} keys per batch)'
|
||||
)
|
||||
|
||||
@@ -50,6 +50,16 @@ concurrency:
|
||||
# Hard admission limits for queued + running pipeline queries.
|
||||
pending_queries: 1000
|
||||
pending_queries_per_workspace: 100
|
||||
webhooks:
|
||||
# Bound database materialization and per-message outbound fan-out.
|
||||
# Existing rows above this limit remain deletable through the management
|
||||
# API, but only this many enabled destinations are dispatched.
|
||||
# Supports WEBHOOKS__MAX_PER_WORKSPACE (hard cap: 64).
|
||||
max_per_workspace: 16
|
||||
# Instance-wide request admission. Delivery fails open when every slot is
|
||||
# occupied instead of retaining an unbounded queue of webhook tasks.
|
||||
# Supports WEBHOOKS__MAX_INFLIGHT_REQUESTS (hard cap: 128).
|
||||
max_inflight_requests: 16
|
||||
cloud:
|
||||
# Operational safety ceilings for the one logical Cloud instance. These
|
||||
# are not subscription entitlements. An authoritative directory update
|
||||
@@ -208,6 +218,9 @@ vdb:
|
||||
request_timeout: 5000 # per-request timeout in ms (glide default 250ms is too low for KNN)
|
||||
storage:
|
||||
use: local
|
||||
# Bound every object materialized into Core memory. Built-in Local/S3
|
||||
# providers enforce this while reading (hard cap: 64 MiB).
|
||||
max_object_read_bytes: 10485760
|
||||
cleanup:
|
||||
# Enable periodic cleanup of local/S3 uploaded files and old log files
|
||||
enabled: true
|
||||
@@ -319,6 +332,10 @@ box:
|
||||
max_sessions: 64
|
||||
max_managed_processes: 64
|
||||
max_completed_processes: 256
|
||||
# Core scans a Workspace before and after quota-enforced executions.
|
||||
# Fail closed instead of repeatedly walking an inode bomb.
|
||||
# Supports BOX__LIMITS__MAX_WORKSPACE_ENTRIES (hard cap: 1000000).
|
||||
max_workspace_entries: 100000
|
||||
# Retained admission fences prevent replay after entitlement expiry or
|
||||
# revocation. Fail closed before that monotonic state can grow without
|
||||
# bound; Cloud may override this with BOX__LIMITS__MAX_ADMISSION_RECORDS.
|
||||
|
||||
Reference in New Issue
Block a user