xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-23 02:27:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): add Workspace multi-tenant foundation (#2353)

Browse Source 
* Document multi-tenant workspace architecture

* Add OSS and commercial workspace boundaries

* docs: redesign multi-tenant workspace architecture

* feat(tenancy): implement workspace isolation

* docs(tenancy): record verification evidence

* docs(tenancy): revise single-instance SaaS topology

* docs(tenancy): refine architecture options

* docs: finalize cloud v2 multi-tenant decisions

* feat(tenancy): establish cloud isolation foundations

* feat(tenancy): harden shared cloud runtime boundaries

* docs(tenancy): record final isolation verification

* fix(tenancy): close isolation and permission gaps

* docs(tenancy): record final isolation verification

* feat(tenancy): connect cloud workspace control plane

* fix(build): install git for pinned SDK

* docs(cloud): update control plane verification

* chore: update multi-tenant SDK pin

* fix(cloud): skip legacy model sync during startup

* test(cloud): preserve minimal model manager fixtures

* fix(cloud): preserve authenticated account context

* fix(cloud): reuse authenticated account for user info

* feat(cloud): complete Workspace settings navigation

* test(web): cover Workspace dropdown menu

* feat(web): place workspace controls in sidebar

* refactor(web): streamline workspace controls

* style(web): format workspace layout test

* fix(cloud): surface runtime and workspace plan status

* fix(plugin): keep runtime identity stable across restarts

* fix(ui): widen and center workspace switcher

* fix(ui): hide roles from workspace switcher

* fix(ui): align workspace switcher with sidebar entries

* feat(workspace): add in-product collaboration and direct Cloud launch

* style: format collaboration changes

* fix(workspace): bind collaboration APIs to tenant UoW

* fix(cloud): preserve Core-owned collaboration state

* test(cloud): require Space identity for invite registration

* feat(cloud): complete secure invitation experience

* style(web): format invitation flows

* fix(cloud): recover box runtime without unscoped skill reload

* feat(oss): enforce invitation account and owner billing flows

* style: format OSS account service

* test(oss): cover invitation logout handoff

* fix(oss): resolve workspace owner in scoped session

* feat(cloud): harden multi-tenant runtime resources

* fix(cloud): bound runtime restart storms

* fix(cloud): eliminate periodic runtime CPU spikes

* fix(cloud): enforce instance capacity ceilings

* fix(cloud): scope public login capability discovery

* fix(cloud): bound tenant maintenance and monitoring work

* fix(runtime): bound tenant resource amplification

* fix(deps): pin green multi-tenant plugin SDK

* fix(cloud): handle unavailable skill capability

* fix(security): require authentication for image file endpoint (H-2)

- Changed /api/v1/files/image from AuthType.NONE to USER_TOKEN_OR_API_KEY
- Added Permission.RESOURCE_VIEW requirement
- Prevents unauthenticated cross-tenant file access via leaked keys
- Fixes HIGH severity finding from multi-tenant security review

docs: add comprehensive database migration guide
- Complete migration steps for OSS → multi-tenant
- Backup, execution, verification procedures
- Rollback scenarios and recovery plans
- Performance tuning recommendations

* test: add comprehensive cross-tenant isolation tests

Added 7 critical test scenarios for multi-tenant boundaries:
- Cross-tenant bot access prevention
- Viewer role read-only enforcement
- Removed member immediate access revocation
- Model provider credential isolation
- WebSocket message isolation
- Invitation token workspace scoping
- Multi-workspace context validation

These tests address P0-2 coverage gaps for:
- workspaces.py (membership & invitation flows)
- user.py (authentication & authorization)
- websocket_chat.py (real-time isolation)
- plugins.py (resource access control)

docs: finalize database migration guide

* fix(security): resolve M-1, M-2, M-3 security findings

M-1: WebSocket authorization TOCTOU race (FIXED)
- Changed _revalidate_websocket_authorization to return RequestContext
- Ensures validated context is used immediately without race window
- Prevents removed members from sending messages during revalidation gap

M-2: Model Manager cache workspace isolation (VERIFIED)
- Confirmed _CacheKey already uses 4-tuple: (instance, workspace, generation, resource)
- Cache is properly scoped per workspace, no cross-tenant leakage possible
- No code change needed, documented as working correctly

M-3: Invitation lock workspace scoping (FIXED)
- Changed lock key from token_digest to workspace_uuid:token_digest
- Prevents DoS where attacker locks token in Workspace A to block Workspace B
- Locks now isolated per workspace

All MEDIUM severity findings from security review now resolved.

* fix(cloud): unblock tenant CI and enforce knowledge quotas

* fix(tenancy): scope rerank model sync

---------

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
RockChinQ
2026-07-30 21:43:35 +08:00
committed by GitHub
parent 463b120923
commit e1ac5e0fc8
468 changed files with 78320 additions and 13137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/langbot/pkg/pipeline/aggregator.py
+236 -147
View File
@@ -1,10 +1,4 @@
"""Message Aggregator Module
This module provides message aggregation/debounce functionality.
When users send multiple messages consecutively, the aggregator will wait
for a configurable delay period and merge them into a single message
before processing.
"""
"""Workspace-scoped message aggregation and debounce support."""
from __future__ import annotations
@@ -13,96 +7,125 @@ import time
import typing
from dataclasses import dataclass, field
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.provider.session as provider_session
import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
import langbot_plugin.api.entities.builtin.platform.events as platform_events
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 ..core.task_boundary import create_detached_task, run_in_workspace_uow
from .pool import ExecutionContextMismatchError
from ..workspace.errors import WorkspaceError, WorkspaceInvariantError
if typing.TYPE_CHECKING:
from ..core import app
# Maximum number of messages to buffer before forcing a flush
MAX_BUFFER_MESSAGES = 10
AggregationKey = tuple[
str,
str,
int,
str,
str | None,
str,
int | str,
]
@dataclass
class PendingMessage:
"""A pending message waiting to be aggregated"""
"""A pending message carrying its trusted execution scope."""
execution_context: ExecutionContext
bot_uuid: str
launcher_type: provider_session.LauncherTypes
launcher_id: typing.Union[int, str]
sender_id: typing.Union[int, str]
launcher_id: int | str
sender_id: int | str
message_event: platform_events.MessageEvent
message_chain: platform_message.MessageChain
adapter: abstract_platform_adapter.AbstractMessagePlatformAdapter
pipeline_uuid: typing.Optional[str]
pipeline_uuid: str | None
routed_by_rule: bool = False
timestamp: float = field(default_factory=time.time)
@dataclass
class SessionBuffer:
"""Buffer for a single session's pending messages"""
"""Pending messages for one scoped aggregation key."""
session_id: str
aggregation_key: AggregationKey
execution_context: ExecutionContext
messages: list[PendingMessage] = field(default_factory=list)
timer_task: typing.Optional[asyncio.Task] = None
timer_task: asyncio.Task | None = None
last_message_time: float = field(default_factory=time.time)
class MessageAggregator:
"""Message aggregator that buffers and merges consecutive messages
This class implements a debounce mechanism for incoming messages.
When a message arrives, it starts a timer. If more messages arrive
before the timer expires, they are buffered. When the timer expires,
all buffered messages are merged and sent to the query pool.
"""
"""Debounce consecutive messages without crossing Workspace boundaries."""
ap: app.Application
buffers: dict[str, SessionBuffer]
"""Session ID -> SessionBuffer mapping"""
buffers: dict[AggregationKey, SessionBuffer]
lock: asyncio.Lock
"""Lock for thread-safe buffer operations"""
def __init__(self, ap: app.Application):
self.ap = ap
self.buffers = {}
self._buffer_counts_by_scope: dict[
tuple[str, str, int],
int,
] = {}
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_session_id(
def _get_aggregation_key(
self,
execution_context: ExecutionContext,
bot_uuid: str,
launcher_type: provider_session.LauncherTypes,
launcher_id: typing.Union[int, str],
) -> str:
"""Generate a unique session ID"""
return f'{bot_uuid}:{launcher_type.value}:{launcher_id}'
launcher_id: int | str,
pipeline_uuid: str | None,
) -> AggregationKey:
"""Build a key that cannot alias another Workspace, bot, or pipeline."""
async def _get_aggregation_config(self, pipeline_uuid: typing.Optional[str]) -> tuple[bool, float]:
"""Get aggregation configuration for a pipeline
return (
execution_context.instance_uuid,
execution_context.workspace_uuid,
execution_context.placement_generation,
bot_uuid,
pipeline_uuid,
launcher_type.value,
launcher_id,
)
async def _get_aggregation_config(
self,
execution_context: ExecutionContext,
pipeline_uuid: str | None,
) -> tuple[bool, float]:
"""Return aggregation enablement and a clamped debounce delay."""
Returns:
tuple: (enabled, delay_seconds)
"""
default_enabled = False
default_delay = 1.5
if pipeline_uuid is None:
return default_enabled, default_delay
# Get pipeline from pipeline manager
pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(pipeline_uuid)
pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(
execution_context,
pipeline_uuid,
)
if pipeline is None:
return default_enabled, default_delay
config = pipeline.pipeline_entity.config or {}
trigger_config = config.get('trigger', {})
aggregation_config = trigger_config.get('message-aggregation', {})
enabled = aggregation_config.get('enabled', default_enabled)
delay_raw = aggregation_config.get('delay', default_delay)
@@ -111,33 +134,31 @@ class MessageAggregator:
except (TypeError, ValueError):
delay = default_delay
# Clamp delay to valid range
delay = max(1.0, min(10.0, delay))
return enabled, delay
return enabled, max(1.0, min(10.0, delay))
async def add_message(
self,
bot_uuid: str,
launcher_type: provider_session.LauncherTypes,
launcher_id: typing.Union[int, str],
sender_id: typing.Union[int, str],
launcher_id: int | str,
sender_id: int | str,
message_event: platform_events.MessageEvent,
message_chain: platform_message.MessageChain,
adapter: abstract_platform_adapter.AbstractMessagePlatformAdapter,
pipeline_uuid: typing.Optional[str] = None,
pipeline_uuid: str | None = None,
routed_by_rule: bool = False,
execution_context: ExecutionContext | None = None,
) -> None:
"""Add a message to the aggregation buffer
"""Buffer or directly enqueue a message in its trusted Workspace."""
If aggregation is disabled for the pipeline, the message is sent
directly to the query pool. Otherwise, it's buffered and will be
merged with other messages from the same session.
"""
enabled, delay = await self._get_aggregation_config(pipeline_uuid)
execution_context = await self.ap.query_pool.resolve_execution_context(
execution_context,
bot_uuid=bot_uuid,
pipeline_uuid=pipeline_uuid,
)
enabled, delay = await self._get_aggregation_config(execution_context, pipeline_uuid)
if not enabled:
# Aggregation disabled, send directly to query pool
await self.ap.query_pool.add_query(
bot_uuid=bot_uuid,
launcher_type=launcher_type,
@@ -148,12 +169,19 @@ class MessageAggregator:
adapter=adapter,
pipeline_uuid=pipeline_uuid,
routed_by_rule=routed_by_rule,
execution_context=execution_context,
)
return
session_id = self._get_session_id(bot_uuid, launcher_type, launcher_id)
aggregation_key = self._get_aggregation_key(
execution_context,
bot_uuid,
launcher_type,
launcher_id,
pipeline_uuid,
)
pending_msg = PendingMessage(
execution_context=execution_context,
bot_uuid=bot_uuid,
launcher_type=launcher_type,
launcher_id=launcher_id,
@@ -166,107 +194,167 @@ class MessageAggregator:
)
force_flush = False
bypass_aggregation = False
async with self.lock:
if session_id in self.buffers:
buffer = self.buffers[session_id]
# Cancel existing timer (just cancel, don't await inside lock)
buffer = self.buffers.get(aggregation_key)
if buffer is None:
scope_key = (
execution_context.instance_uuid,
execution_context.workspace_uuid,
execution_context.placement_generation,
)
workspace_buffer_count = self._buffer_counts_by_scope.get(
scope_key,
0,
)
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
self._buffer_counts_by_scope[scope_key] = workspace_buffer_count + 1
else:
if buffer.execution_context != execution_context:
raise ExecutionContextMismatchError('Aggregation buffer ExecutionContext changed for the same key')
if buffer.timer_task and not buffer.timer_task.done():
buffer.timer_task.cancel()
buffer.messages.append(pending_msg)
else:
buffer = SessionBuffer(
session_id=session_id,
messages=[pending_msg],
)
self.buffers[session_id] = buffer
buffer.last_message_time = time.time()
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,
)
# Check if buffer reached max capacity
if len(buffer.messages) >= MAX_BUFFER_MESSAGES:
force_flush = True
else:
# Start new timer
buffer.timer_task = asyncio.create_task(self._delayed_flush(session_id, delay))
if force_flush:
await self._flush_buffer(session_id)
async def _delayed_flush(self, session_id: str, delay: float) -> None:
"""Wait for delay then flush the buffer"""
try:
await asyncio.sleep(delay)
await self._flush_buffer(session_id)
except asyncio.CancelledError:
# Timer was cancelled, new message arrived
pass
async def _flush_buffer(self, session_id: str) -> None:
"""Flush the buffer for a session, merging all messages"""
async with self.lock:
buffer = self.buffers.pop(session_id, None)
if buffer is None or not buffer.messages:
return
if len(buffer.messages) == 1:
# Only one message, no need to merge
msg = buffer.messages[0]
if bypass_aggregation:
await self.ap.query_pool.add_query(
bot_uuid=msg.bot_uuid,
launcher_type=msg.launcher_type,
launcher_id=msg.launcher_id,
sender_id=msg.sender_id,
message_event=msg.message_event,
message_chain=msg.message_chain,
adapter=msg.adapter,
pipeline_uuid=msg.pipeline_uuid,
routed_by_rule=msg.routed_by_rule,
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
# Merge multiple messages
merged_msg = self._merge_messages(buffer.messages)
if force_flush:
await self._flush_buffer(aggregation_key, execution_context)
async def _delayed_flush(
self,
aggregation_key: AggregationKey,
delay: float,
execution_context: ExecutionContext,
) -> None:
"""Flush after the debounce delay using the captured context."""
try:
await asyncio.sleep(delay)
await run_in_workspace_uow(
self.ap,
execution_context.workspace_uuid,
lambda: self._flush_buffer(aggregation_key, execution_context),
)
except asyncio.CancelledError:
pass
except WorkspaceError as exc:
self.ap.logger.info(
f'Dropped an aggregated message because its Workspace execution binding is stale: {exc}'
)
async def _flush_buffer(
self,
aggregation_key: AggregationKey,
execution_context: ExecutionContext,
) -> None:
"""Flush one buffer only when the captured scope still matches."""
async with self.lock:
buffer = self.buffers.get(aggregation_key)
if buffer is None:
return
if buffer.execution_context != execution_context:
raise ExecutionContextMismatchError('Timer ExecutionContext does not match the aggregation buffer')
self.buffers.pop(aggregation_key)
scope_key = aggregation_key[:3]
scope_count = self._buffer_counts_by_scope.get(scope_key, 0)
if scope_count <= 1:
self._buffer_counts_by_scope.pop(scope_key, None)
else:
self._buffer_counts_by_scope[scope_key] = scope_count - 1
if not buffer.messages:
return
message = buffer.messages[0] if len(buffer.messages) == 1 else self._merge_messages(buffer.messages)
binding = await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
if binding.instance_uuid != execution_context.instance_uuid:
raise WorkspaceInvariantError('Aggregation buffer instance does not match the active Workspace binding')
await self.ap.query_pool.add_query(
bot_uuid=merged_msg.bot_uuid,
launcher_type=merged_msg.launcher_type,
launcher_id=merged_msg.launcher_id,
sender_id=merged_msg.sender_id,
message_event=merged_msg.message_event,
message_chain=merged_msg.message_chain,
adapter=merged_msg.adapter,
pipeline_uuid=merged_msg.pipeline_uuid,
routed_by_rule=merged_msg.routed_by_rule,
bot_uuid=message.bot_uuid,
launcher_type=message.launcher_type,
launcher_id=message.launcher_id,
sender_id=message.sender_id,
message_event=message.message_event,
message_chain=message.message_chain,
adapter=message.adapter,
pipeline_uuid=message.pipeline_uuid,
routed_by_rule=message.routed_by_rule,
execution_context=message.execution_context,
)
def _merge_messages(self, messages: list[PendingMessage]) -> PendingMessage:
"""Merge multiple messages into one
"""Merge message chains after proving all messages share one scope."""
The merged message uses the first message as base and combines
all message chains with newline separators.
The original message_event is kept unmodified to preserve
message metadata (message_id, etc.) for reply/quote.
"""
if not messages:
raise ValueError('At least one pending message is required')
if len(messages) == 1:
return messages[0]
base_msg = messages[0]
base_key = self._get_aggregation_key(
base_msg.execution_context,
base_msg.bot_uuid,
base_msg.launcher_type,
base_msg.launcher_id,
base_msg.pipeline_uuid,
)
for message in messages[1:]:
message_key = self._get_aggregation_key(
message.execution_context,
message.bot_uuid,
message.launcher_type,
message.launcher_id,
message.pipeline_uuid,
)
if message_key != base_key or message.execution_context != base_msg.execution_context:
raise ExecutionContextMismatchError('Cannot merge pending messages from different execution scopes')
# Build merged message chain
merged_chain = platform_message.MessageChain([])
for i, msg in enumerate(messages):
if i > 0:
# Add newline separator between messages
for index, message in enumerate(messages):
if index > 0:
merged_chain.append(platform_message.Plain(text='\n'))
# Copy all components from this message
for component in msg.message_chain:
for component in message.message_chain:
merged_chain.append(component)
# Keep message_event unmodified (preserves original message_id and
# metadata for reply/quote), only pass merged chain separately
return PendingMessage(
execution_context=base_msg.execution_context,
bot_uuid=base_msg.bot_uuid,
launcher_type=base_msg.launcher_type,
launcher_id=base_msg.launcher_id,
@@ -275,22 +363,23 @@ class MessageAggregator:
message_chain=merged_chain,
adapter=base_msg.adapter,
pipeline_uuid=base_msg.pipeline_uuid,
routed_by_rule=any(msg.routed_by_rule for msg in messages),
routed_by_rule=any(message.routed_by_rule for message in messages),
)
async def flush_all(self) -> None:
"""Flush all pending buffers immediately
"""Flush all pending buffers without dropping their captured scopes."""
This is useful during shutdown to ensure no messages are lost.
"""
# Snapshot session IDs and cancel all timers under lock
async with self.lock:
session_ids = list(self.buffers.keys())
for sid in session_ids:
buffer = self.buffers.get(sid)
if buffer and buffer.timer_task and not buffer.timer_task.done():
pending_buffers = [(key, buffer.execution_context) for key, buffer in self.buffers.items()]
for buffer in self.buffers.values():
if buffer.timer_task and not buffer.timer_task.done():
buffer.timer_task.cancel()
# Flush each buffer outside the lock
for session_id in session_ids:
await self._flush_buffer(session_id)
for aggregation_key, execution_context in pending_buffers:
try:
await self._flush_buffer(aggregation_key, execution_context)
except WorkspaceError as exc:
self.ap.logger.info(
'Dropped an aggregated message during shutdown because its '
f'Workspace execution binding is stale: {exc}'
)
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
+129 -39
View File
@@ -5,8 +5,10 @@ import traceback
from ..core import app
from ..core import entities as core_entities
from ..workspace.errors import WorkspaceError, WorkspaceInvariantError
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from .pool import get_query_execution_context
class Controller:
@@ -21,11 +23,95 @@ class Controller:
self.ap = ap
self.semaphore = asyncio.Semaphore(self.ap.instance_config.data['concurrency']['pipeline'])
async def _assert_query_execution_active(
self,
query: pipeline_query.Query,
):
"""Revalidate a queued query immediately before runtime work starts."""
execution_context = get_query_execution_context(query)
binding = await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
if binding.instance_uuid != execution_context.instance_uuid:
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,
*,
selected_session=None,
global_slot_reserved: bool = False,
) -> None:
"""Run one selected query and always release its scheduling slot."""
try:
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
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'
)
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:
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:
@@ -38,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')
@@ -51,46 +139,48 @@ class Controller:
continue
if selected_query:
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
async def _process_query(selected_query: pipeline_query.Query):
async with self.semaphore: # 总并发上限
# find pipeline
# Here firstly find the bot, then find the pipeline, in case the bot adapter's config is not the latest one.
# Like aiocqhttp, once a client is connected, even the adapter was updated and restarted, the existing client connection will not be affected.
pipeline_uuid = selected_query.pipeline_uuid
if pipeline_uuid:
pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(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'
)
else:
self.ap.logger.warning(
f'No pipeline_uuid for query {selected_query.query_id}, query dropped'
)
async with self.ap.query_pool:
(await self.ap.sess_mgr.get_session(selected_query))._semaphore.release()
# 通知其他协程，有新的请求可以处理了
self.ap.query_pool.condition.notify_all()
self.ap.task_mgr.create_task(
_process_query(selected_query),
kind='query',
name=f'query-{selected_query.query_id}',
scopes=[
core_entities.LifecycleControlScope.APPLICATION,
core_entities.LifecycleControlScope.PLATFORM,
],
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 Exception as e:
# traceback.print_exc()
self.ap.logger.error(f'控制器循环出错: {e}')
self.ap.logger.error(f'Traceback: {traceback.format_exc()}')
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
+122 -123
View File
@@ -1,19 +1,32 @@
from __future__ import annotations
import asyncio
import os
import base64
import time
import re
import uuid
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):
@@ -28,28 +41,49 @@ 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,
)
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)
compressed_path, size = self.compress_image(img_path, outfile='temp/{}_compressed.png'.format(int(time.time())))
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:
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):
os.remove(path)
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.
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(
base64=b64.decode('utf-8'),
base64=image_base64,
)
]
@@ -59,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
@@ -118,14 +121,55 @@ 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)
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."""
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:
append_line(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)))
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)))
append_line(rest_text[:point])
rest_text = rest_text[point:]
if rest_text and font.getlength(rest_text) < text_width:
append_line(rest_text)
break
return final_lines
def text_to_image(
self,
text_str: str,
@@ -133,79 +177,34 @@ 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
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=self.get_font(query.pipeline_config['output']['long-text-processing']['font-path']),
)
# 遍历此行,检查是否有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
src/langbot/pkg/pipeline/monitoring_helper.py
+10
View File
@@ -15,6 +15,8 @@ if typing.TYPE_CHECKING:
from ..core import app
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from .pool import get_query_execution_context
class MonitoringHelper:
"""Helper class for monitoring operations"""
@@ -54,6 +56,7 @@ class MonitoringHelper:
# Here we just record None, the full variables will be set when query completes
message_id = await ap.monitoring_service.record_message(
get_query_execution_context(query),
bot_id=bot_id,
bot_name=bot_name,
pipeline_id=pipeline_id,
@@ -74,6 +77,7 @@ class MonitoringHelper:
# Update session activity or create new session if it doesn't exist
# Always pass pipeline info to handle pipeline switches
session_updated = await ap.monitoring_service.update_session_activity(
get_query_execution_context(query),
session_id,
pipeline_id=pipeline_id,
pipeline_name=pipeline_name,
@@ -81,6 +85,7 @@ class MonitoringHelper:
if not session_updated:
# Session doesn't exist, create it
await ap.monitoring_service.record_session_start(
get_query_execution_context(query),
session_id=session_id,
bot_id=bot_id,
bot_name=bot_name,
@@ -118,6 +123,7 @@ class MonitoringHelper:
pass
await ap.monitoring_service.update_message_status(
get_query_execution_context(query),
message_id=message_id,
status='success',
variables=query_variables_str,
@@ -170,6 +176,7 @@ class MonitoringHelper:
return # No response to record
await ap.monitoring_service.record_message(
get_query_execution_context(query),
bot_id=bot_id,
bot_name=bot_name,
pipeline_id=pipeline_id,
@@ -215,6 +222,7 @@ class MonitoringHelper:
# Record error message
message_id = await ap.monitoring_service.record_message(
get_query_execution_context(query),
bot_id=bot_id,
bot_name=bot_name,
pipeline_id=pipeline_id,
@@ -233,6 +241,7 @@ class MonitoringHelper:
# Record error log
await ap.monitoring_service.record_error(
get_query_execution_context(query),
bot_id=bot_id,
bot_name=bot_name,
pipeline_id=pipeline_id,
@@ -271,6 +280,7 @@ class MonitoringHelper:
session_id = f'{query.launcher_type.value if hasattr(query.launcher_type, "value") else query.launcher_type}_{query.launcher_id}'
await ap.monitoring_service.record_llm_call(
get_query_execution_context(query),
bot_id=bot_id,
bot_name=bot_name,
pipeline_id=pipeline_id,
src/langbot/pkg/pipeline/pipelinemgr.py
+276 -17
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import dataclasses
import typing
import traceback
@@ -13,7 +14,11 @@ 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.events as events
from ..utils import importutil
from ..api.http.authz import WorkspaceRequiredError
from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType, RequestContext
from ..workspace.errors import WorkspaceError, WorkspaceInvariantError
from .config_coercion import coerce_pipeline_config
from .pool import get_query_execution_context
import langbot_plugin.api.entities.builtin.provider.session as provider_session
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
@@ -82,15 +87,39 @@ class RuntimePipeline:
enable_all_mcp_servers: bool
"""是否启用所有MCP服务器"""
execution_context: ExecutionContext
workspace_uuid: str
placement_generation: int
def __init__(
self,
ap: app.Application,
pipeline_entity: persistence_pipeline.LegacyPipeline,
stage_containers: list[StageInstContainer],
execution_context: ExecutionContext,
):
if not isinstance(execution_context, ExecutionContext):
raise WorkspaceRequiredError('RuntimePipeline requires an ExecutionContext')
if not execution_context.instance_uuid.strip() or not execution_context.workspace_uuid.strip():
raise WorkspaceRequiredError('RuntimePipeline requires an instance and Workspace')
if execution_context.placement_generation <= 0:
raise WorkspaceRequiredError('RuntimePipeline requires a positive placement generation')
if pipeline_entity.workspace_uuid != execution_context.workspace_uuid:
raise WorkspaceRequiredError('RuntimePipeline entity Workspace does not match its ExecutionContext')
if execution_context.pipeline_uuid not in (None, pipeline_entity.uuid):
raise WorkspaceRequiredError('RuntimePipeline UUID does not match its ExecutionContext')
self.ap = ap
self.pipeline_entity = pipeline_entity
self.stage_containers = stage_containers
self.execution_context = dataclasses.replace(
execution_context,
pipeline_uuid=pipeline_entity.uuid,
)
self.workspace_uuid = self.execution_context.workspace_uuid
self.placement_generation = self.execution_context.placement_generation
# Extract bound plugins and MCP servers from extensions_preferences
extensions_prefs = pipeline_entity.extensions_preferences or {}
@@ -120,7 +149,37 @@ class RuntimePipeline:
mcp_server_list = extensions_prefs.get('mcp_servers', [])
self.bound_mcp_servers = mcp_server_list if mcp_server_list else []
async def _assert_execution_active(
self,
query: pipeline_query.Query | None = None,
) -> ExecutionContext:
"""Fail closed when this runtime or query belongs to a stale placement."""
execution_context = self.execution_context if query is None else get_query_execution_context(query)
if (
execution_context.instance_uuid != self.execution_context.instance_uuid
or execution_context.workspace_uuid != self.workspace_uuid
or execution_context.placement_generation != self.placement_generation
or execution_context.pipeline_uuid != self.pipeline_entity.uuid
):
raise WorkspaceInvariantError('Query execution scope does not match RuntimePipeline')
binding = await self.ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
if binding.instance_uuid != execution_context.instance_uuid:
raise WorkspaceInvariantError('RuntimePipeline instance does not match the active Workspace binding')
return execution_context
async def run(self, query: pipeline_query.Query):
if (
query.instance_uuid != self.execution_context.instance_uuid
or query.workspace_uuid != self.workspace_uuid
or query.placement_generation != self.placement_generation
or query.pipeline_uuid != self.pipeline_entity.uuid
):
raise WorkspaceRequiredError('Query execution scope does not match RuntimePipeline')
await self._assert_execution_active(query)
query.pipeline_config = self.pipeline_entity.config
# Store bound plugins and MCP servers in query for filtering
query.variables['_pipeline_bound_plugins'] = self.bound_plugins
@@ -134,7 +193,11 @@ class RuntimePipeline:
bot_name = 'WebChat'
if query.bot_uuid:
try:
bot = await self.ap.bot_service.get_bot(query.bot_uuid, include_secret=False)
bot = await self.ap.bot_service.get_bot(
query.workspace_uuid,
query.bot_uuid,
include_secret=False,
)
if bot:
bot_name = bot.get('name', 'Unknown')
except Exception:
@@ -150,6 +213,7 @@ class RuntimePipeline:
async def _check_output(self, query: pipeline_query.Query, result: pipeline_entities.StageProcessResult):
"""检查输出"""
await self._assert_execution_active(query)
if result.user_notice:
# 处理str类型
@@ -162,7 +226,9 @@ class RuntimePipeline:
query.message_event, platform_events.GroupMessage
):
result.user_notice.insert(0, platform_message.At(target=query.message_event.sender.id))
if await query.adapter.is_stream_output_supported() and query.resp_messages:
stream_output_supported = await query.adapter.is_stream_output_supported()
await self._assert_execution_active(query)
if stream_output_supported and query.resp_messages:
await query.adapter.reply_message_chunk(
message_source=query.message_event,
bot_message=query.resp_messages[-1],
@@ -186,6 +252,7 @@ class RuntimePipeline:
query.variables['_monitoring_has_error'] = True
# Record error to monitoring system
try:
await self._assert_execution_active(query)
bot_name = query.variables.get('_monitoring_bot_name', 'Unknown')
pipeline_name = query.variables.get('_monitoring_pipeline_name', 'Unknown')
message_id = query.variables.get('_monitoring_message_id', '')
@@ -194,6 +261,7 @@ class RuntimePipeline:
# Update message status to error
if message_id:
await self.ap.monitoring_service.update_message_status(
get_query_execution_context(query),
message_id=message_id,
status='error',
level='error',
@@ -201,6 +269,7 @@ class RuntimePipeline:
# Record error log
await self.ap.monitoring_service.record_error(
get_query_execution_context(query),
bot_id=query.bot_uuid or 'unknown',
bot_name=bot_name,
pipeline_id=self.pipeline_entity.uuid,
@@ -242,6 +311,7 @@ class RuntimePipeline:
i = stage_index
while i < len(self.stage_containers):
await self._assert_execution_active(query)
stage_container = self.stage_containers[i]
query.current_stage_name = stage_container.inst_name # 标记到 Query 对象里
@@ -250,6 +320,7 @@ class RuntimePipeline:
if isinstance(result, typing.Coroutine):
result = await result
await self._assert_execution_active(query)
if isinstance(result, pipeline_entities.StageProcessResult): # 直接返回结果
self.ap.logger.debug(
@@ -265,7 +336,14 @@ class RuntimePipeline:
elif isinstance(result, typing.AsyncGenerator): # 生成器
self.ap.logger.debug(f'Stage {stage_container.inst_name} processed query {query.query_id} gen')
async for sub_result in result:
iterator = result.__aiter__()
while True:
await self._assert_execution_active(query)
try:
sub_result = await anext(iterator)
except StopAsyncIteration:
break
await self._assert_execution_active(query)
self.ap.logger.debug(
f'Stage {stage_container.inst_name} processed query {query.query_id} res {sub_result.result_type}'
)
@@ -283,6 +361,7 @@ class RuntimePipeline:
async def process_query(self, query: pipeline_query.Query):
"""处理请求"""
await self._assert_execution_active(query)
# Get monitoring metadata
bot_name = query.variables.get('_monitoring_bot_name', 'Unknown')
pipeline_name = query.variables.get('_monitoring_pipeline_name', 'Unknown')
@@ -310,6 +389,7 @@ class RuntimePipeline:
query.variables['_monitoring_message_id'] = message_id
# Notify adapter so it can map platform-specific IDs to monitoring message ID
if hasattr(query.adapter, 'on_monitoring_message_created'):
await self._assert_execution_active(query)
await query.adapter.on_monitoring_message_created(query, message_id)
except Exception as e:
self.ap.logger.error(f'Failed to record query start: {e}')
@@ -334,7 +414,9 @@ class RuntimePipeline:
message_chain=query.message_chain,
)
await self._assert_execution_active(query)
event_ctx = await self.ap.plugin_connector.emit_event(event_obj, bound_plugins)
await self._assert_execution_active(query)
if event_ctx.is_prevented_default():
self.ap.logger.debug(
@@ -349,6 +431,7 @@ class RuntimePipeline:
# Record query success only if no error occurred during processing
if not query.variables.get('_monitoring_has_error', False):
try:
await self._assert_execution_active(query)
await monitoring_helper.MonitoringHelper.record_query_success(
ap=self.ap,
message_id=message_id,
@@ -359,6 +442,7 @@ class RuntimePipeline:
# Record bot response message
try:
await self._assert_execution_active(query)
await monitoring_helper.MonitoringHelper.record_query_response(
ap=self.ap,
query=query,
@@ -371,6 +455,8 @@ class RuntimePipeline:
except Exception as e:
self.ap.logger.error(f'Failed to record query response: {e}')
except WorkspaceError as e:
self.ap.logger.info(f'Dropped query {query.query_id} because its Workspace execution binding is stale: {e}')
except Exception as e:
inst_name = query.current_stage_name if query.current_stage_name else 'unknown'
self.ap.logger.error(f'Error processing query {query.query_id} stage={inst_name} : {e}')
@@ -380,6 +466,7 @@ class RuntimePipeline:
try:
from . import monitoring_helper
await self._assert_execution_active(query)
await monitoring_helper.MonitoringHelper.record_query_error(
ap=self.ap,
query=query,
@@ -395,7 +482,7 @@ class RuntimePipeline:
finally:
self.ap.logger.debug(f'Query {query.query_id} processed')
del self.ap.query_pool.cached_queries[query.query_id]
await self.ap.query_pool.remove_query(query)
class PipelineManager:
@@ -409,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()}
@@ -419,25 +549,96 @@ class PipelineManager:
async def load_pipelines_from_db(self):
self.ap.logger.info('Loading pipelines from db...')
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'
if cloud_runtime:
if not callable(list_bindings) or not callable(tenant_uow):
raise RuntimeError('Cloud pipeline loading requires explicit instance discovery and tenant UoWs')
for binding in await list_bindings():
async with tenant_uow(binding.workspace_uuid):
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_pipeline.LegacyPipeline)
.where(persistence_pipeline.LegacyPipeline.workspace_uuid == binding.workspace_uuid)
.order_by(persistence_pipeline.LegacyPipeline.uuid)
)
for pipeline in result.all():
await self.load_pipeline(
ExecutionContext(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
placement_generation=binding.placement_generation,
pipeline_uuid=pipeline.uuid,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
),
pipeline,
_binding_validated=True,
)
return
# Compatibility path for isolated manager tests and older embedders.
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_pipeline.LegacyPipeline))
pipelines = result.all()
# load pipelines
for pipeline in pipelines:
await self.load_pipeline(pipeline)
binding = await self.ap.workspace_service.get_execution_binding(pipeline.workspace_uuid)
await self.load_pipeline(
ExecutionContext(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
placement_generation=binding.placement_generation,
pipeline_uuid=pipeline.uuid,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
),
pipeline,
)
@staticmethod
def _normalize_execution_context(
context: ExecutionContext | RequestContext,
pipeline_uuid: str,
) -> ExecutionContext:
if isinstance(context, RequestContext):
return ExecutionContext.from_request(context, pipeline_uuid=pipeline_uuid)
if not isinstance(context, ExecutionContext):
raise WorkspaceRequiredError('Pipeline runtime operations require an ExecutionContext')
if not context.instance_uuid.strip() or not context.workspace_uuid.strip():
raise WorkspaceRequiredError('Pipeline runtime operations require an instance and Workspace')
if context.placement_generation <= 0:
raise WorkspaceRequiredError('Pipeline runtime operations require a positive placement generation')
if context.pipeline_uuid not in (None, pipeline_uuid):
raise WorkspaceRequiredError('Pipeline UUID does not match its ExecutionContext')
return dataclasses.replace(context, pipeline_uuid=pipeline_uuid)
async def load_pipeline(
self,
context: ExecutionContext | RequestContext,
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)
elif isinstance(pipeline_entity, dict):
pipeline_entity = persistence_pipeline.LegacyPipeline(**pipeline_entity)
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')
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,
getattr(self.ap, 'pipeline_config_meta_trigger', {'name': 'trigger', 'stages': []}),
@@ -454,17 +655,75 @@ class PipelineManager:
for stage_container in stage_containers:
await stage_container.inst.initialize(pipeline_entity.config)
runtime_pipeline = RuntimePipeline(self.ap, pipeline_entity, stage_containers)
self.pipelines.append(runtime_pipeline)
# 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,
)
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, uuid: str) -> RuntimePipeline | None:
for pipeline in self.pipelines:
if pipeline.pipeline_entity.uuid == uuid:
return pipeline
async def get_pipeline_by_uuid(
self,
context: ExecutionContext | RequestContext,
uuid: str,
) -> RuntimePipeline | None:
execution_context = self._normalize_execution_context(context, uuid)
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(self, uuid: str):
for pipeline in self.pipelines:
if pipeline.pipeline_entity.uuid == uuid:
self.pipelines.remove(pipeline)
return
async def remove_pipeline(
self,
context: ExecutionContext | RequestContext,
uuid: str,
) -> None:
execution_context = self._normalize_execution_context(context, uuid)
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
+369 -18
View File
@@ -1,57 +1,287 @@
from __future__ import annotations
import asyncio
import dataclasses
import inspect
import typing
import uuid
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.provider.session as provider_session
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.platform.events as platform_events
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]
QueryCounterKey = tuple[str, str, int]
SingletonContextResolver = typing.Callable[
[],
ExecutionContext | typing.Awaitable[ExecutionContext],
]
class ExecutionContextRequiredError(ValueError):
"""Raised when runtime work is created without a trusted Workspace scope."""
class ExecutionContextMismatchError(ValueError):
"""Raised when entity fields conflict with their trusted execution scope."""
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')
if not isinstance(execution_context.instance_uuid, str) or not execution_context.instance_uuid.strip():
raise ExecutionContextRequiredError('ExecutionContext.instance_uuid is required')
if not isinstance(execution_context.workspace_uuid, str) or not execution_context.workspace_uuid.strip():
raise ExecutionContextRequiredError('ExecutionContext.workspace_uuid is required')
if (
isinstance(execution_context.placement_generation, bool)
or not isinstance(execution_context.placement_generation, int)
or execution_context.placement_generation <= 0
):
raise ExecutionContextRequiredError('ExecutionContext.placement_generation must be a positive integer')
for field_name in ('bot_uuid', 'pipeline_uuid', 'query_uuid'):
value = getattr(execution_context, field_name)
if value is not None and (not isinstance(value, str) or not value.strip()):
raise ExecutionContextRequiredError(f'ExecutionContext.{field_name} must be a non-empty string when set')
def bind_execution_context(
execution_context: ExecutionContext,
*,
bot_uuid: str | None = None,
pipeline_uuid: str | None = None,
query_uuid: str | None = None,
) -> ExecutionContext:
"""Bind runtime entity identifiers without allowing scope substitution."""
_validate_execution_context(execution_context)
requested_fields = {
'bot_uuid': bot_uuid,
'pipeline_uuid': pipeline_uuid,
'query_uuid': query_uuid,
}
updates: dict[str, str] = {}
for field_name, requested_value in requested_fields.items():
if requested_value is None:
continue
if not isinstance(requested_value, str) or not requested_value.strip():
raise ExecutionContextRequiredError(f'{field_name} must be a non-empty string')
current_value = getattr(execution_context, field_name)
if current_value is not None and current_value != requested_value:
raise ExecutionContextMismatchError(f'ExecutionContext.{field_name} does not match the runtime entity')
if current_value is None:
updates[field_name] = requested_value
if not updates:
return execution_context
return dataclasses.replace(execution_context, **updates)
def get_query_execution_context(query: pipeline_query.Query) -> ExecutionContext:
"""Return and validate the trusted context attached to a Query."""
attached_context = getattr(query, '_execution_context', None)
bot_uuid = getattr(query, 'bot_uuid', None)
pipeline_uuid = getattr(query, 'pipeline_uuid', None)
query_uuid = getattr(query, 'query_uuid', None)
if isinstance(attached_context, ExecutionContext):
return bind_execution_context(
attached_context,
bot_uuid=bot_uuid,
pipeline_uuid=pipeline_uuid,
query_uuid=query_uuid,
)
raise ExecutionContextRequiredError('Query is missing its trusted ExecutionContext')
class QueryPool:
"""请求池，请求获得调度进入pipeline之前，保存在这里"""
query_id_counter: int = 0
"""Workspace-scoped queue of requests waiting for pipeline scheduling."""
query_id_counter: int
pool_lock: asyncio.Lock
queries: list[pipeline_query.Query]
cached_queries: dict[int, pipeline_query.Query]
"""Cached queries, used for plugin backward api call, will be removed after the query completely processed"""
cached_queries: dict[QueryCacheKey, pipeline_query.Query]
legacy_query_index: dict[LegacyQueryKey, str]
query_count_by_scope: dict[QueryCounterKey, int]
condition: asyncio.Condition
def __init__(self):
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,
execution_context: ExecutionContext | None,
*,
bot_uuid: str,
pipeline_uuid: str | None,
query_uuid: str | None = None,
) -> ExecutionContext:
"""Resolve an explicit scope or the opt-in OSS singleton scope."""
if execution_context is None:
if self._singleton_context_resolver is None:
raise ExecutionContextRequiredError('ExecutionContext is required; no singleton resolver is configured')
resolved_context = self._singleton_context_resolver()
if inspect.isawaitable(resolved_context):
resolved_context = await resolved_context
execution_context = resolved_context
return bind_execution_context(
execution_context,
bot_uuid=bot_uuid,
pipeline_uuid=pipeline_uuid,
query_uuid=query_uuid,
)
async def add_query(
self,
bot_uuid: str,
launcher_type: provider_session.LauncherTypes,
launcher_id: typing.Union[int, str],
sender_id: typing.Union[int, str],
launcher_id: int | str,
sender_id: int | str,
message_event: platform_events.MessageEvent,
message_chain: platform_message.MessageChain,
adapter: abstract_platform_adapter.AbstractMessagePlatformAdapter,
pipeline_uuid: typing.Optional[str] = None,
pipeline_uuid: str | None = None,
routed_by_rule: bool = False,
variables: typing.Optional[dict[str, typing.Any]] = None,
variables: dict[str, typing.Any] | None = None,
execution_context: ExecutionContext | None = None,
) -> pipeline_query.Query:
"""Create a query and cache it under an opaque, Workspace-scoped key."""
query_uuid = str(uuid.uuid4())
execution_context = await self.resolve_execution_context(
execution_context,
bot_uuid=bot_uuid,
pipeline_uuid=pipeline_uuid,
query_uuid=query_uuid,
)
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:
initial_variables.update(variables)
query = pipeline_query.Query(
instance_uuid=execution_context.instance_uuid,
workspace_uuid=execution_context.workspace_uuid,
placement_generation=execution_context.placement_generation,
bot_uuid=bot_uuid,
query_id=query_id,
query_uuid=query_uuid,
launcher_type=launcher_type,
launcher_id=launcher_id,
sender_id=sender_id,
@@ -63,12 +293,133 @@ class QueryPool:
adapter=adapter,
pipeline_uuid=pipeline_uuid,
)
# langbot-plugin 0.4.13 ignores these forward-compatible fields.
# Attach them explicitly until the Workspace-aware SDK is released.
object.__setattr__(query, 'instance_uuid', execution_context.instance_uuid)
object.__setattr__(query, 'workspace_uuid', execution_context.workspace_uuid)
object.__setattr__(
query,
'placement_generation',
execution_context.placement_generation,
)
object.__setattr__(query, 'query_uuid', query_uuid)
object.__setattr__(query, '_execution_context', execution_context)
self.queries.append(query)
self.cached_queries[query_id] = 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 = (
execution_context.instance_uuid,
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
def get_query_count(self, execution_context: ExecutionContext) -> int:
"""Return the lifetime query count for one active placement scope."""
_validate_execution_context(execution_context)
return self.query_count_by_scope.get(
(
execution_context.instance_uuid,
execution_context.workspace_uuid,
execution_context.placement_generation,
),
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,
query_uuid: str,
) -> pipeline_query.Query | None:
"""Return a query only from the explicitly selected Workspace."""
async with self.pool_lock:
return self.cached_queries.get((workspace_uuid, query_uuid))
async def require_query(
self,
workspace_uuid: str,
query_uuid: str,
) -> pipeline_query.Query:
"""Return a scoped query or raise without checking other Workspaces."""
query = await self.get_query(workspace_uuid, query_uuid)
if query is None:
raise QueryNotFoundError(f'Query {query_uuid!r} was not found in Workspace {workspace_uuid!r}')
return query
async def get_query_by_legacy_id(
self,
workspace_uuid: str,
query_id: int,
) -> pipeline_query.Query | None:
"""Resolve a legacy integer ID within one explicit Workspace."""
async with self.pool_lock:
query_uuid = self.legacy_query_index.get((workspace_uuid, query_id))
if query_uuid is None:
return None
return self.cached_queries.get((workspace_uuid, query_uuid))
async def remove_query(self, query: pipeline_query.Query) -> bool:
"""Remove a query and both of its Workspace-scoped indexes."""
execution_context = get_query_execution_context(query)
query_uuid = execution_context.query_uuid
if query_uuid is None:
raise ExecutionContextRequiredError('Query.query_uuid is required for removal')
async with self.pool_lock:
cache_key = (execution_context.workspace_uuid, query_uuid)
cached_query = self.cached_queries.get(cache_key)
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,
)
for index, queued_query in enumerate(self.queries):
if queued_query is query:
self.queries.pop(index)
break
plugin_diagnostics.discard_query_state(query)
return True
async def __aenter__(self):
await self.pool_lock.acquire()
return self
src/langbot/pkg/pipeline/preproc/preproc.py
+21 -5
View File
@@ -8,6 +8,7 @@ import langbot_plugin.api.entities.events as events
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.platform.events as platform_events
from ...pipeline.pool import get_query_execution_context
@stage.stage_class('PreProcessor')
@@ -70,7 +71,10 @@ class PreProcessor(stage.PipelineStage):
if primary_uuid:
try:
llm_model = await self.ap.model_mgr.get_model_by_uuid(primary_uuid)
llm_model = await self.ap.model_mgr.get_model_by_uuid(
get_query_execution_context(query),
primary_uuid,
)
except ValueError:
self.ap.logger.warning(f'LLM model {primary_uuid} not found or not configured')
@@ -79,7 +83,10 @@ class PreProcessor(stage.PipelineStage):
valid_fallbacks = []
for fb_uuid in fallback_uuids:
try:
await self.ap.model_mgr.get_model_by_uuid(fb_uuid)
await self.ap.model_mgr.get_model_by_uuid(
get_query_execution_context(query),
fb_uuid,
)
valid_fallbacks.append(fb_uuid)
except ValueError:
self.ap.logger.warning(f'Fallback model {fb_uuid} not found, skipping')
@@ -131,6 +138,7 @@ class PreProcessor(stage.PipelineStage):
bound_mcp_servers = query.variables.get('_pipeline_bound_mcp_servers', None)
include_mcp_resource_tools = query.variables.get('_pipeline_mcp_resource_agent_read_enabled', True)
all_tools = await self.ap.tool_mgr.get_all_tools(
get_query_execution_context(query),
bound_plugins,
bound_mcp_servers,
include_skill_authoring=include_skill_authoring,
@@ -149,6 +157,7 @@ class PreProcessor(stage.PipelineStage):
bound_mcp_servers = query.variables.get('_pipeline_bound_mcp_servers', None)
include_mcp_resource_tools = query.variables.get('_pipeline_mcp_resource_agent_read_enabled', True)
all_tools = await self.ap.tool_mgr.get_all_tools(
get_query_execution_context(query),
bound_plugins,
bound_mcp_servers,
include_skill_authoring=include_skill_authoring,
@@ -279,7 +288,13 @@ class PreProcessor(stage.PipelineStage):
# relied on this injection; without it the LLM never discovers
# the skills are there and just calls native tools instead.
if selected_runner == 'local-agent' and self.ap.skill_mgr:
pipeline_data = await self.ap.pipeline_service.get_pipeline(query.pipeline_uuid)
skill_execution_context = get_query_execution_context(query)
await self.ap.skill_mgr.ensure_loaded(skill_execution_context)
pipeline_data = await self.ap.pipeline_service.get_pipeline(
query.workspace_uuid,
query.pipeline_uuid,
include_secret=True,
)
extensions_prefs = (pipeline_data or {}).get('extensions_preferences', {})
enable_all_skills = extensions_prefs.get('enable_all_skills', True)
@@ -291,6 +306,7 @@ class PreProcessor(stage.PipelineStage):
query.variables['_pipeline_bound_skills'] = bound_skills
skill_addition = self.ap.skill_mgr.build_skill_aware_prompt_addition(
skill_execution_context,
bound_skills=bound_skills,
)
if skill_addition:
@@ -319,13 +335,13 @@ class PreProcessor(stage.PipelineStage):
f'Skill index injected into system prompt: '
f'pipeline={query.pipeline_uuid} '
f'bound_skills={bound_skills or "all"} '
f'loaded_skills={len(self.ap.skill_mgr.skills)}'
f'loaded_skills={len(self.ap.skill_mgr.get_skills(skill_execution_context))}'
)
else:
self.ap.logger.debug(
f'No skills available for prompt injection: '
f'pipeline={query.pipeline_uuid} '
f'loaded_skills={len(self.ap.skill_mgr.skills)} '
f'loaded_skills={len(self.ap.skill_mgr.get_skills(skill_execution_context))} '
f'bound_skills={bound_skills}'
)
src/langbot/pkg/pipeline/process/handlers/chat.py
+43 -1
View File
@@ -19,12 +19,32 @@ from ....provider import runners
import langbot_plugin.api.entities.builtin.provider.session as provider_session
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
from ...pool import get_query_execution_context
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,
@@ -86,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:
@@ -106,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()
@@ -118,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
@@ -144,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)
@@ -158,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}')
@@ -180,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()
@@ -198,7 +237,10 @@ class ChatMessageHandler(handler.MessageHandler):
model_name = None
try:
if runner_name == 'local-agent' and getattr(query, 'use_llm_model_uuid', None):
m = await self.ap.model_mgr.get_model_by_uuid(query.use_llm_model_uuid)
m = await self.ap.model_mgr.get_model_by_uuid(
get_query_execution_context(query),
query.use_llm_model_uuid,
)
if m and getattr(m, 'model_entity', None):
model_name = getattr(m.model_entity, 'name', None)
except Exception:
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)