feat(cloud): harden multi-tenant runtime resources

src/langbot/pkg/provider/modelmgr/modelmgr.py

@@ -52,6 +52,128 @@ class ModelManager:
         self.rerank_model_dict = {}
         self.requester_components = []
         self.requester_dict = {}
+        self._scope_generations: dict[tuple[str, str], int] = {}
+        self._provider_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
+        self._llm_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
+        self._embedding_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
+        self._rerank_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
+
+    def _cache_index(self, cache: dict) -> dict[tuple[str, str], set[_CacheKey]]:
+        if cache is self.provider_dict:
+            return self._provider_keys_by_scope
+        if cache is self.llm_model_dict:
+            return self._llm_keys_by_scope
+        if cache is self.embedding_model_dict:
+            return self._embedding_keys_by_scope
+        if cache is self.rerank_model_dict:
+            return self._rerank_keys_by_scope
+        raise ValueError('Unknown model runtime cache')
+
+    def _cache_set(self, cache: dict, key: _CacheKey, value: object) -> None:
+        cache[key] = value
+        self._cache_index(cache).setdefault(key[:2], set()).add(key)
+
+    def _cache_pop(self, cache: dict, key: _CacheKey) -> object | None:
+        removed = cache.pop(key, None)
+        scope = key[:2]
+        index = self._cache_index(cache)
+        keys = index.get(scope)
+        if keys is not None:
+            keys.discard(key)
+            if not keys:
+                index.pop(scope, None)
+        if not any(
+            scope in candidate
+            for candidate in (
+                self._provider_keys_by_scope,
+                self._llm_keys_by_scope,
+                self._embedding_keys_by_scope,
+                self._rerank_keys_by_scope,
+            )
+        ):
+            self._scope_generations.pop(scope, None)
+        return removed
+
+    def _observe_execution_context(
+        self,
+        context: ExecutionContext,
+    ) -> tuple[requester.RuntimeProvider, ...]:
+        """Prune superseded runtime objects when a Workspace generation advances."""
+
+        scope = (context.instance_uuid, context.workspace_uuid)
+        previous_generation = self._scope_generations.get(scope)
+        if previous_generation is not None and context.placement_generation < previous_generation:
+            raise WorkspaceInvariantError('Model runtime placement generation rolled back')
+        if previous_generation == context.placement_generation:
+            return ()
+        retired_providers: list[requester.RuntimeProvider] = []
+        if previous_generation is not None:
+            for cache, index in (
+                (self.provider_dict, self._provider_keys_by_scope),
+                (self.llm_model_dict, self._llm_keys_by_scope),
+                (self.embedding_model_dict, self._embedding_keys_by_scope),
+                (self.rerank_model_dict, self._rerank_keys_by_scope),
+            ):
+                for key in index.pop(scope, ()):
+                    removed = cache.pop(key, None)
+                    if cache is self.provider_dict and removed is not None:
+                        retired_providers.append(removed)
+        self._scope_generations[scope] = context.placement_generation
+        return tuple(retired_providers)
+
+    async def _close_runtime_providers(
+        self,
+        providers: tuple[requester.RuntimeProvider, ...] | list[requester.RuntimeProvider],
+    ) -> None:
+        """Close each retired requester once without blocking other cleanup."""
+
+        seen: set[int] = set()
+        for provider in providers:
+            provider_id = id(provider)
+            if provider_id in seen:
+                continue
+            seen.add(provider_id)
+            try:
+                await provider.requester.aclose()
+            except Exception as exc:
+                self.ap.logger.warning(
+                    f'Failed to close model requester for provider {provider.provider_entity.uuid}: {exc}'
+                )
+
+    async def _observe_and_close_execution_context(
+        self,
+        context: ExecutionContext,
+        *,
+        retain_empty: bool = True,
+    ) -> None:
+        await self._close_runtime_providers(self._observe_execution_context(context))
+        if not retain_empty:
+            scope = (context.instance_uuid, context.workspace_uuid)
+            if not any(
+                scope in candidate
+                for candidate in (
+                    self._provider_keys_by_scope,
+                    self._llm_keys_by_scope,
+                    self._embedding_keys_by_scope,
+                    self._rerank_keys_by_scope,
+                )
+            ):
+                self._scope_generations.pop(scope, None)
+
+    async def shutdown(self) -> None:
+        """Release every requester owned by the model runtime cache."""
+
+        providers = list(self.provider_dict.values())
+        self.provider_dict = {}
+        self.llm_model_dict = {}
+        self.embedding_model_dict = {}
+        self.rerank_model_dict = {}
+        self._scope_generations = {}
+        self._provider_keys_by_scope = {}
+        self._llm_keys_by_scope = {}
+        self._embedding_keys_by_scope = {}
+        self._rerank_keys_by_scope = {}
+        await self._close_runtime_providers(providers)
 
     @staticmethod
     def _get_litellm_provider_from_manifest(component: engine.Component | None) -> str | None:
@@ -137,7 +259,17 @@ class ModelManager:
         if supplied_instance_uuid is not None and supplied_instance_uuid != binding.instance_uuid:
             raise WorkspaceInvariantError('Runtime context belongs to another LangBot instance')
 
-        return self._context_from_binding(binding, trigger_principal=trigger_principal)
+        execution_context = self._context_from_binding(binding, trigger_principal=trigger_principal)
+        scope = (
+            execution_context.instance_uuid,
+            execution_context.workspace_uuid,
+        )
+        if scope in self._scope_generations:
+            await self._observe_and_close_execution_context(
+                execution_context,
+                retain_empty=False,
+            )
+        return execution_context
 
     async def initialize(self) -> None:
         self.requester_components = self.ap.discover.get_components_by_kind('LLMAPIRequester')
@@ -199,10 +331,16 @@ class ModelManager:
         """Load every active projected Workspace into isolated runtime caches."""
 
         self.ap.logger.info('Loading models from db...')
+        await self._close_runtime_providers(list(self.provider_dict.values()))
         self.provider_dict = {}
         self.llm_model_dict = {}
         self.embedding_model_dict = {}
         self.rerank_model_dict = {}
+        self._scope_generations = {}
+        self._provider_keys_by_scope = {}
+        self._llm_keys_by_scope = {}
+        self._embedding_keys_by_scope = {}
+        self._rerank_keys_by_scope = {}
 
         list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
         tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
@@ -233,6 +371,7 @@ class ModelManager:
                 binding,
                 trigger_principal=PrincipalContext(principal_type=PrincipalType.SYSTEM),
             )
+            await self._observe_and_close_execution_context(resolved)
             contexts[workspace_uuid] = resolved
             return resolved
 
@@ -243,7 +382,11 @@ class ModelManager:
             try:
                 context = await context_for(provider_entity.workspace_uuid)
                 runtime_provider = await self._build_provider(context, provider_entity)
-                self.provider_dict[self._cache_key(context, provider_entity.uuid)] = runtime_provider
+                self._cache_set(
+                    self.provider_dict,
+                    self._cache_key(context, provider_entity.uuid),
+                    runtime_provider,
+                )
             except provider_errors.RequesterNotFoundError as exc:
                 self.ap.logger.warning(
                     f'Requester {exc.requester_name} not found, skipping provider {provider_entity.uuid}'
@@ -282,7 +425,11 @@ class ModelManager:
                     )
                     continue
                 runtime_model = builder(context, model_entity, provider)
-                cache[self._cache_key(context, model_entity.uuid)] = runtime_model
+                self._cache_set(
+                    cache,
+                    self._cache_key(context, model_entity.uuid),
+                    runtime_model,
+                )
             except Exception as exc:
                 self.ap.logger.error(f'Failed to load model {model_entity.uuid}: {exc}\n{traceback.format_exc()}')
 
@@ -294,10 +441,20 @@ class ModelManager:
                 persistence_model.ModelProvider.workspace_uuid == context.workspace_uuid
             )
         )
-        for provider_entity in providers_result.all():
+        provider_entities = providers_result.all()
+        if provider_entities:
+            # Empty Workspaces are the dominant SaaS registration case. Do
+            # not retain one generation record per account until the
+            # Workspace owns an actual runtime model resource.
+            await self._observe_and_close_execution_context(context)
+        for provider_entity in provider_entities:
             try:
                 runtime_provider = await self._build_provider(context, provider_entity)
-                self.provider_dict[self._cache_key(context, provider_entity.uuid)] = runtime_provider
+                self._cache_set(
+                    self.provider_dict,
+                    self._cache_key(context, provider_entity.uuid),
+                    runtime_provider,
+                )
             except provider_errors.RequesterNotFoundError as exc:
                 self.ap.logger.warning(
                     f'Requester {exc.requester_name} not found, skipping provider {provider_entity.uuid}'
@@ -337,7 +494,11 @@ class ModelManager:
                     )
                     continue
                 runtime_model = builder(context, model_entity, provider)
-                cache[self._cache_key(context, model_entity.uuid)] = runtime_model
+                self._cache_set(
+                    cache,
+                    self._cache_key(context, model_entity.uuid),
+                    runtime_model,
+                )
             except Exception as exc:
                 self.ap.logger.error(f'Failed to load model {model_entity.uuid}: {exc}\n{traceback.format_exc()}')
 
@@ -562,7 +723,12 @@ class ModelManager:
         execution_context = await self.resolve_execution_context(context)
         self._ensure_same_scope(execution_context, provider.execution_context, resource='Provider')
         self._ensure_entity_workspace(provider.provider_entity, execution_context, resource='Provider')
-        self.provider_dict[self._cache_key(execution_context, provider.provider_entity.uuid)] = provider
+        self._observe_execution_context(execution_context)
+        self._cache_set(
+            self.provider_dict,
+            self._cache_key(execution_context, provider.provider_entity.uuid),
+            provider,
+        )
 
     async def get_provider_by_uuid(
         self,
@@ -578,7 +744,12 @@ class ModelManager:
 
     async def remove_provider(self, context: TenantContext, provider_uuid: str) -> None:
         execution_context = await self.resolve_execution_context(context)
-        self.provider_dict.pop(self._cache_key(execution_context, provider_uuid), None)
+        removed = self._cache_pop(
+            self.provider_dict,
+            self._cache_key(execution_context, provider_uuid),
+        )
+        if removed is not None:
+            await self._close_runtime_providers([removed])
 
     async def reload_provider(self, context: TenantContext, provider_uuid: str) -> None:
         execution_context = await self.resolve_execution_context(context)
@@ -593,12 +764,26 @@ class ModelManager:
             raise provider_errors.ProviderNotFoundError(provider_uuid)
 
         new_provider = await self._build_provider(execution_context, provider_entity)
-        cache_prefix = self._cache_key(execution_context, '')[:3]
-        for cache in (self.llm_model_dict, self.embedding_model_dict, self.rerank_model_dict):
-            for key, model in cache.items():
-                if key[:3] == cache_prefix and model.provider.provider_entity.uuid == provider_uuid:
+        scope = (execution_context.instance_uuid, execution_context.workspace_uuid)
+        for cache, index in (
+            (self.llm_model_dict, self._llm_keys_by_scope),
+            (self.embedding_model_dict, self._embedding_keys_by_scope),
+            (self.rerank_model_dict, self._rerank_keys_by_scope),
+        ):
+            for key in tuple(index.get(scope, ())):
+                model = cache.get(key)
+                if model is not None and model.provider.provider_entity.uuid == provider_uuid:
                     model.provider = new_provider
-        self.provider_dict[self._cache_key(execution_context, provider_uuid)] = new_provider
+        self._observe_execution_context(execution_context)
+        provider_key = self._cache_key(execution_context, provider_uuid)
+        old_provider = self.provider_dict.get(provider_key)
+        self._cache_set(
+            self.provider_dict,
+            provider_key,
+            new_provider,
+        )
+        if old_provider is not None and old_provider is not new_provider:
+            await self._close_runtime_providers([old_provider])
 
     @staticmethod
     def _coerce_model(model_info: _ModelEntity | sqlalchemy.Row, entity_type: type[_ModelEntity]) -> _ModelEntity:
@@ -689,7 +874,12 @@ class ModelManager:
     async def cache_llm_model(self, context: TenantContext, model: requester.RuntimeLLMModel) -> None:
         execution_context = await self.resolve_execution_context(context)
         self._ensure_same_scope(execution_context, model.execution_context, resource='LLM model')
-        self.llm_model_dict[self._cache_key(execution_context, model.model_entity.uuid)] = model
+        self._observe_execution_context(execution_context)
+        self._cache_set(
+            self.llm_model_dict,
+            self._cache_key(execution_context, model.model_entity.uuid),
+            model,
+        )
 
     async def cache_embedding_model(
         self,
@@ -698,12 +888,22 @@ class ModelManager:
     ) -> None:
         execution_context = await self.resolve_execution_context(context)
         self._ensure_same_scope(execution_context, model.execution_context, resource='Embedding model')
-        self.embedding_model_dict[self._cache_key(execution_context, model.model_entity.uuid)] = model
+        self._observe_execution_context(execution_context)
+        self._cache_set(
+            self.embedding_model_dict,
+            self._cache_key(execution_context, model.model_entity.uuid),
+            model,
+        )
 
     async def cache_rerank_model(self, context: TenantContext, model: requester.RuntimeRerankModel) -> None:
         execution_context = await self.resolve_execution_context(context)
         self._ensure_same_scope(execution_context, model.execution_context, resource='Rerank model')
-        self.rerank_model_dict[self._cache_key(execution_context, model.model_entity.uuid)] = model
+        self._observe_execution_context(execution_context)
+        self._cache_set(
+            self.rerank_model_dict,
+            self._cache_key(execution_context, model.model_entity.uuid),
+            model,
+        )
 
     async def get_model_by_uuid(self, context: TenantContext, model_uuid: str) -> requester.RuntimeLLMModel:
         execution_context = await self.resolve_execution_context(context)
@@ -739,15 +939,24 @@ class ModelManager:
 
     async def remove_llm_model(self, context: TenantContext, model_uuid: str) -> None:
         execution_context = await self.resolve_execution_context(context)
-        self.llm_model_dict.pop(self._cache_key(execution_context, model_uuid), None)
+        self._cache_pop(
+            self.llm_model_dict,
+            self._cache_key(execution_context, model_uuid),
+        )
 
     async def remove_embedding_model(self, context: TenantContext, model_uuid: str) -> None:
         execution_context = await self.resolve_execution_context(context)
-        self.embedding_model_dict.pop(self._cache_key(execution_context, model_uuid), None)
+        self._cache_pop(
+            self.embedding_model_dict,
+            self._cache_key(execution_context, model_uuid),
+        )
 
     async def remove_rerank_model(self, context: TenantContext, model_uuid: str) -> None:
         execution_context = await self.resolve_execution_context(context)
-        self.rerank_model_dict.pop(self._cache_key(execution_context, model_uuid), None)
+        self._cache_pop(
+            self.rerank_model_dict,
+            self._cache_key(execution_context, model_uuid),
+        )
 
     def get_available_requesters_info(self, model_type: str) -> list[dict]:
         if model_type:

src/langbot/pkg/provider/modelmgr/requester.py

@@ -463,6 +463,17 @@ class ProviderAPIRequester(metaclass=abc.ABCMeta):
     async def initialize(self):
         pass
 
+    async def aclose(self) -> None:
+        """Release requester-owned clients when its runtime provider retires.
+
+        Most built-in requesters are currently stateless, but provider
+        extensions may own connection pools or background resources. Keeping
+        the lifecycle hook on the base class lets Workspace generation changes
+        and application shutdown retire them deterministically.
+        """
+
+        return None
+
     async def scan_models(self, api_key: str | None = None) -> dict[str, typing.Any] | list[dict[str, typing.Any]]:
         """Scan models supported by the provider.
 

src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py

@@ -8,6 +8,7 @@ import litellm
 from litellm import acompletion, aembedding, arerank
 
 from .. import errors, requester
+from ....utils import httpclient
 import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
 import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
 import langbot_plugin.api.entities.builtin.provider.message as provider_message
@@ -955,14 +956,17 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
         rerank_url = f'{base_url}/rerank'
 
         try:
-            async with httpx.AsyncClient(timeout=timeout) as client:
+            async with httpx.AsyncClient(
+                timeout=timeout,
+                event_hooks=httpclient.httpx_response_limit_hooks(),
+            ) as client:
                 resp = await client.post(rerank_url, headers=headers, json=payload)
                 resp.raise_for_status()
-                data = resp.json()
+                data = await httpclient.parse_json_response(resp)
         except httpx.HTTPStatusError as e:
             body = ''
             try:
-                body = e.response.text
+                body = await httpclient.response_text(e.response)
             except Exception:
                 pass
             raise errors.RequesterError(f'rerank 请求失败 (HTTP {e.response.status_code}): {body or str(e)}')
@@ -998,10 +1002,14 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
         models_url = f'{base_url}/models'
 
         try:
-            async with httpx.AsyncClient(trust_env=True, timeout=timeout) as client:
+            async with httpx.AsyncClient(
+                trust_env=True,
+                timeout=timeout,
+                event_hooks=httpclient.httpx_response_limit_hooks(),
+            ) as client:
                 response = await client.get(models_url, headers=headers)
                 response.raise_for_status()
-                payload = response.json()
+                payload = await httpclient.parse_json_response(response)
 
             models = []
             for item in payload.get('data', []):

src/langbot/pkg/provider/runner.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+import asyncio
 import typing
 from typing import TYPE_CHECKING
 
@@ -11,6 +12,32 @@ if TYPE_CHECKING:
 
 
 preregistered_runners: list[typing.Type[RequestRunner]] = []
+_DEFAULT_SYNC_ITERATION_LIMIT = 100_000
+
+_T = typing.TypeVar('_T')
+
+
+def _next_sync(iterator: typing.Iterator[_T]) -> tuple[bool, _T | None]:
+    try:
+        return True, next(iterator)
+    except StopIteration:
+        return False, None
+
+
+async def iterate_sync(
+    iterable: typing.Iterable[_T],
+    *,
+    max_items: int = _DEFAULT_SYNC_ITERATION_LIMIT,
+) -> typing.AsyncGenerator[_T, None]:
+    """Consume a blocking SDK iterator without stalling the event loop."""
+
+    iterator = iter(iterable)
+    for _ in range(max(max_items, 1)):
+        has_item, item = await asyncio.to_thread(_next_sync, iterator)
+        if not has_item:
+            return
+        yield typing.cast(_T, item)
+    raise RuntimeError('Synchronous provider stream exceeded the event limit')
 
 
 def runner_class(name: str):
@@ -43,3 +70,6 @@ class RequestRunner(abc.ABC):
     ) -> typing.AsyncGenerator[provider_message.Message | provider_message.MessageChunk, None]:
         """运行请求"""
         pass
+
+    async def aclose(self) -> None:
+        """Release request-scoped resources after one runner invocation."""

src/langbot/pkg/provider/runners/cozeapi.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import typing
 import json
-import base64
 
 from langbot.pkg.provider import runner
 from langbot.pkg.core import app
@@ -11,6 +10,16 @@ from langbot.pkg.utils import image
 import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
 from langbot.libs.coze_server_api.client import AsyncCozeAPIClient
 
+_MAX_COZE_GENERATED_CHARS = 1024 * 1024
+_MAX_COZE_MEDIA_BYTES = 10 * 1024 * 1024
+
+
+def _append_bounded(current: str, addition: typing.Any) -> str:
+    addition = str(addition or '')
+    if len(current) + len(addition) > _MAX_COZE_GENERATED_CHARS:
+        raise ValueError('Coze response exceeds the runtime limit')
+    return current + addition
+
 
 @runner.runner_class('coze-api')
 class CozeAPIRunner(runner.RequestRunner):
@@ -77,7 +86,10 @@ class CozeAPIRunner(runner.RequestRunner):
                     content_parts.append({'type': 'text', 'text': ce.text})
                 elif ce.type == 'image_base64':
                     image_b64, image_format = await image.extract_b64_and_format(ce.image_base64)
-                    file_bytes = base64.b64decode(image_b64)
+                    file_bytes = await image.decode_base64_limited(
+                        image_b64,
+                        max_bytes=_MAX_COZE_MEDIA_BYTES,
+                    )
                     file_id = await self._get_file_id(file_bytes)
                     content_parts.append({'type': 'image', 'file_id': file_id})
                 elif ce.type == 'file':
@@ -144,7 +156,7 @@ class CozeAPIRunner(runner.RequestRunner):
                 auto_save_history=self.auto_save_history,
                 stream=True,
             ):
-                self.ap.logger.debug(f'coze-chat-stream: {chunk}')
+                self.ap.logger.debug(f'coze-chat-stream: {str(chunk)[:1000]}')
 
                 event_type = chunk.get('event')
                 data = chunk.get('data', {})
@@ -153,11 +165,17 @@ class CozeAPIRunner(runner.RequestRunner):
                 if event_type == 'conversation.message.delta':
                     # 收集内容
                     if 'content' in data:
-                        full_content += data.get('content', '')
+                        full_content = _append_bounded(
+                            full_content,
+                            data.get('content', ''),
+                        )
 
                     # 收集推理内容（如果有）
                     if 'reasoning_content' in data:
-                        full_reasoning += data.get('reasoning_content', '')
+                        full_reasoning = _append_bounded(
+                            full_reasoning,
+                            data.get('reasoning_content', ''),
+                        )
 
                 elif event_type.split('.')[-1] == 'done':  # 本地部署coze时，结束event不为done
                     # 保存会话ID
@@ -179,6 +197,8 @@ class CozeAPIRunner(runner.RequestRunner):
                 remove_think = self.pipeline_config.get('output', {}).get('misc', {}).get('remove-think', False)
                 if not remove_think:
                     content = f'<think>\n{full_reasoning}\n</think>\n{content}'.strip()
+                    if len(content) > _MAX_COZE_GENERATED_CHARS:
+                        raise ValueError('Coze response exceeds the runtime limit')
 
             # 一次性返回完整内容
             yield provider_message.Message(
@@ -227,7 +247,7 @@ class CozeAPIRunner(runner.RequestRunner):
                 auto_save_history=self.auto_save_history,
                 stream=True,
             ):
-                self.ap.logger.debug(f'coze-chat-stream-chunk: {chunk}')
+                self.ap.logger.debug(f'coze-chat-stream-chunk: {str(chunk)[:1000]}')
 
                 event_type = chunk.get('event')
                 data = chunk.get('data', {})
@@ -263,7 +283,7 @@ class CozeAPIRunner(runner.RequestRunner):
                     error_msg = f'Coze API错误: {data.get("message", "未知错误")}'
                     yield provider_message.MessageChunk(role='assistant', content=error_msg, finish_reason='error')
                     return
-                full_content += content
+                full_content = _append_bounded(full_content, content)
                 if message_idx % 8 == 0 or is_final:
                     if full_content:
                         yield provider_message.MessageChunk(role='assistant', content=full_content, is_final=is_final)
@@ -286,3 +306,6 @@ class CozeAPIRunner(runner.RequestRunner):
         else:
             async for msg in self._chat_messages(query):
                 yield msg
+
+    async def aclose(self) -> None:
+        await self.coze.close()

src/langbot/pkg/provider/runners/dashscopeapi.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import typing
 import re
 
@@ -10,6 +11,9 @@ from ...core import app
 import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
 import langbot_plugin.api.entities.builtin.provider.message as provider_message
 
+_MAX_DASHSCOPE_RESPONSE_CHARS = 1024 * 1024
+_MAX_DASHSCOPE_REFERENCES = 1024
+
 
 class DashscopeAPIError(Exception):
     """Dashscope API 请求失败"""
@@ -19,6 +23,13 @@ class DashscopeAPIError(Exception):
         super().__init__(self.message)
 
 
+def _append_bounded(current: str, addition: typing.Any) -> str:
+    addition = str(addition or '')
+    if len(current) + len(addition) > _MAX_DASHSCOPE_RESPONSE_CHARS:
+        raise DashscopeAPIError('Dashscope response exceeds the runtime limit')
+    return current + addition
+
+
 @runner.runner_class('dashscope-app-api')
 class DashScopeAPIRunner(runner.RequestRunner):
     "阿里云百炼DashsscopeAPI对话请求器"
@@ -111,18 +122,16 @@ class DashScopeAPIRunner(runner.RequestRunner):
         if remove_think:
             has_thoughts = False
         # 发送对话请求
-        response = dashscope.Application.call(
-            api_key=self.api_key,  # 智能体应用的API Key
-            app_id=self.app_id,  # 智能体应用的ID
-            prompt=plain_text,  # 用户输入的文本信息
-            stream=True,  # 流式输出
-            incremental_output=True,  # 增量输出，使用流式输出需要开启增量输出
-            session_id=query.session.using_conversation.uuid,  # 会话ID用于，多轮对话
+        response = await asyncio.to_thread(
+            dashscope.Application.call,
+            api_key=self.api_key,
+            app_id=self.app_id,
+            prompt=plain_text,
+            stream=True,
+            incremental_output=True,
+            session_id=query.session.using_conversation.uuid,
             enable_thinking=has_thoughts,
             has_thoughts=has_thoughts,
-            # rag_options={                                     # 主要用于文件交互，暂不支持
-            #     "session_file_ids": ["FILE_ID1"],             # FILE_ID1 替换为实际的临时文件ID,逗号隔开多个
-            # }
         )
         idx_chunk = 0
         try:
@@ -131,7 +140,7 @@ class DashScopeAPIRunner(runner.RequestRunner):
         except AttributeError:
             is_stream = False
         if is_stream:
-            for chunk in response:
+            async for chunk in runner.iterate_sync(response):
                 if chunk.get('status_code') != 200:
                     raise DashscopeAPIError(
                         f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '
@@ -145,15 +154,27 @@ class DashScopeAPIRunner(runner.RequestRunner):
                 if stream_think and stream_think[0].get('thought'):
                     if not think_start:
                         think_start = True
-                        pending_content += f'<think>\n{stream_think[0].get("thought")}'
+                        pending_content = _append_bounded(
+                            pending_content,
+                            f'<think>\n{stream_think[0].get("thought")}',
+                        )
                     else:
                         # 继续输出 reasoning_content
-                        pending_content += stream_think[0].get('thought')
+                        pending_content = _append_bounded(
+                            pending_content,
+                            stream_think[0].get('thought'),
+                        )
                 elif think_start and (not stream_think or stream_think[0].get('thought') == '') and not think_end:
                     think_end = True
-                    pending_content += '\n</think>\n'
+                    pending_content = _append_bounded(
+                        pending_content,
+                        '\n</think>\n',
+                    )
                 if stream_output.get('text') is not None:
-                    pending_content += stream_output.get('text')
+                    pending_content = _append_bounded(
+                        pending_content,
+                        stream_output.get('text'),
+                    )
                 # 是否是流式最后一个chunk
                 is_final = False if stream_output.get('finish_reason', False) == 'null' else True
 
@@ -162,12 +183,14 @@ class DashScopeAPIRunner(runner.RequestRunner):
 
                 # 从模型传出的参考资料信息中提取用于替换的字典
                 if references_dict_list is not None:
-                    for doc in references_dict_list:
+                    for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]:
                         if doc.get('index_id') is not None:
                             references_dict[doc.get('index_id')] = doc.get('doc_name')
 
                     # 将参考资料替换到文本中
                     pending_content = self._replace_references(pending_content, references_dict)
+                    if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS:
+                        raise DashscopeAPIError('Dashscope response exceeds the runtime limit')
 
                 if idx_chunk % 8 == 0 or is_final:
                     yield provider_message.MessageChunk(
@@ -178,7 +201,7 @@ class DashScopeAPIRunner(runner.RequestRunner):
             # 保存当前会话的session_id用于下次对话的语境
             query.session.using_conversation.uuid = stream_output.get('session_id')
         else:
-            for chunk in response:
+            async for chunk in runner.iterate_sync(response):
                 if chunk.get('status_code') != 200:
                     raise DashscopeAPIError(
                         f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '
@@ -192,15 +215,27 @@ class DashScopeAPIRunner(runner.RequestRunner):
                 if stream_think and stream_think[0].get('thought'):
                     if not think_start:
                         think_start = True
-                        pending_content += f'<think>\n{stream_think[0].get("thought")}'
+                        pending_content = _append_bounded(
+                            pending_content,
+                            f'<think>\n{stream_think[0].get("thought")}',
+                        )
                     else:
                         # 继续输出 reasoning_content
-                        pending_content += stream_think[0].get('thought')
+                        pending_content = _append_bounded(
+                            pending_content,
+                            stream_think[0].get('thought'),
+                        )
                 elif think_start and (not stream_think or stream_think[0].get('thought') == '') and not think_end:
                     think_end = True
-                    pending_content += '\n</think>\n'
+                    pending_content = _append_bounded(
+                        pending_content,
+                        '\n</think>\n',
+                    )
                 if stream_output.get('text') is not None:
-                    pending_content += stream_output.get('text')
+                    pending_content = _append_bounded(
+                        pending_content,
+                        stream_output.get('text'),
+                    )
 
             # 保存当前会话的session_id用于下次对话的语境
             query.session.using_conversation.uuid = stream_output.get('session_id')
@@ -210,12 +245,14 @@ class DashScopeAPIRunner(runner.RequestRunner):
 
             # 从模型传出的参考资料信息中提取用于替换的字典
             if references_dict_list is not None:
-                for doc in references_dict_list:
+                for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]:
                     if doc.get('index_id') is not None:
                         references_dict[doc.get('index_id')] = doc.get('doc_name')
 
                 # 将参考资料替换到文本中
-                pending_content = self._replace_references(pending_content, references_dict)
+            pending_content = self._replace_references(pending_content, references_dict)
+            if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS:
+                raise DashscopeAPIError('Dashscope response exceeds the runtime limit')
 
             yield provider_message.Message(
                 role='assistant',
@@ -240,18 +277,16 @@ class DashScopeAPIRunner(runner.RequestRunner):
         biz_params.update(query.variables)
 
         # 发送对话请求
-        response = dashscope.Application.call(
-            api_key=self.api_key,  # 智能体应用的API Key
-            app_id=self.app_id,  # 智能体应用的ID
-            prompt=plain_text,  # 用户输入的文本信息
-            stream=True,  # 流式输出
-            incremental_output=True,  # 增量输出，使用流式输出需要开启增量输出
-            session_id=query.session.using_conversation.uuid,  # 会话ID用于，多轮对话
-            biz_params=biz_params,  # 工作流应用的自定义输入参数传递
-            flow_stream_mode='message_format',  # 消息模式，输出/结束节点的流式结果
-            # rag_options={                                     # 主要用于文件交互，暂不支持
-            #     "session_file_ids": ["FILE_ID1"],             # FILE_ID1 替换为实际的临时文件ID,逗号隔开多个
-            # }
+        response = await asyncio.to_thread(
+            dashscope.Application.call,
+            api_key=self.api_key,
+            app_id=self.app_id,
+            prompt=plain_text,
+            stream=True,
+            incremental_output=True,
+            session_id=query.session.using_conversation.uuid,
+            biz_params=biz_params,
+            flow_stream_mode='message_format',
         )
 
         # 处理API返回的流式输出
@@ -262,7 +297,7 @@ class DashScopeAPIRunner(runner.RequestRunner):
             is_stream = False
         idx_chunk = 0
         if is_stream:
-            for chunk in response:
+            async for chunk in runner.iterate_sync(response):
                 if chunk.get('status_code') != 200:
                     raise DashscopeAPIError(
                         f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '
@@ -273,7 +308,10 @@ class DashScopeAPIRunner(runner.RequestRunner):
                 # 获取流式传输的output
                 stream_output = chunk.get('output', {})
                 if stream_output.get('workflow_message') is not None:
-                    pending_content += stream_output.get('workflow_message').get('message').get('content')
+                    pending_content = _append_bounded(
+                        pending_content,
+                        stream_output.get('workflow_message').get('message').get('content'),
+                    )
                 # if stream_output.get('text') is not None:
                 #     pending_content += stream_output.get('text')
 
@@ -284,12 +322,14 @@ class DashScopeAPIRunner(runner.RequestRunner):
 
                 # 从模型传出的参考资料信息中提取用于替换的字典
                 if references_dict_list is not None:
-                    for doc in references_dict_list:
+                    for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]:
                         if doc.get('index_id') is not None:
                             references_dict[doc.get('index_id')] = doc.get('doc_name')
 
                     # 将参考资料替换到文本中
                     pending_content = self._replace_references(pending_content, references_dict)
+                    if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS:
+                        raise DashscopeAPIError('Dashscope response exceeds the runtime limit')
                 if idx_chunk % 8 == 0 or is_final:
                     yield provider_message.MessageChunk(
                         role='assistant',
@@ -301,7 +341,7 @@ class DashScopeAPIRunner(runner.RequestRunner):
             query.session.using_conversation.uuid = stream_output.get('session_id')
 
         else:
-            for chunk in response:
+            async for chunk in runner.iterate_sync(response):
                 if chunk.get('status_code') != 200:
                     raise DashscopeAPIError(
                         f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '
@@ -312,7 +352,10 @@ class DashScopeAPIRunner(runner.RequestRunner):
                 # 获取流式传输的output
                 stream_output = chunk.get('output', {})
                 if stream_output.get('text') is not None:
-                    pending_content += stream_output.get('text')
+                    pending_content = _append_bounded(
+                        pending_content,
+                        stream_output.get('text'),
+                    )
 
                 is_final = False if stream_output.get('finish_reason', False) == 'null' else True
 
@@ -324,12 +367,14 @@ class DashScopeAPIRunner(runner.RequestRunner):
 
             # 从模型传出的参考资料信息中提取用于替换的字典
             if references_dict_list is not None:
-                for doc in references_dict_list:
+                for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]:
                     if doc.get('index_id') is not None:
                         references_dict[doc.get('index_id')] = doc.get('doc_name')
 
                 # 将参考资料替换到文本中
-                pending_content = self._replace_references(pending_content, references_dict)
+            pending_content = self._replace_references(pending_content, references_dict)
+            if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS:
+                raise DashscopeAPIError('Dashscope response exceeds the runtime limit')
 
             yield provider_message.Message(
                 role='assistant',

src/langbot/pkg/provider/runners/difysvapi.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+import asyncio
+import heapq
 import typing
 import json
 import time
 import uuid
-import base64
 import mimetypes
 import os
 import re
@@ -16,11 +17,9 @@ from langbot.pkg.provider import runner
 from langbot.pkg.core import app
 import langbot_plugin.api.entities.builtin.provider.message as provider_message
 import langbot_plugin.api.entities.builtin.platform.message as platform_message
-from langbot.pkg.utils import image
+from langbot.pkg.utils import httpclient, image
 import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
 from langbot.libs.dify_service_api.v1 import client, errors
-import httpx
-
 
 # Module-level store for paused-workflow form state. The key includes the full
 # execution scope before the bot, pipeline, adapter, and launcher dimensions;
@@ -29,8 +28,29 @@ import httpx
 # placement generations.
 PendingFormKey = tuple[str, str, int, str, str, str, str, str]
 _PENDING_FORMS: dict[PendingFormKey, 'OrderedDict[str, dict[str, typing.Any]]'] = {}
+_PENDING_FORM_EXPIRY_HEAP: list[tuple[float, int, PendingFormKey, str]] = []
+_PENDING_FORM_ACTIVE_COUNT = 0
+_PENDING_FORM_REVISION = 0
 _PENDING_FORM_DEFAULT_TTL = 30 * 60  # 30 minutes safety cap
+_PENDING_FORM_MAX_SESSIONS = 4096
+_PENDING_FORM_MAX_PER_SESSION = 16
+_PENDING_FORM_HEAP_COMPACT_FLOOR = 64
+_PENDING_FORM_HEAP_MAX_MULTIPLIER = 4
+_PENDING_FORM_REVISION_KEY = '_langbot_cache_revision'
 _STREAM_FORM_PLACEHOLDER = '\u200b'
+_MAX_DIFY_UPLOAD_BYTES = 10 * 1024 * 1024
+
+
+def _read_local_file_limited(path: str) -> bytes:
+    """Read a local platform attachment without allowing an oversized allocation."""
+
+    if os.path.getsize(path) > _MAX_DIFY_UPLOAD_BYTES:
+        raise ValueError('Dify upload file exceeds the size limit')
+    with open(path, 'rb') as file:
+        content = file.read(_MAX_DIFY_UPLOAD_BYTES + 1)
+    if len(content) > _MAX_DIFY_UPLOAD_BYTES:
+        raise ValueError('Dify upload file exceeds the size limit')
+    return content
 
 
 def _merge_stream_text(accumulated: str, incoming: typing.Any) -> str:
@@ -65,19 +85,100 @@ def _session_key_from_query(query: pipeline_query.Query) -> PendingFormKey:
     )
 
 
+def _synchronize_pending_form_cache_if_externally_cleared() -> None:
+    """Keep test/debug direct cache clears from retaining stale heap entries."""
+
+    global _PENDING_FORM_ACTIVE_COUNT
+    if _PENDING_FORMS:
+        return
+    _PENDING_FORM_EXPIRY_HEAP.clear()
+    _PENDING_FORM_ACTIVE_COUNT = 0
+
+
+def _pending_form_entry_is_current(
+    expires_at: float,
+    revision: int,
+    session_key: PendingFormKey,
+    form_token: str,
+) -> bool:
+    forms = _PENDING_FORMS.get(session_key)
+    if forms is None:
+        return False
+    stored = forms.get(form_token)
+    if stored is None:
+        return False
+    return stored.get(_PENDING_FORM_REVISION_KEY) == revision and stored.get('_expires_at') == expires_at
+
+
+def _peek_valid_pending_form_expiry(
+    *,
+    pop: bool = False,
+) -> tuple[float, int, PendingFormKey, str] | None:
+    while _PENDING_FORM_EXPIRY_HEAP:
+        entry = _PENDING_FORM_EXPIRY_HEAP[0]
+        if _pending_form_entry_is_current(*entry):
+            if pop:
+                heapq.heappop(_PENDING_FORM_EXPIRY_HEAP)
+            return entry
+        heapq.heappop(_PENDING_FORM_EXPIRY_HEAP)
+    return None
+
+
+def _drop_pending_form(session_key: PendingFormKey, form_token: str) -> None:
+    global _PENDING_FORM_ACTIVE_COUNT
+    forms = _PENDING_FORMS.get(session_key)
+    if forms is None or forms.pop(form_token, None) is None:
+        return
+    _PENDING_FORM_ACTIVE_COUNT = max(_PENDING_FORM_ACTIVE_COUNT - 1, 0)
+    if not forms:
+        _PENDING_FORMS.pop(session_key, None)
+
+
+def _drop_pending_form_session(session_key: PendingFormKey) -> None:
+    global _PENDING_FORM_ACTIVE_COUNT
+    forms = _PENDING_FORMS.pop(session_key, None)
+    if forms is not None:
+        _PENDING_FORM_ACTIVE_COUNT = max(
+            _PENDING_FORM_ACTIVE_COUNT - len(forms),
+            0,
+        )
+
+
+def _compact_pending_form_expiry_heap_if_needed() -> None:
+    max_heap_entries = max(
+        _PENDING_FORM_HEAP_COMPACT_FLOOR,
+        _PENDING_FORM_ACTIVE_COUNT * _PENDING_FORM_HEAP_MAX_MULTIPLIER,
+    )
+    if len(_PENDING_FORM_EXPIRY_HEAP) <= max_heap_entries:
+        return
+    _PENDING_FORM_EXPIRY_HEAP[:] = [
+        (
+            float(stored['_expires_at']),
+            int(stored[_PENDING_FORM_REVISION_KEY]),
+            session_key,
+            form_token,
+        )
+        for session_key, forms in _PENDING_FORMS.items()
+        for form_token, stored in forms.items()
+    ]
+    heapq.heapify(_PENDING_FORM_EXPIRY_HEAP)
+
+
 def _prune_pending_forms(now: float | None = None) -> None:
+    _synchronize_pending_form_cache_if_externally_cleared()
     if now is None:
         now = time.time()
-    for session_key in list(_PENDING_FORMS.keys()):
-        forms = _PENDING_FORMS[session_key]
-        expired_tokens = [token for token, data in forms.items() if data.get('_expires_at', 0) <= now]
-        for token in expired_tokens:
-            forms.pop(token, None)
-        if not forms:
-            _PENDING_FORMS.pop(session_key, None)
+    while True:
+        entry = _peek_valid_pending_form_expiry()
+        if entry is None or entry[0] > now:
+            break
+        _, _, session_key, form_token = _peek_valid_pending_form_expiry(pop=True)
+        _drop_pending_form(session_key, form_token)
+    _compact_pending_form_expiry_heap_if_needed()
 
 
 def _set_pending_form(session_key: PendingFormKey, form_data: dict[str, typing.Any]) -> None:
+    global _PENDING_FORM_ACTIVE_COUNT, _PENDING_FORM_REVISION
     _prune_pending_forms()
     if isinstance(session_key, tuple) and len(session_key) == 8:
         form_data['pipeline_uuid'] = session_key[4]
@@ -88,11 +189,31 @@ def _set_pending_form(session_key: PendingFormKey, form_data: dict[str, typing.A
     except (TypeError, ValueError):
         expiration_ts = 0.0
     stored['_expires_at'] = expiration_ts or (time.time() + _PENDING_FORM_DEFAULT_TTL)
+    _PENDING_FORM_REVISION += 1
+    stored[_PENDING_FORM_REVISION_KEY] = _PENDING_FORM_REVISION
     form_token = str(stored.get('form_token') or '')
     forms = _PENDING_FORMS.setdefault(session_key, OrderedDict())
     # Re-insert at the end so this becomes the "latest" entry
-    forms.pop(form_token, None)
+    if forms.pop(form_token, None) is None:
+        _PENDING_FORM_ACTIVE_COUNT += 1
     forms[form_token] = stored
+    heapq.heappush(
+        _PENDING_FORM_EXPIRY_HEAP,
+        (
+            stored['_expires_at'],
+            _PENDING_FORM_REVISION,
+            session_key,
+            form_token,
+        ),
+    )
+    while len(forms) > _PENDING_FORM_MAX_PER_SESSION:
+        oldest_token = next(iter(forms))
+        _drop_pending_form(session_key, oldest_token)
+    if len(_PENDING_FORMS) > _PENDING_FORM_MAX_SESSIONS:
+        oldest_entry = _peek_valid_pending_form_expiry()
+        if oldest_entry is not None:
+            _drop_pending_form_session(oldest_entry[2])
+    _compact_pending_form_expiry_heap_if_needed()
 
 
 def _get_pending_form_by_token(session_key: PendingFormKey, form_token: str) -> dict[str, typing.Any] | None:
@@ -144,11 +265,11 @@ def _clear_pending_form(session_key: PendingFormKey, form_token: str | None = No
     if not forms:
         return
     if form_token is None:
-        _PENDING_FORMS.pop(session_key, None)
+        _drop_pending_form_session(session_key)
+        _compact_pending_form_expiry_heap_if_needed()
         return
-    forms.pop(form_token, None)
-    if not forms:
-        _PENDING_FORMS.pop(session_key, None)
+    _drop_pending_form(session_key, form_token)
+    _compact_pending_form_expiry_heap_if_needed()
 
 
 def _format_human_input_text(
@@ -721,6 +842,9 @@ class DifyServiceAPIRunner(runner.RequestRunner):
             base_url=self.pipeline_config['ai']['dify-service-api']['base-url'],
         )
 
+    async def aclose(self) -> None:
+        await self.dify_client.aclose()
+
     def _process_thinking_content(
         self,
         content: str,
@@ -796,13 +920,16 @@ class DifyServiceAPIRunner(runner.RequestRunner):
         async def download_file(file_url: str) -> tuple[bytes, str]:
             """Download file from url (supports data url)."""
 
-            async with httpx.AsyncClient() as client_session:
-                resp = await client_session.get(file_url)
+            client_session = httpclient.get_session()
+            async with client_session.get(file_url, timeout=120) as resp:
                 resp.raise_for_status()
                 content_type = (
                     resp.headers.get('content-type') or mimetypes.guess_type(file_url)[0] or 'application/octet-stream'
                 )
-                return resp.content, content_type
+                return (
+                    await httpclient.read_limited(resp, max_bytes=_MAX_DIFY_UPLOAD_BYTES),
+                    content_type,
+                )
 
         def _detect_file_type(content_type: str) -> str:
             """Map MIME to dify file type."""
@@ -820,7 +947,10 @@ class DifyServiceAPIRunner(runner.RequestRunner):
                     plain_text += ce.text
                 elif ce.type == 'image_base64':
                     image_b64, image_format = await image.extract_b64_and_format(ce.image_base64)
-                    file_bytes = base64.b64decode(image_b64)
+                    file_bytes = await image.decode_base64_limited(
+                        image_b64,
+                        max_bytes=_MAX_DIFY_UPLOAD_BYTES,
+                    )
                     image_id = await upload_file_bytes(f'img.{image_format}', file_bytes, f'image/{image_format}')
                     upload_files.append({'type': 'image', 'id': image_id})
                 elif ce.type == 'file_url':
@@ -840,7 +970,10 @@ class DifyServiceAPIRunner(runner.RequestRunner):
                     content_type = 'application/octet-stream'
                     if ';' in header:
                         content_type = header.split(';')[0][5:] or content_type
-                    file_bytes = base64.b64decode(b64_data)
+                    file_bytes = await image.decode_base64_limited(
+                        b64_data,
+                        max_bytes=_MAX_DIFY_UPLOAD_BYTES,
+                    )
                     file_id = await upload_file_bytes(file_name, file_bytes, content_type)
                     file_type = _detect_file_type(content_type)
                     upload_files.append({'type': file_type, 'id': file_id})
@@ -865,15 +998,19 @@ class DifyServiceAPIRunner(runner.RequestRunner):
         }
 
     async def _download_file_for_form(self, file_url: str) -> tuple[bytes, str, str]:
-        async with httpx.AsyncClient() as client_session:
-            resp = await client_session.get(file_url)
+        client_session = httpclient.get_session()
+        async with client_session.get(file_url, timeout=120) as resp:
             resp.raise_for_status()
             content_type = (
                 resp.headers.get('content-type') or mimetypes.guess_type(file_url)[0] or 'application/octet-stream'
             )
             parsed = urlparse(file_url)
             file_name = os.path.basename(parsed.path) or 'file'
-            return resp.content, content_type, file_name
+            return (
+                await httpclient.read_limited(resp, max_bytes=_MAX_DIFY_UPLOAD_BYTES),
+                content_type,
+                file_name,
+            )
 
     async def _platform_file_to_dify(self, item: typing.Any, user: str) -> dict | None:
         try:
@@ -890,13 +1027,15 @@ class DifyServiceAPIRunner(runner.RequestRunner):
                         content_type = header.split(';', 1)[0][5:] or content_type
                     return await self._upload_file_bytes_for_user(
                         file_name,
-                        base64.b64decode(b64_data),
+                        await image.decode_base64_limited(
+                            b64_data,
+                            max_bytes=_MAX_DIFY_UPLOAD_BYTES,
+                        ),
                         content_type,
                         user,
                     )
                 if item.path:
-                    with open(item.path, 'rb') as f:
-                        file_bytes = f.read()
+                    file_bytes = await asyncio.to_thread(_read_local_file_limited, str(item.path))
                     content_type = mimetypes.guess_type(str(item.path))[0] or 'application/octet-stream'
                     file_name = item.name or os.path.basename(str(item.path)) or 'file'
                     return await self._upload_file_bytes_for_user(file_name, file_bytes, content_type, user)

src/langbot/pkg/provider/runners/langflowapi.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import codecs
 import typing
 import json
 import httpx
@@ -11,6 +12,44 @@ from ...core import app
 import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
 import langbot_plugin.api.entities.builtin.provider.message as provider_message
 
+_MAX_LANGFLOW_LINE_CHARS = 1024 * 1024
+_MAX_LANGFLOW_TOTAL_BYTES = 16 * 1024 * 1024
+_MAX_LANGFLOW_RESPONSE_BYTES = 1024 * 1024
+
+
+async def _iter_limited_lines(
+    response: httpx.Response,
+) -> typing.AsyncGenerator[str, None]:
+    decoder = codecs.getincrementaldecoder('utf-8')('replace')
+    buffer = ''
+    total_bytes = 0
+    async for chunk in response.aiter_bytes(chunk_size=8192):
+        total_bytes += len(chunk)
+        if total_bytes > _MAX_LANGFLOW_TOTAL_BYTES:
+            raise ValueError('Langflow stream exceeds the runtime limit')
+        buffer += decoder.decode(chunk)
+        while '\n' in buffer:
+            line, buffer = buffer.split('\n', 1)
+            if len(line) > _MAX_LANGFLOW_LINE_CHARS:
+                raise ValueError('Langflow event exceeds the runtime limit')
+            yield line.rstrip('\r')
+        if len(buffer) > _MAX_LANGFLOW_LINE_CHARS:
+            raise ValueError('Langflow event exceeds the runtime limit')
+    buffer += decoder.decode(b'', final=True)
+    if buffer:
+        if len(buffer) > _MAX_LANGFLOW_LINE_CHARS:
+            raise ValueError('Langflow event exceeds the runtime limit')
+        yield buffer.rstrip('\r')
+
+
+async def _read_limited_response(response: httpx.Response) -> bytes:
+    body = bytearray()
+    async for chunk in response.aiter_bytes(chunk_size=8192):
+        body.extend(chunk)
+        if len(body) > _MAX_LANGFLOW_RESPONSE_BYTES:
+            raise ValueError('Langflow response exceeds the runtime limit')
+    return bytes(body)
+
 
 @runner.runner_class('langflow-api')
 class LangflowAPIRunner(runner.RequestRunner):
@@ -99,7 +138,7 @@ class LangflowAPIRunner(runner.RequestRunner):
                     accumulated_content = ''
                     message_count = 0
 
-                    async for line in response.aiter_lines():
+                    async for line in _iter_limited_lines(response):
                         data_str = line
 
                         if data_str.startswith('data: '):
@@ -144,11 +183,15 @@ class LangflowAPIRunner(runner.RequestRunner):
                     yield provider_message.MessageChunk(role='assistant', content=accumulated_content, is_final=True)
             else:
                 # 非流式请求
-                response = await client.post(url, json=payload, headers=headers, timeout=120.0)
-                response.raise_for_status()
-
-                # 解析响应
-                response_data = response.json()
+                async with client.stream(
+                    'POST',
+                    url,
+                    json=payload,
+                    headers=headers,
+                    timeout=120.0,
+                ) as response:
+                    response.raise_for_status()
+                    response_data = json.loads(await _read_limited_response(response))
 
                 # 提取消息内容
                 # 根据Langflow API文档，响应结构可能在outputs[0].outputs[0].outputs.message.message中

src/langbot/pkg/provider/runners/n8nsvapi.py

@@ -12,6 +12,8 @@ from ...core import app
 import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
 import langbot_plugin.api.entities.builtin.provider.message as provider_message
 
+_MAX_N8N_RESPONSE_CHARS = 1024 * 1024
+
 
 class N8nAPIError(Exception):
     """N8n API 请求失败"""
@@ -94,6 +96,8 @@ class N8nServiceAPIRunner(runner.RequestRunner):
                 else:
                     chunk_str = str(raw_chunk)
 
+                if len(full_text) + len(chunk_str) > _MAX_N8N_RESPONSE_CHARS:
+                    raise N8nAPIError('n8n response exceeds the runtime limit')
                 full_text += chunk_str
                 buffer += chunk_str
 
@@ -112,7 +116,9 @@ class N8nServiceAPIRunner(runner.RequestRunner):
 
                         if obj.get('type') == 'item' and 'content' in obj:
                             chunk_idx += 1
-                            content = obj['content']
+                            content = str(obj['content'])
+                            if len(full_content) + len(content) > _MAX_N8N_RESPONSE_CHARS:
+                                raise N8nAPIError('n8n response exceeds the runtime limit')
                             full_content += content
                         elif obj.get('type') == 'end':
                             is_final = True
@@ -128,6 +134,8 @@ class N8nServiceAPIRunner(runner.RequestRunner):
                     except json.JSONDecodeError:
                         # buffer 末尾可能是一个不完整的 JSON，等待更多数据
                         break
+            except N8nAPIError:
+                raise
             except Exception as e:
                 # 记录解析失败并继续接收后续 chunk
                 try:
@@ -255,7 +263,12 @@ class N8nServiceAPIRunner(runner.RequestRunner):
                 self.webhook_url, json=payload, headers=headers, auth=auth, timeout=self.timeout
             ) as response:
                 if response.status != 200:
-                    error_text = await response.text()
+                    error_text = (
+                        await httpclient.read_limited(
+                            response,
+                            max_bytes=_MAX_N8N_RESPONSE_CHARS,
+                        )
+                    ).decode('utf-8', errors='replace')
                     self.ap.logger.error(f'n8n webhook call failed: {response.status}, {error_text}')
                     raise Exception(f'n8n webhook call failed: {response.status}, {error_text}')
 

src/langbot/pkg/provider/runners/tboxapi.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+import asyncio
 import typing
 import json
-import base64
+import logging
 import tempfile
 import os
 
@@ -11,10 +12,13 @@ from tboxsdk.model.file import File, FileType
 
 from .. import runner
 from ...core import app
-from ...utils import image
+from ...utils import bounded_executor, image
 import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
 import langbot_plugin.api.entities.builtin.provider.message as provider_message
 
+_MAX_TBOX_RESPONSE_CHARS = 1024 * 1024
+_MAX_TBOX_MEDIA_BYTES = 10 * 1024 * 1024
+
 
 class TboxAPIError(Exception):
     """TBox API 请求失败"""
@@ -24,6 +28,19 @@ class TboxAPIError(Exception):
         super().__init__(self.message)
 
 
+def _append_bounded(current: str, addition: typing.Any) -> str:
+    addition = str(addition or '')
+    if len(current) + len(addition) > _MAX_TBOX_RESPONSE_CHARS:
+        raise TboxAPIError('Tbox response exceeds the runtime limit')
+    return current + addition
+
+
+def _write_temp_media(file_bytes: bytes, suffix: str) -> str:
+    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp_file:
+        tmp_file.write(file_bytes)
+        return tmp_file.name
+
+
 @runner.runner_class('tbox-app-api')
 class TboxAPIRunner(runner.RequestRunner):
     "蚂蚁百宝箱API对话请求器"
@@ -42,6 +59,7 @@ class TboxAPIRunner(runner.RequestRunner):
         self.api_key = self.pipeline_config['ai']['tbox-app-api']['api-key']
 
         # 初始化Tbox client
+        logging.getLogger('tbox.client').setLevel(logging.WARNING)
         self.tbox_client = TboxClient(authorization=self.api_key)
 
     async def _preprocess_user_message(self, query: pipeline_query.Query) -> tuple[str, list[str]]:
@@ -59,19 +77,29 @@ class TboxAPIRunner(runner.RequestRunner):
                     plain_text += ce.text
                 elif ce.type == 'image_base64':
                     image_b64, image_format = await image.extract_b64_and_format(ce.image_base64)
-                    # 创建临时文件
-                    file_bytes = base64.b64decode(image_b64)
+                    file_bytes = await image.decode_base64_limited(
+                        image_b64,
+                        max_bytes=_MAX_TBOX_MEDIA_BYTES,
+                    )
+                    tmp_file_path: str | None = None
                     try:
-                        with tempfile.NamedTemporaryFile(suffix=f'.{image_format}', delete=False) as tmp_file:
-                            tmp_file.write(file_bytes)
-                            tmp_file_path = tmp_file.name
-                        file_upload_resp = self.tbox_client.upload_file(tmp_file_path)
+                        tmp_file_path = await asyncio.to_thread(
+                            _write_temp_media,
+                            file_bytes,
+                            f'.{image_format}',
+                        )
+                        file_upload_resp = await asyncio.to_thread(
+                            self.tbox_client.upload_file,
+                            tmp_file_path,
+                        )
                         image_id = file_upload_resp.get('data', '')
                         image_ids.append(image_id)
                     finally:
-                        # 清理临时文件
-                        if os.path.exists(tmp_file_path):
-                            os.unlink(tmp_file_path)
+                        if tmp_file_path and os.path.exists(tmp_file_path):
+                            await bounded_executor.run_blocking_cleanup(
+                                os.unlink,
+                                tmp_file_path,
+                            )
         elif isinstance(query.user_message.content, str):
             plain_text = query.user_message.content
 
@@ -98,18 +126,23 @@ class TboxAPIRunner(runner.RequestRunner):
             files = [File(file_id=image_id, type=FileType.IMAGE) for image_id in image_ids]
 
         # 发送对话请求
-        response = self.tbox_client.chat(
-            app_id=self.app_id,  # Tbox中智能体应用的ID
-            user_id=query.bot_uuid,  # 用户ID
-            query=plain_text,  # 用户输入的文本信息
-            stream=is_stream,  # 是否流式输出
-            conversation_id=conversation_id,  # 会话ID，为None时Tbox会自动创建一个新会话
-            files=files,  # 图片内容
+        response = await asyncio.to_thread(
+            self.tbox_client.chat,
+            app_id=self.app_id,
+            user_id=query.bot_uuid,
+            query=plain_text,
+            stream=is_stream,
+            conversation_id=conversation_id,
+            files=files,
         )
 
         if is_stream:
             # 解析Tbox流式输出内容，并发送给上游
-            for chunk in self._process_stream_message(response, query, remove_think):
+            async for chunk in self._process_stream_message(
+                response,
+                query,
+                remove_think,
+            ):
                 yield chunk
         else:
             message = self._process_non_stream_message(response, query, remove_think)
@@ -127,13 +160,16 @@ class TboxAPIRunner(runner.RequestRunner):
         thinking_content = payload.get('reasoningContent', [])
         result = ''
         if thinking_content and not remove_think:
-            result += f'<think>\n{thinking_content[0].get("text", "")}\n</think>\n'
+            result = _append_bounded(
+                result,
+                f'<think>\n{thinking_content[0].get("text", "")}\n</think>\n',
+            )
         content = payload.get('result', [])
         if content:
-            result += content[0].get('chunk', '')
+            result = _append_bounded(result, content[0].get('chunk', ''))
         return result
 
-    def _process_stream_message(
+    async def _process_stream_message(
         self, response: typing.Generator[dict], query: pipeline_query.Query, remove_think: bool
     ):
         idx_msg = 0
@@ -141,7 +177,7 @@ class TboxAPIRunner(runner.RequestRunner):
         conversation_id = None
         think_start = False
         think_end = False
-        for chunk in response:
+        async for chunk in runner.iterate_sync(response):
             if chunk.get('type', '') == 'chunk':
                 """
                 Tbox返回的消息内容chunk结构
@@ -149,7 +185,10 @@ class TboxAPIRunner(runner.RequestRunner):
                 """
                 # 如果包含思考过程，拼接</think>
                 if think_start and not think_end:
-                    pending_content += '\n</think>\n'
+                    pending_content = _append_bounded(
+                        pending_content,
+                        '\n</think>\n',
+                    )
                     think_end = True
 
                 payload = chunk.get('payload', {})
@@ -158,7 +197,10 @@ class TboxAPIRunner(runner.RequestRunner):
                     query.session.using_conversation.uuid = conversation_id
                 if payload.get('text'):
                     idx_msg += 1
-                    pending_content += payload.get('text')
+                    pending_content = _append_bounded(
+                        pending_content,
+                        payload.get('text'),
+                    )
             elif chunk.get('type', '') == 'thinking' and not remove_think:
                 """
                 Tbox返回的思考过程chunk结构
@@ -170,9 +212,15 @@ class TboxAPIRunner(runner.RequestRunner):
                     content = payload.get('ext_data', {}).get('text')
                     if not think_start:
                         think_start = True
-                        pending_content += f'<think>\n{content}'
+                        pending_content = _append_bounded(
+                            pending_content,
+                            f'<think>\n{content}',
+                        )
                     else:
-                        pending_content += content
+                        pending_content = _append_bounded(
+                            pending_content,
+                            content,
+                        )
             elif chunk.get('type', '') == 'error':
                 raise TboxAPIError(
                     f'Tbox API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '

src/langbot/pkg/provider/runners/weknoraapi.py

@@ -10,6 +10,15 @@ import langbot_plugin.api.entities.builtin.provider.message as provider_message
 import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
 from langbot.libs.weknora_api import client, errors
 
+_MAX_WEKNORA_GENERATED_CHARS = 1024 * 1024
+
+
+def _append_bounded(current: str, addition: typing.Any) -> str:
+    addition = str(addition or '')
+    if len(current) + len(addition) > _MAX_WEKNORA_GENERATED_CHARS:
+        raise errors.WeKnoraAPIError('WeKnora response exceeds the runtime limit')
+    return current + addition
+
 
 @runner.runner_class('weknora-api')
 class WeKnoraAPIRunner(runner.RequestRunner):
@@ -94,7 +103,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
             web_search_enabled=web_search_enabled,
             timeout=timeout,
         ):
-            self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk))
+            self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk)[:1000])
 
             response_type = chunk.get('response_type', '')
             content = chunk.get('content', '')
@@ -120,7 +129,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
 
             elif response_type == 'answer':
                 if content:
-                    full_answer += content
+                    full_answer = _append_bounded(full_answer, content)
 
             elif response_type == 'error':
                 raise errors.WeKnoraAPIError(f'WeKnora 服务错误: {content}')
@@ -158,14 +167,14 @@ class WeKnoraAPIRunner(runner.RequestRunner):
             knowledge_base_ids=knowledge_base_ids,
             timeout=timeout,
         ):
-            self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk))
+            self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk)[:1000])
 
             response_type = chunk.get('response_type', '')
             content = chunk.get('content', '')
 
             if response_type == 'answer':
                 if content:
-                    full_answer += content
+                    full_answer = _append_bounded(full_answer, content)
 
             elif response_type == 'error':
                 raise errors.WeKnoraAPIError(f'WeKnora 服务错误: {content}')
@@ -207,7 +216,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
             web_search_enabled=web_search_enabled,
             timeout=timeout,
         ):
-            self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk))
+            self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk)[:1000])
 
             response_type = chunk.get('response_type', '')
             content = chunk.get('content', '')
@@ -235,7 +244,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
             elif response_type == 'answer':
                 message_idx += 1
                 if content:
-                    pending_answer += content
+                    pending_answer = _append_bounded(pending_answer, content)
 
                 if done:
                     is_final = True
@@ -288,7 +297,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
             knowledge_base_ids=knowledge_base_ids,
             timeout=timeout,
         ):
-            self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk))
+            self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk)[:1000])
 
             response_type = chunk.get('response_type', '')
             content = chunk.get('content', '')
@@ -297,7 +306,7 @@ class WeKnoraAPIRunner(runner.RequestRunner):
             if response_type == 'answer':
                 message_idx += 1
                 if content:
-                    pending_answer += content
+                    pending_answer = _append_bounded(pending_answer, content)
 
                 if done:
                     is_final = True

src/langbot/pkg/provider/session/sessionmgr.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import asyncio
 import dataclasses
+import heapq
+import time
 
 from langbot_plugin.api.entities.builtin.provider import message as provider_message, prompt as provider_prompt
 import langbot_plugin.api.entities.builtin.provider.session as provider_session
@@ -24,6 +26,10 @@ SessionKey = tuple[
     str,
     int | str,
 ]
+SessionExpiryEntry = tuple[float, int, SessionKey]
+
+_SESSION_EXPIRY_HEAP_MIN_LIMIT = 64
+_SESSION_EXPIRY_HEAP_ACTIVE_MULTIPLIER = 4
 
 
 def _query_session_key(query: pipeline_query.Query) -> tuple[SessionKey, ExecutionContext]:
@@ -49,11 +55,215 @@ class SessionManager:
 
     ap: app.Application
 
-    session_list: list[provider_session.Session]
-
     def __init__(self, ap: app.Application):
         self.ap = ap
-        self.session_list = []
+        self._legacy_sessions: list[provider_session.Session] = []
+        self._session_index: dict[SessionKey, provider_session.Session] = {}
+        self._session_keys_by_workspace: dict[str, set[SessionKey]] = {}
+        self._session_expiry_heap: list[SessionExpiryEntry] = []
+        self._next_access_revision = 0
+
+    @property
+    def session_list(self) -> list[provider_session.Session]:
+        """Compatibility view for API services that enumerate sessions."""
+
+        return [
+            *self._legacy_sessions,
+            *self._session_index.values(),
+        ]
+
+    @session_list.setter
+    def session_list(self, sessions: list[provider_session.Session]) -> None:
+        """Replace the cache while keeping the O(1) index consistent."""
+
+        session_values = list(sessions)
+        self._legacy_sessions = []
+        self._session_index = {}
+        self._session_keys_by_workspace = {}
+        self._session_expiry_heap = []
+        self._next_access_revision = 0
+        now = time.monotonic()
+        for session in session_values:
+            key = getattr(session, '_langbot_session_key', None)
+            if isinstance(key, tuple) and len(key) == 6:
+                self._session_index[key] = session
+                self._session_keys_by_workspace.setdefault(key[1], set()).add(key)
+                last_accessed = getattr(session, '_langbot_last_accessed', None)
+                if last_accessed is None:
+                    last_accessed = now
+                self._touch_session(
+                    session,
+                    key,
+                    float(last_accessed),
+                    compact=False,
+                )
+            else:
+                self._legacy_sessions.append(session)
+        self._compact_session_expiry_heap(force=True)
+
+    def _retention_config(self) -> dict:
+        instance_config = getattr(getattr(self.ap, 'instance_config', None), 'data', {})
+        if not isinstance(instance_config, dict):
+            return {}
+        config = instance_config.get('system', {}).get('session_retention', {})
+        return config if isinstance(config, dict) else {}
+
+    def _positive_config_int(self, name: str, default: int) -> int:
+        try:
+            value = int(self._retention_config().get(name, default))
+        except (TypeError, ValueError):
+            value = default
+        return max(value, 1)
+
+    @staticmethod
+    def _session_is_idle(session: provider_session.Session) -> bool:
+        semaphore = getattr(session, '_semaphore', None)
+        concurrency = getattr(session, '_langbot_session_concurrency', None)
+        if semaphore is None or not isinstance(concurrency, int):
+            return True
+        return getattr(semaphore, '_value', -1) == concurrency
+
+    def _remove_session(self, session: provider_session.Session) -> None:
+        key = getattr(session, '_langbot_session_key', None)
+        if isinstance(key, tuple) and len(key) == 6 and self._session_index.get(key) is session:
+            self._session_index.pop(key, None)
+            workspace_keys = self._session_keys_by_workspace.get(key[1])
+            if workspace_keys is not None:
+                workspace_keys.discard(key)
+                if not workspace_keys:
+                    self._session_keys_by_workspace.pop(key[1], None)
+        else:
+            try:
+                self._legacy_sessions.remove(session)
+            except ValueError:
+                pass
+
+    def _touch_session(
+        self,
+        session: provider_session.Session,
+        key: SessionKey,
+        now: float,
+        *,
+        compact: bool = True,
+    ) -> None:
+        self._next_access_revision += 1
+        revision = self._next_access_revision
+        object.__setattr__(session, '_langbot_last_accessed', now)
+        object.__setattr__(session, '_langbot_access_revision', revision)
+        heapq.heappush(
+            self._session_expiry_heap,
+            (now, revision, key),
+        )
+        if compact:
+            self._compact_session_expiry_heap()
+
+    def _compact_session_expiry_heap(self, *, force: bool = False) -> None:
+        limit = max(
+            len(self._session_index) * _SESSION_EXPIRY_HEAP_ACTIVE_MULTIPLIER,
+            _SESSION_EXPIRY_HEAP_MIN_LIMIT,
+        )
+        if not force and len(self._session_expiry_heap) <= limit:
+            return
+        self._session_expiry_heap = [
+            (
+                float(getattr(session, '_langbot_last_accessed', 0.0)),
+                int(getattr(session, '_langbot_access_revision', 0)),
+                key,
+            )
+            for key, session in self._session_index.items()
+        ]
+        heapq.heapify(self._session_expiry_heap)
+
+    def _pop_current_expiry_entry(
+        self,
+    ) -> tuple[float, int, SessionKey, provider_session.Session] | None:
+        while self._session_expiry_heap:
+            last_accessed, revision, key = heapq.heappop(self._session_expiry_heap)
+            session = self._session_index.get(key)
+            if session is None:
+                continue
+            if getattr(session, '_langbot_access_revision', None) != revision:
+                continue
+            return last_accessed, revision, key, session
+        return None
+
+    def _prune_expired_sessions(self, now: float) -> None:
+        idle_ttl = self._positive_config_int('idle_ttl_seconds', 86400)
+        cutoff = now - idle_ttl
+        while self._session_expiry_heap:
+            last_accessed, _, _ = self._session_expiry_heap[0]
+            if last_accessed > cutoff:
+                break
+            current = self._pop_current_expiry_entry()
+            if current is None:
+                break
+            last_accessed, revision, key, session = current
+            if last_accessed > cutoff:
+                heapq.heappush(
+                    self._session_expiry_heap,
+                    (last_accessed, revision, key),
+                )
+                break
+            if self._session_is_idle(session):
+                self._remove_session(session)
+                continue
+            # The session became active without another cache lookup. Give it
+            # a fresh TTL instead of repeatedly examining the same expired
+            # entry or losing its future expiry record.
+            self._touch_session(session, key, now)
+
+    def _prune_workspace_capacity(
+        self,
+        workspace_uuid: str,
+        max_entries_per_workspace: int,
+    ) -> None:
+        workspace_keys = self._session_keys_by_workspace.get(workspace_uuid, set())
+        overflow = len(workspace_keys) - max_entries_per_workspace + 1
+        if overflow <= 0:
+            return
+        idle_workspace_sessions = sorted(
+            (
+                session
+                for key in tuple(workspace_keys)
+                if (session := self._session_index.get(key)) is not None and self._session_is_idle(session)
+            ),
+            key=lambda session: float(getattr(session, '_langbot_last_accessed', 0.0)),
+        )
+        for session in idle_workspace_sessions[:overflow]:
+            self._remove_session(session)
+
+    def _evict_oldest_idle_session(self, now: float) -> bool:
+        # At most one current entry per active session is examined. Stale heap
+        # revisions do not count and are discarded in O(log N).
+        current_probes = 0
+        max_probes = len(self._session_index)
+        while current_probes < max_probes:
+            current = self._pop_current_expiry_entry()
+            if current is None:
+                return False
+            _, _, key, session = current
+            current_probes += 1
+            if self._session_is_idle(session):
+                self._remove_session(session)
+                return True
+            self._touch_session(session, key, now)
+        return False
+
+    def _prune_sessions(self, now: float, workspace_uuid: str) -> None:
+        self._prune_expired_sessions(now)
+        max_entries_per_workspace = self._positive_config_int('max_entries_per_workspace', 200)
+        self._prune_workspace_capacity(
+            workspace_uuid,
+            max_entries_per_workspace,
+        )
+
+        max_entries = self._positive_config_int('max_entries', 2000)
+        overflow = len(self._session_index) - max_entries + 1
+        if overflow <= 0:
+            return
+        for _ in range(overflow):
+            if not self._evict_oldest_idle_session(now):
+                break
 
     async def initialize(self):
         pass
@@ -61,9 +271,25 @@ class SessionManager:
     async def get_session(self, query: pipeline_query.Query) -> provider_session.Session:
         """获取会话"""
         session_key, execution_context = _query_session_key(query)
-        for session in self.session_list:
-            if getattr(session, '_langbot_session_key', None) == session_key:
-                return session
+        now = time.monotonic()
+        session = self._session_index.get(session_key)
+        if session is not None:
+            self._touch_session(session, session_key, now)
+            return session
+
+        self._prune_sessions(now, execution_context.workspace_uuid)
+        max_entries_per_workspace = self._positive_config_int('max_entries_per_workspace', 200)
+        workspace_entries = len(
+            self._session_keys_by_workspace.get(
+                execution_context.workspace_uuid,
+                (),
+            )
+        )
+        if workspace_entries >= max_entries_per_workspace:
+            raise RuntimeError(f'Workspace session cache capacity reached ({max_entries_per_workspace})')
+        max_entries = self._positive_config_int('max_entries', 2000)
+        if len(self._session_index) >= max_entries:
+            raise RuntimeError(f'Session cache capacity reached ({max_entries})')
 
         session_concurrency = self.ap.instance_config.data['concurrency']['session']
 
@@ -93,8 +319,14 @@ class SessionManager:
         object.__setattr__(session, 'bot_uuid', query.bot_uuid)
         object.__setattr__(session, '_execution_context', session_context)
         object.__setattr__(session, '_langbot_session_key', session_key)
+        object.__setattr__(session, '_langbot_session_concurrency', session_concurrency)
         session._semaphore = asyncio.Semaphore(session_concurrency)
-        self.session_list.append(session)
+        self._session_index[session_key] = session
+        self._session_keys_by_workspace.setdefault(
+            execution_context.workspace_uuid,
+            set(),
+        ).add(session_key)
+        self._touch_session(session, session_key, now)
         return session
 
     async def get_conversation(
@@ -144,6 +376,47 @@ class SessionManager:
                 bot_uuid=bot_uuid,
             )
             session.conversations.append(conversation)
+            max_conversations = self._positive_config_int('max_conversations_per_session', 20)
+            if len(session.conversations) > max_conversations:
+                del session.conversations[:-max_conversations]
             session.using_conversation = conversation
 
         return session.using_conversation
+
+    def trim_conversation_messages(
+        self,
+        conversation: provider_session.Conversation,
+        *,
+        max_rounds: int,
+    ) -> None:
+        """Bound retained process-local history after a completed turn."""
+
+        try:
+            max_rounds = int(max_rounds)
+        except (TypeError, ValueError):
+            max_rounds = 10
+        max_rounds = max(max_rounds, 1)
+        max_messages = self._positive_config_int('max_messages_per_conversation', 100)
+
+        kept_reversed = []
+        user_rounds = 0
+        for message in reversed(conversation.messages):
+            if user_rounds >= max_rounds:
+                break
+            kept_reversed.append(message)
+            if getattr(message, 'role', None) == 'user':
+                user_rounds += 1
+        retained = list(reversed(kept_reversed))[-max_messages:]
+        # Binary payloads are needed for the current model call, but retaining
+        # them in process-local history makes a few image/file turns consume
+        # hundreds of MB. Historical URL and text references remain intact.
+        for message in retained:
+            content = getattr(message, 'content', None)
+            if not isinstance(content, list):
+                continue
+            for element in content:
+                if getattr(element, 'image_base64', None) is not None:
+                    element.image_base64 = None
+                if getattr(element, 'file_base64', None) is not None:
+                    element.file_base64 = None
+        conversation.messages = retained

src/langbot/pkg/provider/tools/loaders/mcp.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
-import base64
 import enum
 import json
 import math
 import re
 import time
 import typing
+import ipaddress
+from urllib.parse import urlparse
 from contextlib import AsyncExitStack, asynccontextmanager
 from datetime import timedelta
 import traceback
@@ -50,6 +51,7 @@ MCP_TOOL_READ_RESOURCE = 'langbot_mcp_read_resource'
 
 MCP_RESOURCE_DISCOVERY_MAX_PAGES = 20
 MCP_RESOURCE_CACHE_TTL_SECONDS = 30
+MCP_RESOURCE_CACHE_MAX_ENTRIES = 32
 MCP_RESOURCE_PREVIEW_MAX_BYTES = 64 * 1024
 MCP_RESOURCE_AGENT_READ_MAX_BYTES = 64 * 1024
 MCP_RESOURCE_AGENT_READ_MAX_TOKENS = 12000
@@ -134,10 +136,13 @@ def _truncate_text(text: str, max_bytes: int, max_tokens: int | None = None) ->
 
 
 def _blob_size(blob: str) -> int:
-    try:
-        return len(base64.b64decode(blob, validate=False))
-    except Exception:
+    # MCP BlobResourceContents is schema-validated base64 without whitespace.
+    # Compute decoded size in O(1) without allocating a second binary copy.
+    encoded_chars = len(blob)
+    if encoded_chars % 4:
         return len(blob.encode('utf-8', errors='ignore'))
+    padding = 2 if blob.endswith('==') else 1 if blob.endswith('=') else 0
+    return max((encoded_chars // 4) * 3 - padding, 0)
 
 
 def _resource_to_dict(resource: mcp_types.Resource | mcp_types.ResourceLink) -> dict:
@@ -434,12 +439,24 @@ class RuntimeMCPSession:
         await self._box_stdio_runtime.initialize()
 
     async def _init_sse_server(self):
+        trust_env = self._remote_http_trust_env()
+
+        def httpx_client_factory(headers=None, timeout=None, auth=None):
+            return httpx.AsyncClient(
+                headers=headers,
+                timeout=timeout,
+                auth=auth,
+                follow_redirects=True,
+                trust_env=trust_env,
+            )
+
         sse_transport = await self.exit_stack.enter_async_context(
             sse_client(
                 self.server_config['url'],
                 headers=self.server_config.get('headers', {}),
                 timeout=self.server_config.get('timeout', 10),
                 sse_read_timeout=self.server_config.get('ssereadtimeout', 30),
+                httpx_client_factory=httpx_client_factory,
             )
         )
 
@@ -449,6 +466,18 @@ class RuntimeMCPSession:
 
         await self.session.initialize()
 
+    def _remote_http_trust_env(self) -> bool:
+        configured = self.server_config.get('trust_env')
+        if isinstance(configured, bool):
+            return configured
+        hostname = (urlparse(str(self.server_config.get('url', ''))).hostname or '').lower()
+        if hostname == 'localhost':
+            return False
+        try:
+            return not ipaddress.ip_address(hostname).is_loopback
+        except ValueError:
+            return True
+
     @asynccontextmanager
     async def _streamable_http_session(self) -> typing.AsyncIterator[ClientSession]:
         """Enter a fully initialized Streamable HTTP session as one context.
@@ -465,6 +494,7 @@ class RuntimeMCPSession:
             headers=self.server_config.get('headers', {}),
             timeout=self.server_config.get('timeout', 10),
             follow_redirects=True,
+            trust_env=self._remote_http_trust_env(),
         ) as http_client:
             async with streamable_http_client(
                 self.server_config['url'],
@@ -1215,6 +1245,9 @@ class RuntimeMCPSession:
 
         cache_key = (uri, max_bytes, max_tokens, include_blob)
         now = time.time()
+        for expired_key, entry in tuple(self._resource_cache.items()):
+            if now - entry.get('cached_at', 0) > MCP_RESOURCE_CACHE_TTL_SECONDS:
+                self._resource_cache.pop(expired_key, None)
         cached = self._resource_cache.get(cache_key)
         if cached and now - cached.get('cached_at', 0) <= MCP_RESOURCE_CACHE_TTL_SECONDS:
             envelope = {
@@ -1314,6 +1347,12 @@ class RuntimeMCPSession:
             'warnings': warnings,
         }
         await self._assert_execution_active()
+        if cache_key not in self._resource_cache and len(self._resource_cache) >= MCP_RESOURCE_CACHE_MAX_ENTRIES:
+            oldest_key = min(
+                self._resource_cache,
+                key=lambda key: self._resource_cache[key].get('cached_at', 0),
+            )
+            self._resource_cache.pop(oldest_key, None)
         self._resource_cache[cache_key] = {'cached_at': now, 'envelope': envelope}
         self._record_resource_read_trace(query, envelope)
         return envelope
@@ -1498,17 +1537,249 @@ class MCPLoader(loader.ToolLoader):
     在此加载器中管理所有与 MCP Server 的连接。
     """
 
-    sessions: dict[tuple[str, str, int, str], RuntimeMCPSession]
-
-    _last_listed_functions: list[resource_tool.LLMTool]
+    _sessions: dict[tuple[str, str, int, str], RuntimeMCPSession]
 
     _hosted_mcp_tasks: list[asyncio.Task]
 
     def __init__(self, ap: app.Application):
         super().__init__(ap)
         self.sessions = {}
-        self._last_listed_functions = []
         self._hosted_mcp_tasks = []
+        self._hosted_mcp_tasks_by_scope: dict[
+            tuple[str, str, int],
+            set[asyncio.Task],
+        ] = {}
+        self._host_dispatch_tasks: set[asyncio.Task] = set()
+        config = getattr(getattr(ap, 'instance_config', None), 'data', {})
+        mcp_config = config.get('mcp', {}) if isinstance(config, dict) else {}
+        raw_lifecycle_concurrency = mcp_config.get('lifecycle_concurrency', 16) if isinstance(mcp_config, dict) else 16
+        if (
+            isinstance(raw_lifecycle_concurrency, bool)
+            or not isinstance(raw_lifecycle_concurrency, int)
+            or raw_lifecycle_concurrency < 1
+        ):
+            raw_lifecycle_concurrency = 16
+        self._lifecycle_concurrency = min(
+            raw_lifecycle_concurrency,
+            128,
+        )
+        self._lifecycle_semaphore = asyncio.Semaphore(self._lifecycle_concurrency)
+
+    @property
+    def sessions(
+        self,
+    ) -> dict[tuple[str, str, int, str], RuntimeMCPSession]:
+        return self._sessions
+
+    @sessions.setter
+    def sessions(self, sessions: dict) -> None:
+        """Compatibility setter that rebuilds the per-scope session index."""
+
+        self._sessions = sessions
+        self._session_keys_by_scope: dict[
+            tuple[str, str, int],
+            set[tuple[str, str, int, str]],
+        ] = {}
+        self._scope_generations: dict[tuple[str, str], int] = {}
+        for key in sessions:
+            if not isinstance(key, tuple) or len(key) != 4:
+                continue
+            scope_key = key[:3]
+            self._session_keys_by_scope.setdefault(scope_key, set()).add(key)
+            self._scope_generations[scope_key[:2]] = scope_key[2]
+
+    def _register_session(
+        self,
+        context: TenantContext,
+        server_name: str,
+        session: RuntimeMCPSession,
+    ) -> None:
+        scope_key = self._scope_key(context)
+        workspace_scope = scope_key[:2]
+        previous_generation = self._scope_generations.get(workspace_scope)
+        if previous_generation is not None and previous_generation != scope_key[2]:
+            raise WorkspaceInvariantError('MCP session registration crossed a Workspace generation')
+        key = (*scope_key, server_name)
+        self._sessions[key] = session
+        self._session_keys_by_scope.setdefault(scope_key, set()).add(key)
+        self._scope_generations[workspace_scope] = scope_key[2]
+
+    def _pop_session(
+        self,
+        context: TenantContext,
+        server_name: str,
+    ) -> RuntimeMCPSession | None:
+        scope_key = self._scope_key(context)
+        key = (*scope_key, server_name)
+        session = self._sessions.pop(key, None)
+        keys = self._session_keys_by_scope.get(scope_key)
+        if keys is not None:
+            keys.discard(key)
+            if not keys:
+                self._session_keys_by_scope.pop(scope_key, None)
+        self._drop_empty_scope(scope_key)
+        return session
+
+    def _drop_empty_scope(self, scope_key: tuple[str, str, int]) -> None:
+        if (
+            scope_key not in self._session_keys_by_scope
+            and scope_key not in self._hosted_mcp_tasks_by_scope
+            and self._scope_generations.get(scope_key[:2]) == scope_key[2]
+        ):
+            self._scope_generations.pop(scope_key[:2], None)
+
+    def track_hosted_task(
+        self,
+        task: asyncio.Task,
+        context: TenantContext,
+    ) -> asyncio.Task:
+        """Track a host task without retaining it after completion."""
+
+        scope_key = self._scope_key(context)
+        workspace_scope = scope_key[:2]
+        previous_generation = self._scope_generations.get(workspace_scope)
+        if previous_generation is not None and previous_generation != scope_key[2]:
+            task.cancel()
+            raise WorkspaceInvariantError('MCP host task crossed a Workspace generation')
+        self._scope_generations[workspace_scope] = scope_key[2]
+        self._hosted_mcp_tasks.append(task)
+        self._hosted_mcp_tasks_by_scope.setdefault(scope_key, set()).add(task)
+
+        def discard(completed: asyncio.Task) -> None:
+            try:
+                self._hosted_mcp_tasks.remove(completed)
+            except ValueError:
+                pass
+            tasks = self._hosted_mcp_tasks_by_scope.get(scope_key)
+            if tasks is not None:
+                tasks.discard(completed)
+                if not tasks:
+                    self._hosted_mcp_tasks_by_scope.pop(scope_key, None)
+            self._drop_empty_scope(scope_key)
+
+        task.add_done_callback(discard)
+        return task
+
+    def _track_host_dispatch_task(self, task: asyncio.Task) -> None:
+        """Track the bounded startup dispatcher without retaining it."""
+
+        self._host_dispatch_tasks.add(task)
+
+        def discard(completed: asyncio.Task) -> None:
+            self._host_dispatch_tasks.discard(completed)
+            if completed.cancelled():
+                return
+            exception = completed.exception()
+            if exception is not None:
+                self.ap.logger.error(
+                    f'MCP startup dispatcher failed: {exception}',
+                )
+
+        task.add_done_callback(discard)
+
+    async def _retire_runtime_scope(
+        self,
+        scope_key: tuple[str, str, int],
+    ) -> None:
+        tasks = tuple(self._hosted_mcp_tasks_by_scope.pop(scope_key, ()))
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+        keys = tuple(self._session_keys_by_scope.pop(scope_key, ()))
+        sessions = [session for key in keys if (session := self._sessions.pop(key, None)) is not None]
+        await self._shutdown_sessions(sessions)
+        self._scope_generations.pop(scope_key[:2], None)
+
+    async def _observe_execution_context(
+        self,
+        context: ExecutionContext,
+    ) -> None:
+        workspace_scope = (
+            context.instance_uuid,
+            context.workspace_uuid,
+        )
+        previous_generation = self._scope_generations.get(workspace_scope)
+        if previous_generation is None:
+            return
+        if context.placement_generation < previous_generation:
+            raise WorkspaceInvariantError('MCP runtime placement generation rolled back')
+        if context.placement_generation == previous_generation:
+            return
+        await self._retire_runtime_scope((*workspace_scope, previous_generation))
+
+    async def _reset_runtime_state(self) -> None:
+        """Cancel host tasks and close sessions before reload or shutdown."""
+
+        dispatch_tasks = tuple(self._host_dispatch_tasks)
+        self._host_dispatch_tasks.clear()
+        for task in dispatch_tasks:
+            if not task.done():
+                task.cancel()
+        if dispatch_tasks:
+            await asyncio.gather(*dispatch_tasks, return_exceptions=True)
+
+        tasks = tuple(self._hosted_mcp_tasks)
+        self._hosted_mcp_tasks.clear()
+        self._hosted_mcp_tasks_by_scope.clear()
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+        sessions = tuple(self._sessions.values())
+        self.sessions = {}
+        await self._shutdown_sessions(sessions)
+
+    async def _shutdown_sessions(
+        self,
+        sessions: typing.Iterable[RuntimeMCPSession],
+    ) -> None:
+        """Close MCP sessions in bounded batches to avoid shutdown storms."""
+
+        session_list = list(sessions)
+        for offset in range(0, len(session_list), self._lifecycle_concurrency):
+            batch = session_list[offset : offset + self._lifecycle_concurrency]
+            results = await asyncio.gather(
+                *(session.shutdown() for session in batch),
+                return_exceptions=True,
+            )
+            for session, result in zip(batch, results, strict=True):
+                if isinstance(result, BaseException):
+                    self.ap.logger.error(f'Error shutting down MCP session {session.server_name}: {result}')
+
+    async def _host_server_configs_bounded(
+        self,
+        server_configs: typing.Sequence[tuple[ExecutionContext, dict],],
+    ) -> None:
+        """Create at most one lifecycle batch of MCP host tasks at a time."""
+
+        for offset in range(0, len(server_configs), self._lifecycle_concurrency):
+            batch = server_configs[offset : offset + self._lifecycle_concurrency]
+            tasks: list[asyncio.Task] = []
+            for execution_context, config in batch:
+                task = create_detached_task(
+                    self.host_mcp_server(execution_context, config),
+                    after_commit_manager=getattr(
+                        self.ap,
+                        'persistence_mgr',
+                        None,
+                    ),
+                    workspace_uuid=execution_context.workspace_uuid,
+                )
+                tasks.append(task)
+                try:
+                    self.track_hosted_task(task, execution_context)
+                except WorkspaceInvariantError as exc:
+                    self.ap.logger.warning(
+                        f'Skipping stale MCP startup task for {execution_context.workspace_uuid}: {exc}'
+                    )
+                    continue
+            if tasks:
+                await asyncio.gather(*tasks, return_exceptions=True)
 
     async def _assert_execution_active(
         self,
@@ -1523,6 +1794,7 @@ class MCPLoader(loader.ToolLoader):
         )
         if binding.instance_uuid != execution_context.instance_uuid:
             raise WorkspaceInvariantError('MCP caller instance does not match the active Workspace binding')
+        await self._observe_execution_context(execution_context)
         return execution_context
 
     async def initialize(self):
@@ -1531,9 +1803,36 @@ class MCPLoader(loader.ToolLoader):
     async def load_mcp_servers_from_db(self):
         self.ap.logger.info('Loading MCP servers from db...')
 
-        self.sessions = {}
+        await self._reset_runtime_state()
+
+        pending_hosts: list[tuple[ExecutionContext, dict]] = []
+
+        async def queue_server(binding, server) -> None:
+            config = self.ap.persistence_mgr.serialize_model(
+                persistence_mcp.MCPServer,
+                server,
+            )
+            if config.get('mode') == 'stdio' and not stdio_mcp_enabled(self.ap):
+                self.ap.logger.info(
+                    f'Skipping disabled stdio MCP server {server.uuid}; '
+                    'the persisted configuration is retained but no process is launched'
+                )
+                return
+            try:
+                if binding is None:
+                    binding = await self.ap.workspace_service.get_execution_binding(server.workspace_uuid)
+                execution_context = ExecutionContext(
+                    instance_uuid=binding.instance_uuid,
+                    workspace_uuid=binding.workspace_uuid,
+                    placement_generation=binding.placement_generation,
+                )
+            except Exception as exc:
+                self.ap.logger.warning(
+                    f'Skipping MCP server {server.uuid}: Workspace execution binding is unavailable: {exc}'
+                )
+                return
+            pending_hosts.append((execution_context, config))
 
-        server_configs: list[tuple[typing.Any, typing.Any, dict]] = []
         list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
         tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
         cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
@@ -1548,51 +1847,19 @@ class MCPLoader(loader.ToolLoader):
                         .order_by(persistence_mcp.MCPServer.uuid)
                     )
                     for server in result.all():
-                        server_configs.append(
-                            (
-                                binding,
-                                server,
-                                self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server),
-                            )
-                        )
+                        await queue_server(binding, server)
         else:
             # Compatibility path for isolated loader tests and older embedders.
             result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_mcp.MCPServer))
             for server in result.all():
-                server_configs.append(
-                    (
-                        None,
-                        server,
-                        self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server),
-                    )
-                )
+                await queue_server(None, server)
 
-        for binding, server, config in server_configs:
-            if config.get('mode') == 'stdio' and not stdio_mcp_enabled(self.ap):
-                self.ap.logger.info(
-                    f'Skipping disabled stdio MCP server {server.uuid}; '
-                    'the persisted configuration is retained but no process is launched'
-                )
-                continue
-            try:
-                if binding is None:
-                    binding = await self.ap.workspace_service.get_execution_binding(server.workspace_uuid)
-                execution_context = ExecutionContext(
-                    instance_uuid=binding.instance_uuid,
-                    workspace_uuid=binding.workspace_uuid,
-                    placement_generation=binding.placement_generation,
-                )
-            except Exception as exc:
-                self.ap.logger.warning(
-                    f'Skipping MCP server {server.uuid}: Workspace execution binding is unavailable: {exc}'
-                )
-                continue
-
-            task = create_detached_task(
-                self.host_mcp_server(execution_context, config),
+        if pending_hosts:
+            dispatch_task = create_detached_task(
+                self._host_server_configs_bounded(pending_hosts),
                 after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
             )
-            self._hosted_mcp_tasks.append(task)
+            self._track_host_dispatch_task(dispatch_task)
 
     @staticmethod
     def _scope_key(context: TenantContext) -> tuple[str, str, int]:
@@ -1609,9 +1876,25 @@ class MCPLoader(loader.ToolLoader):
 
     def _sessions_for_context(self, context: TenantContext) -> list[RuntimeMCPSession]:
         scope_key = self._scope_key(context)
-        return [session for key, session in self.sessions.items() if key[:3] == scope_key]
+        return [
+            session
+            for key in self._session_keys_by_scope.get(scope_key, ())
+            if (session := self._sessions.get(key)) is not None
+        ]
 
-    async def host_mcp_server(self, context: TenantContext, server_config: dict):
+    async def host_mcp_server(
+        self,
+        context: TenantContext,
+        server_config: dict,
+    ) -> None:
+        async with self._lifecycle_semaphore:
+            await self._host_mcp_server(context, server_config)
+
+    async def _host_mcp_server(
+        self,
+        context: TenantContext,
+        server_config: dict,
+    ) -> None:
         requested_context = _execution_context_from_tenant(context)
         execution_context = await run_in_workspace_uow(
             self.ap,
@@ -1627,7 +1910,17 @@ class MCPLoader(loader.ToolLoader):
         try:
             session = await self.load_mcp_server(execution_context, server_config)
             await self._assert_execution_active(execution_context)
-            self.sessions[self._session_key(execution_context, server_config['name'])] = session
+            old_session = self._pop_session(
+                execution_context,
+                server_config['name'],
+            )
+            if old_session is not None:
+                await old_session.shutdown()
+            self._register_session(
+                execution_context,
+                server_config['name'],
+                session,
+            )
         except Exception as e:
             self.ap.logger.error(
                 f'Failed to load MCP server from db: {server_config["name"]}({server_config["uuid"]}): {e}\n{traceback.format_exc()}'
@@ -1876,8 +2169,6 @@ class MCPLoader(loader.ToolLoader):
         if include_resource_tools and self._eligible_resource_sessions_for_bound(context, bound_mcp_servers):
             all_functions.extend(self._mcp_synthetic_resource_tools())
 
-        self._last_listed_functions = all_functions
-
         return all_functions
 
     async def get_tool_catalog(
@@ -2140,7 +2431,9 @@ class MCPLoader(loader.ToolLoader):
             self.ap.logger.warning(f'MCP server {server_name} not found in sessions, skipping removal')
             return
 
-        session = self.sessions.pop(key)
+        session = self._pop_session(context, server_name)
+        if session is None:
+            return
         await session.shutdown()
         self.ap.logger.info(f'Removed MCP server: {server_name}')
 
@@ -2180,22 +2473,5 @@ class MCPLoader(loader.ToolLoader):
         """关闭所有工具"""
         self.ap.logger.info('Shutting down all MCP sessions...')
 
-        hosted_tasks = [task for task in self._hosted_mcp_tasks if not task.done()]
-        for task in hosted_tasks:
-            task.cancel()
-        if hosted_tasks:
-            await asyncio.gather(*hosted_tasks, return_exceptions=True)
-        self._hosted_mcp_tasks.clear()
-
-        async def shutdown_session(session: RuntimeMCPSession) -> None:
-            try:
-                await session.shutdown()
-                self.ap.logger.debug(f'Shutdown MCP session: {session.server_name}')
-            except Exception as e:
-                self.ap.logger.error(
-                    f'Error shutting down MCP session {session.server_name}: {e}\n{traceback.format_exc()}'
-                )
-
-        await asyncio.gather(*(shutdown_session(session) for session in list(self.sessions.values())))
-        self.sessions.clear()
+        await self._reset_runtime_state()
         self.ap.logger.info('All MCP sessions shutdown complete')

src/langbot/pkg/provider/tools/loaders/mcp_stdio.py

@@ -6,10 +6,12 @@ import os
 import shutil
 import shlex
 import threading
+import weakref
 from contextlib import suppress, AsyncExitStack, asynccontextmanager
 from typing import TYPE_CHECKING, Any
 
 import pydantic
+from ....utils import bounded_executor
 from mcp import ClientSession
 from mcp.client.websocket import websocket_client
 from ....box.workspace import (
@@ -27,7 +29,7 @@ if TYPE_CHECKING:
     from .mcp import RuntimeMCPSession
 
 
-_WORKSPACE_COPY_LOCKS: dict[str, threading.Lock] = {}
+_WORKSPACE_COPY_LOCKS: weakref.WeakValueDictionary[str, threading.Lock] = weakref.WeakValueDictionary()
 _WORKSPACE_COPY_LOCKS_GUARD = threading.Lock()
 
 
@@ -536,7 +538,11 @@ class BoxStdioSessionRuntime:
             return
         try:
             process_host_root = os.path.join(self._shared_workspace_host_path(), '.mcp', self.process_id)
-            await asyncio.to_thread(shutil.rmtree, process_host_root, True)
+            await bounded_executor.run_blocking_cleanup(
+                shutil.rmtree,
+                process_host_root,
+                True,
+            )
         except Exception as exc:
             self.ap.logger.warning(
                 f'MCP server {self.server_name}: failed to clean staged workspace '

src/langbot/pkg/provider/tools/loaders/native.py

@@ -1,24 +1,29 @@
 from __future__ import annotations
 
+import asyncio
 import base64
 import contextlib
 import errno
+import heapq
 import json
 import os
 import posixpath
 import stat
+import time
 from collections.abc import Iterator
 from dataclasses import dataclass
 from pathlib import PurePosixPath
 
 import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
 from langbot_plugin.api.entities.events import pipeline_query
+import regex
 
 from .. import loader
 from ..errors import ToolNotFoundError
 from .availability import is_box_backend_available
 from . import skill as skill_loader
 from ....api.http.context import ExecutionContext
+from ....utils.bounded_executor import run_blocking_atomic
 
 EXEC_TOOL_NAME = 'exec'
 READ_TOOL_NAME = 'read'
@@ -36,10 +41,18 @@ _DEFAULT_READ_MAX_LINES = 2000
 _MAX_READ_MAX_LINES = 10000
 _DEFAULT_TOOL_RESULT_MAX_BYTES = 50 * 1024
 _BOX_FILE_SCRIPT_MAX_BYTES = 2048
+_MAX_HOST_EDIT_FILE_BYTES = 1024 * 1024
 _GLOB_MAX_MATCHES = 100
+_FILE_WALK_MAX_ENTRIES = 100_000
+_DIRECTORY_MAX_ENTRIES = 10_000
 _GREP_MAX_MATCHES = 200
 _GREP_MAX_FILES = 5000
 _GREP_MAX_LINE_CHARS = 500
+_GREP_MAX_SCAN_LINE_CHARS = 1024 * 1024
+_GREP_MAX_FILE_SCAN_CHARS = 10 * 1024 * 1024
+_GREP_MAX_TOTAL_SCAN_CHARS = 50 * 1024 * 1024
+_GREP_MAX_PATTERN_CHARS = 1024
+_GREP_REGEX_TIMEOUT_SECONDS = 0.25
 
 _DIRECTORY_OPEN_FLAGS = (
     os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0) | getattr(os, 'O_NOFOLLOW', 0) | getattr(os, 'O_CLOEXEC', 0)
@@ -433,7 +446,19 @@ class NativeToolLoader(loader.ToolLoader):
             with _open_location_fd(root_fd, location.relative_parts, os.O_RDONLY) as target_fd:
                 metadata = os.fstat(target_fd)
                 if stat.S_ISDIR(metadata.st_mode):
-                    return self._build_directory_result(os.listdir(target_fd))
+                    entries: list[str] = []
+                    truncated = False
+                    with os.scandir(target_fd) as iterator:
+                        for entry in iterator:
+                            if len(entries) >= _DIRECTORY_MAX_ENTRIES:
+                                truncated = True
+                                break
+                            entries.append(entry.name)
+                    return self._build_directory_result(
+                        entries,
+                        total=len(entries) + int(truncated),
+                        force_truncated_by='entries' if truncated else None,
+                    )
                 if not stat.S_ISREG(metadata.st_mode):
                     raise ValueError('Path must reference a regular file or directory.')
                 return self._read_text_file_preview(target_fd, parameters, metadata=metadata)
@@ -479,10 +504,16 @@ class NativeToolLoader(loader.ToolLoader):
 
         with _open_host_root(location, create=False) as root_fd:
             with _open_location_fd(root_fd, location.relative_parts, os.O_RDWR) as target_fd:
-                if not stat.S_ISREG(os.fstat(target_fd).st_mode):
+                metadata = os.fstat(target_fd)
+                if not stat.S_ISREG(metadata.st_mode):
                     return False, 'File not found.'
-                with os.fdopen(os.dup(target_fd), 'r', encoding='utf-8', errors='replace') as file_obj:
-                    content = file_obj.read()
+                if metadata.st_size > _MAX_HOST_EDIT_FILE_BYTES:
+                    return False, f'File exceeds the {_MAX_HOST_EDIT_FILE_BYTES}-byte edit limit.'
+                with os.fdopen(os.dup(target_fd), 'rb') as file_obj:
+                    raw_content = file_obj.read(_MAX_HOST_EDIT_FILE_BYTES + 1)
+                if len(raw_content) > _MAX_HOST_EDIT_FILE_BYTES:
+                    return False, f'File exceeds the {_MAX_HOST_EDIT_FILE_BYTES}-byte edit limit.'
+                content = raw_content.decode('utf-8', errors='replace')
                 count = content.count(old_string)
                 if count == 0:
                     return False, 'old_string not found in file.'
@@ -490,6 +521,8 @@ class NativeToolLoader(loader.ToolLoader):
                     return False, f'old_string matches {count} locations; provide a more unique string.'
 
                 payload = content.replace(old_string, new_string, 1).encode('utf-8')
+                if len(payload) > _MAX_HOST_EDIT_FILE_BYTES:
+                    return False, f'Edited file exceeds the {_MAX_HOST_EDIT_FILE_BYTES}-byte limit.'
                 os.ftruncate(target_fd, 0)
                 os.lseek(target_fd, 0, os.SEEK_SET)
                 self._write_all(target_fd, payload)
@@ -520,11 +553,19 @@ class NativeToolLoader(loader.ToolLoader):
         return any(candidate and PurePosixPath(relative_path).match(candidate) for candidate in candidates)
 
     def _glob_host_location(self, location: _HostLocation, pattern: str, sandbox_base: str) -> dict:
-        hits: list[tuple[str, float]] = []
+        newest_hits: list[tuple[float, str]] = []
+        total = 0
+        entries_seen = 0
+        scan_truncated = False
 
-        def walk(directory_fd: int, prefix: str) -> None:
+        def walk(directory_fd: int, prefix: str) -> bool:
+            nonlocal entries_seen, scan_truncated, total
             with os.scandir(directory_fd) as entries:
                 for entry in entries:
+                    entries_seen += 1
+                    if entries_seen > _FILE_WALK_MAX_ENTRIES:
+                        scan_truncated = True
+                        return True
                     name = entry.name
                     if name in _SKIP_DIRS:
                         continue
@@ -536,11 +577,17 @@ class NativeToolLoader(loader.ToolLoader):
                         metadata = os.fstat(child_fd)
                         relative = f'{prefix}/{name}' if prefix else name
                         if self._rglob_matches(relative, pattern):
-                            hits.append((relative, metadata.st_mtime))
-                        if stat.S_ISDIR(metadata.st_mode):
-                            walk(child_fd, relative)
+                            total += 1
+                            candidate = (metadata.st_mtime, relative)
+                            if len(newest_hits) < _GLOB_MAX_MATCHES:
+                                heapq.heappush(newest_hits, candidate)
+                            elif candidate > newest_hits[0]:
+                                heapq.heapreplace(newest_hits, candidate)
+                        if stat.S_ISDIR(metadata.st_mode) and walk(child_fd, relative):
+                            return True
                     finally:
                         os.close(child_fd)
+            return False
 
         with _open_host_root(location, create=False) as root_fd:
             with _open_location_fd(root_fd, location.relative_parts, os.O_RDONLY) as target_fd:
@@ -548,12 +595,11 @@ class NativeToolLoader(loader.ToolLoader):
                     return {'ok': False, 'error': f'Path is not a directory: {sandbox_base}'}
                 walk(target_fd, '')
 
-        hits.sort(key=lambda item: item[1], reverse=True)
-        total = len(hits)
+        hits = sorted(newest_hits, reverse=True)
         sandbox_paths: list[str] = []
         output_bytes = 0
         truncated_by_bytes = False
-        for relative, _mtime in hits[:_GLOB_MAX_MATCHES]:
+        for _mtime, relative in hits:
             sandbox_path = self._sandbox_child_path(sandbox_base, relative)
             entry_bytes = len(sandbox_path.encode('utf-8')) + (1 if sandbox_paths else 0)
             if output_bytes + entry_bytes > _DEFAULT_TOOL_RESULT_MAX_BYTES:
@@ -567,27 +613,57 @@ class NativeToolLoader(loader.ToolLoader):
             'matches': sandbox_paths,
             'preview': '\n'.join(sandbox_paths),
             'total': total,
-            'truncated': total > len(sandbox_paths) or truncated_by_bytes,
-            'truncated_by': 'bytes' if truncated_by_bytes else ('matches' if total > len(sandbox_paths) else None),
+            'truncated': scan_truncated or total > len(sandbox_paths) or truncated_by_bytes,
+            'truncated_by': (
+                'scan'
+                if scan_truncated
+                else ('bytes' if truncated_by_bytes else ('matches' if total > len(sandbox_paths) else None))
+            ),
         }
 
     def _grep_host_location(
         self,
         location: _HostLocation,
-        regex,
+        pattern: str,
         include: str | None,
         sandbox_base: str,
     ) -> dict:
+        try:
+            compiled = regex.compile(pattern)
+        except regex.error as exc:
+            return {'ok': False, 'error': f'Invalid regex: {exc}'}
+
         matches: list[dict] = []
         output_bytes = 0
         truncated_by: str | None = None
         files_seen = 0
+        entries_seen = 0
+        total_chars_seen = 0
+        deadline = time.monotonic() + _GREP_REGEX_TIMEOUT_SECONDS
 
         def grep_file(file_fd: int, sandbox_path: str) -> bool:
-            nonlocal output_bytes, truncated_by
+            nonlocal output_bytes, total_chars_seen, truncated_by
+            file_chars_seen = 0
             with os.fdopen(os.dup(file_fd), 'r', encoding='utf-8', errors='ignore') as handle:
-                for lineno, line in enumerate(handle, 1):
-                    if not regex.search(line):
+                lineno = 0
+                while True:
+                    line = handle.readline(_GREP_MAX_SCAN_LINE_CHARS + 1)
+                    if not line:
+                        break
+                    lineno += 1
+                    line_chars = len(line)
+                    file_chars_seen += line_chars
+                    total_chars_seen += line_chars
+                    if file_chars_seen > _GREP_MAX_FILE_SCAN_CHARS or total_chars_seen > _GREP_MAX_TOTAL_SCAN_CHARS:
+                        truncated_by = 'scan'
+                        return True
+                    if line_chars > _GREP_MAX_SCAN_LINE_CHARS and not line.endswith('\n'):
+                        truncated_by = truncated_by or 'line'
+                        return False
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        raise TimeoutError
+                    if not compiled.search(line, timeout=remaining, concurrent=True):
                         continue
                     content, line_truncated = self._truncate_grep_line(line.rstrip())
                     entry = {'file': sandbox_path, 'line': lineno, 'content': content}
@@ -605,9 +681,13 @@ class NativeToolLoader(loader.ToolLoader):
             return False
 
         def walk(directory_fd: int, prefix: str) -> bool:
-            nonlocal files_seen
+            nonlocal entries_seen, files_seen, truncated_by
             with os.scandir(directory_fd) as entries:
                 for entry in entries:
+                    entries_seen += 1
+                    if entries_seen > _FILE_WALK_MAX_ENTRIES:
+                        truncated_by = 'scan'
+                        return True
                     name = entry.name
                     if name in _SKIP_DIRS:
                         continue
@@ -637,13 +717,16 @@ class NativeToolLoader(loader.ToolLoader):
 
         with _open_host_root(location, create=False) as root_fd:
             with _open_location_fd(root_fd, location.relative_parts, os.O_RDONLY) as target_fd:
-                metadata = os.fstat(target_fd)
-                if stat.S_ISREG(metadata.st_mode):
-                    grep_file(target_fd, sandbox_base)
-                elif stat.S_ISDIR(metadata.st_mode):
-                    walk(target_fd, '')
-                else:
-                    return {'ok': False, 'error': f'Path not found: {sandbox_base}'}
+                try:
+                    metadata = os.fstat(target_fd)
+                    if stat.S_ISREG(metadata.st_mode):
+                        grep_file(target_fd, sandbox_base)
+                    elif stat.S_ISDIR(metadata.st_mode):
+                        walk(target_fd, '')
+                    else:
+                        return {'ok': False, 'error': f'Path not found: {sandbox_base}'}
+                except TimeoutError:
+                    return {'ok': False, 'error': 'Regex search timed out'}
 
         return {
             'ok': True,
@@ -701,9 +784,24 @@ if not path.startswith('/workspace'):
 elif not os.path.exists(path):
     print(json.dumps({{'ok': False, 'error': f'File not found: {{path}}'}}))
 elif os.path.isdir(path):
-    entries = sorted(os.listdir(path))
+    entries = []
+    directory_truncated = False
+    with os.scandir(path) as iterator:
+        for entry in iterator:
+            if len(entries) >= {_DIRECTORY_MAX_ENTRIES}:
+                directory_truncated = True
+                break
+            entries.append(entry.name)
+    entries.sort()
     content = '\\n'.join(entries)
-    print(json.dumps({{'ok': True, 'content': content, 'is_directory': True, 'total': len(entries), 'truncated': False}}))
+    print(json.dumps({{
+        'ok': True,
+        'content': content,
+        'is_directory': True,
+        'total': len(entries) + int(directory_truncated),
+        'truncated': directory_truncated,
+        'truncated_by': 'entries' if directory_truncated else None,
+    }}))
 elif encoding == 'base64':
     size_bytes = os.path.getsize(path)
     with open(path, 'rb') as f:
@@ -824,7 +922,7 @@ else:
 
     async def _glob_workspace_via_box(self, path: str, pattern: str, query: pipeline_query.Query) -> dict:
         script = f"""
-import json, os
+import heapq, json, os
 from pathlib import Path
 path = {json.dumps(path)}
 pattern = {json.dumps(pattern)}
@@ -835,12 +933,28 @@ elif not os.path.isdir(path):
     print(json.dumps({{'ok': False, 'error': f'Path is not a directory: {{path}}'}}))
 else:
     base = Path(path)
-    hits = [
-        item for item in base.rglob(pattern)
-        if not any(part in skip_dirs for part in item.parts)
-    ]
-    hits.sort(key=lambda item: item.stat().st_mtime if item.exists() else 0, reverse=True)
-    shown = hits[:{_GLOB_MAX_MATCHES}]
+    newest_hits = []
+    total = 0
+    entries_seen = 0
+    scan_truncated = False
+    for item in base.rglob(pattern):
+        entries_seen += 1
+        if entries_seen > {_FILE_WALK_MAX_ENTRIES}:
+            scan_truncated = True
+            break
+        if any(part in skip_dirs for part in item.parts):
+            continue
+        total += 1
+        try:
+            mtime = item.stat().st_mtime
+        except OSError:
+            mtime = 0
+        candidate = (mtime, str(item))
+        if len(newest_hits) < {_GLOB_MAX_MATCHES}:
+            heapq.heappush(newest_hits, candidate)
+        elif candidate > newest_hits[0]:
+            heapq.heapreplace(newest_hits, candidate)
+    shown = [Path(item_path) for _mtime, item_path in sorted(newest_hits, reverse=True)]
     matches = []
     output_bytes = 0
     truncated_by_bytes = False
@@ -857,9 +971,12 @@ else:
         'ok': True,
         'matches': matches,
         'preview': '\\n'.join(matches),
-        'total': len(hits),
-        'truncated': len(hits) > len(matches) or truncated_by_bytes,
-        'truncated_by': 'bytes' if truncated_by_bytes else ('matches' if len(hits) > len(matches) else None),
+        'total': total,
+        'truncated': scan_truncated or total > len(matches) or truncated_by_bytes,
+        'truncated_by': (
+            'scan' if scan_truncated
+            else ('bytes' if truncated_by_bytes else ('matches' if total > len(matches) else None))
+        ),
     }}))
 """.strip()
         return await self._run_workspace_file_script(script, query)
@@ -872,12 +989,15 @@ else:
         query: pipeline_query.Query,
     ) -> dict:
         script = f"""
-import json, os, re
+import json, os, re, signal, time
 from pathlib import Path
 path = {json.dumps(path)}
 pattern = {json.dumps(pattern)}
 include = {json.dumps(include)}
 skip_dirs = {json.dumps(sorted(_SKIP_DIRS))}
+def regex_timeout(_signum, _frame):
+    raise TimeoutError
+signal.signal(signal.SIGALRM, regex_timeout)
 try:
     regex = re.compile(pattern)
 except re.error as exc:
@@ -888,6 +1008,17 @@ else:
     elif not os.path.exists(path):
         print(json.dumps({{'ok': False, 'error': f'Path not found: {{path}}'}}))
     else:
+        regex_deadline = time.monotonic() + {_GREP_REGEX_TIMEOUT_SECONDS}
+        def bounded_search(value):
+            remaining = regex_deadline - time.monotonic()
+            if remaining <= 0:
+                raise TimeoutError
+            signal.setitimer(signal.ITIMER_REAL, remaining)
+            try:
+                return regex.search(value)
+            finally:
+                signal.setitimer(signal.ITIMER_REAL, 0)
+
         base = Path(path)
         if base.is_file():
             files = [base]
@@ -904,14 +1035,37 @@ else:
         matches = []
         output_bytes = 0
         truncated_by = None
+        total_chars_seen = 0
         for fp in files:
             try:
                 handle = fp.open('r', encoding='utf-8', errors='ignore')
             except OSError:
                 continue
+            file_chars_seen = 0
             with handle:
-                for lineno, line in enumerate(handle, 1):
-                    if regex.search(line):
+                lineno = 0
+                while True:
+                    line = handle.readline({_GREP_MAX_SCAN_LINE_CHARS + 1})
+                    if not line:
+                        break
+                    lineno += 1
+                    file_chars_seen += len(line)
+                    total_chars_seen += len(line)
+                    if (
+                        file_chars_seen > {_GREP_MAX_FILE_SCAN_CHARS}
+                        or total_chars_seen > {_GREP_MAX_TOTAL_SCAN_CHARS}
+                    ):
+                        truncated_by = 'scan'
+                        break
+                    if len(line) > {_GREP_MAX_SCAN_LINE_CHARS} and not line.endswith('\\n'):
+                        truncated_by = truncated_by or 'line'
+                        break
+                    try:
+                        matched = bounded_search(line)
+                    except TimeoutError:
+                        print(json.dumps({{'ok': False, 'error': 'Regex search timed out'}}))
+                        raise SystemExit(0)
+                    if matched:
                         if base.is_file():
                             file_path = path
                         else:
@@ -934,9 +1088,9 @@ else:
                         if len(matches) >= {_GREP_MAX_MATCHES}:
                             truncated_by = truncated_by or 'matches'
                             break
-                if truncated_by == 'bytes' or len(matches) >= {_GREP_MAX_MATCHES}:
+                if truncated_by in ('bytes', 'scan') or len(matches) >= {_GREP_MAX_MATCHES}:
                     break
-            if truncated_by == 'bytes' or len(matches) >= {_GREP_MAX_MATCHES}:
+            if truncated_by in ('bytes', 'scan') or len(matches) >= {_GREP_MAX_MATCHES}:
                 break
 
         print(json.dumps({{
@@ -966,7 +1120,7 @@ else:
                 host_location = None
             if host_location is not None:
                 try:
-                    return self._read_host_location(host_location, parameters)
+                    return await asyncio.to_thread(self._read_host_location, host_location, parameters)
                 except FileNotFoundError:
                     pass
 
@@ -998,7 +1152,7 @@ else:
         if self._should_use_box_workspace_files(host_location.selected_skill):
             return await self._read_workspace_via_box(path, parameters, query)
         try:
-            return self._read_host_location(host_location, parameters)
+            return await asyncio.to_thread(self._read_host_location, host_location, parameters)
         except (FileNotFoundError, NotADirectoryError):
             return {'ok': False, 'error': f'File not found: {path}'}
 
@@ -1031,7 +1185,7 @@ else:
         if self._should_use_box_workspace_files(host_location.selected_skill):
             return await self._write_workspace_via_box(path, content, parameters, query)
         try:
-            self._write_host_location(host_location, content, parameters)
+            await run_blocking_atomic(self._write_host_location, host_location, content, parameters)
         except ValueError as exc:
             return {'ok': False, 'error': str(exc)}
         self._refresh_skill_from_disk(query, host_location.selected_skill)
@@ -1091,7 +1245,12 @@ else:
         if self._should_use_box_workspace_files(host_location.selected_skill):
             return await self._edit_workspace_via_box(path, old_string, new_string, query)
         try:
-            changed, error = self._edit_host_location(host_location, old_string, new_string)
+            changed, error = await run_blocking_atomic(
+                self._edit_host_location,
+                host_location,
+                old_string,
+                new_string,
+            )
         except (FileNotFoundError, NotADirectoryError):
             return {'ok': False, 'error': f'File not found: {path}'}
         if not changed:
@@ -1364,7 +1523,7 @@ else:
         if self._should_use_box_workspace_files(host_location.selected_skill):
             return await self._glob_workspace_via_box(path, pattern, query)
         try:
-            return self._glob_host_location(host_location, pattern, path)
+            return await asyncio.to_thread(self._glob_host_location, host_location, pattern, path)
         except (FileNotFoundError, NotADirectoryError):
             return {'ok': False, 'error': f'Path is not a directory: {path}'}
 
@@ -1374,12 +1533,8 @@ else:
         include = parameters.get('include')
         self.ap.logger.info(f'grep tool invoked: query_id={query.query_id} pattern={pattern} path={path}')
 
-        import re
-
-        try:
-            regex = re.compile(pattern)
-        except re.error as e:
-            return {'ok': False, 'error': f'Invalid regex: {e}'}
+        if not isinstance(pattern, str) or len(pattern) > _GREP_MAX_PATTERN_CHARS:
+            return {'ok': False, 'error': f'Regex patterns may contain at most {_GREP_MAX_PATTERN_CHARS} characters'}
 
         host_location = self._resolve_host_location(
             query,
@@ -1390,7 +1545,13 @@ else:
         if self._should_use_box_workspace_files(host_location.selected_skill):
             return await self._grep_workspace_via_box(path, pattern, include, query)
         try:
-            return self._grep_host_location(host_location, regex, include, path)
+            return await asyncio.to_thread(
+                self._grep_host_location,
+                host_location,
+                pattern,
+                include,
+                path,
+            )
         except (FileNotFoundError, NotADirectoryError):
             return {'ok': False, 'error': f'Path not found: {path}'}
 
@@ -1430,18 +1591,24 @@ else:
             normalized['truncated_by'] = 'bytes'
         return normalized
 
-    def _build_directory_result(self, entries: list[str]) -> dict:
+    def _build_directory_result(
+        self,
+        entries: list[str],
+        *,
+        total: int | None = None,
+        force_truncated_by: str | None = None,
+    ) -> dict:
         sorted_entries = sorted(str(entry) for entry in entries)
         content = '\n'.join(sorted_entries)
         preview = self._truncate_text_to_bytes(content, _DEFAULT_TOOL_RESULT_MAX_BYTES)
-        truncated = preview != content
+        truncated_by = force_truncated_by or ('bytes' if preview != content else None)
         return {
             'ok': True,
             'content': preview,
             'is_directory': True,
-            'total': len(sorted_entries),
-            'truncated': truncated,
-            'truncated_by': 'bytes' if truncated else None,
+            'total': len(sorted_entries) if total is None else total,
+            'truncated': truncated_by is not None,
+            'truncated_by': truncated_by,
         }
 
     def _read_text_file_preview(self, file_fd: int, parameters: dict, *, metadata: os.stat_result) -> dict: