Merge pull request #74 from YuJunZhiXue/trae

Trae

Author: YuJunZhiXue

backend/adapter/cli_proxy.py

@@ -7,6 +7,7 @@
 
 from backend.adapter.standard_request import StandardRequest, CLAUDE_CODE_OPENAI_PROFILE
 from backend.core.config import resolve_model
+from backend.runtime.visible_text import sanitize_visible_text, sanitize_visible_text_blocks
 from backend.services.model_modes import parse_model_mode
 from backend.services.prompt_builder import messages_to_prompt
 from backend.services.workspace_context import derive_workspace_root
@@ -189,6 +190,7 @@ def to_openai_response(execution, standard_request: StandardRequest) -> dict:
         Returns:
             dict: OpenAI 格式的响应
         """
+        visible_text = sanitize_visible_text(execution.state.answer_text)
         return {
             "id": f"chatcmpl-{execution.chat_id[:12]}",
             "object": "chat.completion",
@@ -199,15 +201,15 @@ def to_openai_response(execution, standard_request: StandardRequest) -> dict:
                     "index": 0,
                     "message": {
                         "role": "assistant",
-                        "content": execution.state.answer_text,
+                        "content": visible_text,
                     },
                     "finish_reason": "stop",
                 }
             ],
             "usage": {
                 "prompt_tokens": len(standard_request.prompt),
-                "completion_tokens": len(execution.state.answer_text),
-                "total_tokens": len(standard_request.prompt) + len(execution.state.answer_text),
+                "completion_tokens": len(visible_text),
+                "total_tokens": len(standard_request.prompt) + len(visible_text),
             },
         }
 
@@ -231,10 +233,10 @@ def to_anthropic_response(execution, standard_request: StandardRequest, msg_id:
 
         # 添加思考内容
         if execution.state.reasoning_text:
-            content_blocks.append({"type": "thinking", "thinking": execution.state.reasoning_text})
+            content_blocks.append({"type": "thinking", "thinking": sanitize_visible_text(execution.state.reasoning_text)})
 
         # 添加工具调用块
-        content_blocks.extend(directive.tool_blocks)
+        content_blocks.extend(sanitize_visible_text_blocks(directive.tool_blocks))
         visible_text = tool_directive_visible_text(directive, execution.state.answer_text)
         if (
             directive.stop_reason != "tool_use"
@@ -269,11 +271,12 @@ def to_gemini_response(execution, standard_request: StandardRequest) -> dict:
         Returns:
             dict: Gemini 格式的响应
         """
+        visible_text = sanitize_visible_text(execution.state.answer_text)
         return {
             "candidates": [
                 {
                     "content": {
-                        "parts": [{"text": execution.state.answer_text}],
+                        "parts": [{"text": visible_text}],
                         "role": "model",
                     },
                     "finishReason": "STOP",
@@ -282,8 +285,8 @@ def to_gemini_response(execution, standard_request: StandardRequest) -> dict:
             ],
             "usageMetadata": {
                 "promptTokenCount": len(standard_request.prompt),
-                "candidatesTokenCount": len(execution.state.answer_text),
-                "totalTokenCount": len(standard_request.prompt) + len(execution.state.answer_text),
+                "candidatesTokenCount": len(visible_text),
+                "totalTokenCount": len(standard_request.prompt) + len(visible_text),
             },
         }
 

backend/api/admin.py

@@ -269,8 +269,11 @@ async def get_settings(request: Request):
         "max_inflight_per_account": backend_settings.MAX_INFLIGHT_PER_ACCOUNT,
         "global_max_inflight": getattr(acc_pool, "global_max_inflight", 0),
         "max_queue_size": getattr(acc_pool, "max_queue_size", 0),
+        "account_ready_set_threshold": backend_settings.ACCOUNT_READY_SET_THRESHOLD,
+        "account_ready_set_enabled": getattr(acc_pool, "ready_set_enabled", False),
         "chat_id_pool_target": pool.target if pool else 0,
         "chat_id_pool_ttl_seconds": pool.ttl if pool else 0,
+        "chat_id_pool_max_concurrency": pool.max_concurrency if pool else 0,
         "model_aliases": safe_map,
     }
 
@@ -286,6 +289,15 @@ async def update_settings(data: dict, request: Request):
                 pool.set_max_inflight(val)
         except (TypeError, ValueError):
             pass
+    if "account_ready_set_threshold" in data:
+        try:
+            val = max(1, int(data["account_ready_set_threshold"]))
+            settings.ACCOUNT_READY_SET_THRESHOLD = val
+            pool = getattr(request.app.state, "account_pool", None)
+            if pool is not None and hasattr(pool, "_reset_concurrency_limits"):
+                pool._reset_concurrency_limits()
+        except (TypeError, ValueError):
+            pass
     if "global_max_inflight" in data:
         try:
             val = int(data["global_max_inflight"])
@@ -294,12 +306,13 @@ async def update_settings(data: dict, request: Request):
                 pool.global_max_inflight = val
         except (TypeError, ValueError):
             pass
-    if "chat_id_pool_target" in data or "chat_id_pool_ttl_seconds" in data:
+    if "chat_id_pool_target" in data or "chat_id_pool_ttl_seconds" in data or "chat_id_pool_max_concurrency" in data:
         cp = getattr(request.app.state, "chat_id_pool", None)
         if cp is not None:
             await cp.apply_config(
                 target=data.get("chat_id_pool_target"),
                 ttl_seconds=data.get("chat_id_pool_ttl_seconds"),
+                max_concurrency=data.get("chat_id_pool_max_concurrency"),
             )
     if "model_aliases" in data:
         MODEL_MAP.clear()

backend/api/anthropic.py

@@ -20,6 +20,7 @@
     request_max_attempts,
     tool_directive_visible_text,
 )
+from backend.runtime.visible_text import VisibleTextSanitizer, sanitize_visible_text, sanitize_visible_text_blocks
 from backend.services.auth_quota import resolve_auth_context
 from backend.services.completion_bridge import force_fresh_chat_after_empty_response, is_empty_upstream_response
 from backend.services.context_attachment_manager import prepare_context_attachments, derive_session_key
@@ -75,6 +76,8 @@ def __init__(self, *, msg_id: str, model_name: str, prompt: str):
         self.block_index = 0
         self.current_block: dict[str, object] = {"type": None, "index": None, "tool_call_id": None}
         self.opened_tool_calls: set[str] = set()
+        self.answer_sanitizer = VisibleTextSanitizer()
+        self.thinking_sanitizer = VisibleTextSanitizer()
 
     def ensure_message_start(self) -> None:
         if not self.pending_chunks:
@@ -122,12 +125,24 @@ def open_tool_block(self, tool_call_id: str, tool_name: str) -> int:
         return index
 
     def append_thinking_delta(self, text_chunk: str) -> None:
+        text_chunk = self.thinking_sanitizer.feed(text_chunk)
+        if not text_chunk:
+            return
         index = self.open_textual_block("thinking")
         self.pending_chunks.append(
             stream_presenter.anthropic_content_block_delta(index, {"type": "thinking_delta", "thinking": text_chunk})
         )
 
     def buffer_answer_text(self, text_chunk: str) -> None:
+        thinking_tail = self.thinking_sanitizer.flush()
+        if thinking_tail:
+            index = self.open_textual_block("thinking")
+            self.pending_chunks.append(
+                stream_presenter.anthropic_content_block_delta(index, {"type": "thinking_delta", "thinking": thinking_tail})
+            )
+        text_chunk = self.answer_sanitizer.feed(text_chunk)
+        if not text_chunk:
+            return
         index = self.open_textual_block("text")
         self.answer_text_buffer.append((index, text_chunk))
 
@@ -150,6 +165,8 @@ def flush_answer_text(self) -> None:
 
     def clear_answer_text(self) -> None:
         self.answer_text_buffer = []
+        self.flushed_answer_text = ""
+        self.answer_sanitizer.reset()
 
     def answer_text(self) -> str:
         return "".join(text_chunk for _, text_chunk in self.answer_text_buffer)
@@ -178,6 +195,18 @@ def buffer_missing_answer_tail(self, final_text: str) -> None:
             self.answer_text_buffer = []
             self.buffer_answer_text(final_text)
 
+    def flush_text_sanitizers(self) -> None:
+        answer_tail = self.answer_sanitizer.flush()
+        if answer_tail:
+            index = self.open_textual_block("text")
+            self.answer_text_buffer.append((index, answer_tail))
+        thinking_tail = self.thinking_sanitizer.flush()
+        if thinking_tail:
+            index = self.open_textual_block("thinking")
+            self.pending_chunks.append(
+                stream_presenter.anthropic_content_block_delta(index, {"type": "thinking_delta", "thinking": thinking_tail})
+            )
+
 
 def _build_standard_request(req_data: dict) -> StandardRequest:
     """浣跨敤 CLIProxy 杩涜鍗忚杞崲"""
@@ -448,6 +477,7 @@ async def on_delta(evt, text_chunk, _):
 
                             directive = build_tool_directive(standard_request, execution.state, history_messages=history_messages)
                             visible_text = tool_directive_visible_text(directive, execution.state.answer_text)
+                            stream_state.flush_text_sanitizers()
                             if directive.stop_reason != "tool_use":
                                 stream_state.buffer_missing_answer_tail(visible_text)
                             if (
@@ -607,8 +637,8 @@ async def on_delta(evt, text_chunk, _):
                     _log_response_tool_blocks("json_response", directive.tool_blocks)
                     content_blocks: list[dict] = []
                     if execution.state.reasoning_text:
-                        content_blocks.append({"type": "thinking", "thinking": execution.state.reasoning_text})
-                    content_blocks.extend(directive.tool_blocks)
+                        content_blocks.append({"type": "thinking", "thinking": sanitize_visible_text(execution.state.reasoning_text)})
+                    content_blocks.extend(sanitize_visible_text_blocks(directive.tool_blocks))
                     if (
                         directive.stop_reason != "tool_use"
                         and visible_text

backend/api/gemini.py

@@ -11,6 +11,7 @@
 from backend.core.request_logging import new_request_id, request_context, update_request_context
 from backend.runtime import stream_presenter
 from backend.runtime.execution import collect_completion_run, cleanup_runtime_resources
+from backend.runtime.visible_text import VisibleTextSanitizer, sanitize_visible_text
 from backend.services.auth_quota import resolve_auth_context
 from backend.services.completion_bridge import force_fresh_chat_after_empty_response, is_empty_upstream_response
 from backend.services.token_calc import calculate_usage
@@ -74,7 +75,8 @@ async def gemini_generate_content(model: str, request: Request):
             log.error(f"Gemini proxy failed: {e}")
             raise HTTPException(status_code=500, detail=str(e))
 
-        usage = calculate_usage(content, execution.state.answer_text)
+        visible_text = sanitize_visible_text(execution.state.answer_text)
+        usage = calculate_usage(content, visible_text)
         users = await users_db.get()
         for u in users:
             if u["id"] == token:
@@ -83,13 +85,13 @@ async def gemini_generate_content(model: str, request: Request):
         await users_db.save(users)
         await cleanup_runtime_resources(client, execution.acc, execution.chat_id)
 
-        log.info(f"[Gemini] Request complete. Generated {len(execution.state.answer_text)} characters.")
+        log.info(f"[Gemini] Request complete. Generated {len(visible_text)} visible characters.")
         return JSONResponse(
             {
                 "candidates": [
                     {
                         "content": {
-                            "parts": [{"text": execution.state.answer_text}],
+                            "parts": [{"text": visible_text}],
                             "role": "model",
                         }
                     }
@@ -109,10 +111,13 @@ async def gemini_stream_generate_content(model: str, request: Request):
 
         async def generate():
             queue: asyncio.Queue[str | None] = asyncio.Queue()
+            answer_sanitizer = VisibleTextSanitizer()
 
             async def on_delta(evt, text_chunk, _):
                 if text_chunk and evt.get("phase") == "answer":
-                    await queue.put(stream_presenter.gemini_text_chunk(text_chunk))
+                    visible_chunk = answer_sanitizer.feed(text_chunk)
+                    if visible_chunk:
+                        await queue.put(stream_presenter.gemini_text_chunk(visible_chunk))
 
             async def runner():
                 execution = None
@@ -129,18 +134,22 @@ async def runner():
                         await cleanup_runtime_resources(client, execution.acc, execution.chat_id, preserve_chat=False)
                         raise RuntimeError("empty upstream response")
 
-                    usage = calculate_usage(content, execution.state.answer_text)
+                    visible_text = sanitize_visible_text(execution.state.answer_text)
+                    usage = calculate_usage(content, visible_text)
                     users = await users_db.get()
                     for u in users:
                         if u["id"] == token:
                             u["used_tokens"] += usage["total_tokens"]
                             break
                     await users_db.save(users)
                     await cleanup_runtime_resources(client, execution.acc, execution.chat_id)
-                    log.info(f"[Gemini] Request complete. Generated {len(execution.state.answer_text)} characters.")
+                    log.info(f"[Gemini] Request complete. Generated {len(visible_text)} visible characters.")
                 except Exception as e:
                     await queue.put(json.dumps({"error": str(e)}) + "\n")
                 finally:
+                    visible_tail = answer_sanitizer.flush()
+                    if visible_tail:
+                        await queue.put(stream_presenter.gemini_text_chunk(visible_tail))
                     await queue.put(None)
 
             task = asyncio.create_task(runner())

backend/api/responses.py

@@ -14,6 +14,7 @@
 from backend.core.request_logging import new_request_id, request_context, update_request_context
 from backend.core.request_trace import log_test_prompt, prompt_tail
 from backend.runtime.execution import build_tool_directive, build_usage_delta_factory, request_max_attempts, tool_directive_visible_text
+from backend.runtime.visible_text import VisibleTextSanitizer, sanitize_visible_text
 from backend.services.attachment_preprocessor import preprocess_attachments
 from backend.services.auth_quota import resolve_auth_context
 from backend.services.client_profiles import detect_openai_client_profile
@@ -457,7 +458,7 @@ def build_responses_payload(
         output_text = visible_text
         content: list[dict[str, Any]] = []
         if execution.state.reasoning_text:
-            content.append({"type": "reasoning_text", "text": execution.state.reasoning_text})
+            content.append({"type": "reasoning_text", "text": sanitize_visible_text(execution.state.reasoning_text)})
         content.append({"type": "output_text", "text": output_text, "annotations": []})
         payload["output"] = [{
             "id": f"msg_{uuid.uuid4().hex[:12]}",
@@ -490,7 +491,10 @@ def __init__(self, *, response_id: str, created: int, model_name: str, prompt: s
         self.started_text = False
         self.pending_chunks: list[str] = []
         self.answer_fragments: list[str] = []
+        self.visible_answer_fragments: list[str] = []
         self.reasoning_fragments: list[str] = []
+        self.answer_sanitizer = VisibleTextSanitizer()
+        self.reasoning_sanitizer = VisibleTextSanitizer()
         self.tool_calls_emitted = False
 
     def initial_chunks(self) -> list[str]:
@@ -523,6 +527,9 @@ def _ensure_text_item(self) -> None:
 
     def on_delta(self, evt: dict[str, Any], text_chunk: str | None, tool_calls: list[dict[str, Any]] | None) -> None:
         if text_chunk and evt.get("phase") in ("think", "thinking_summary"):
+            text_chunk = self.reasoning_sanitizer.feed(text_chunk)
+            if not text_chunk:
+                return
             self.reasoning_fragments.append(text_chunk)
             self.pending_chunks.append(_sse("response.reasoning_text.delta", {
                 "response_id": self.response_id,
@@ -534,8 +541,12 @@ def on_delta(self, evt: dict[str, Any], text_chunk: str | None, tool_calls: list
             return
 
         if text_chunk and evt.get("phase") == "answer":
-            self._ensure_text_item()
             self.answer_fragments.append(text_chunk)
+            text_chunk = self.answer_sanitizer.feed(text_chunk)
+            if not text_chunk:
+                return
+            self._ensure_text_item()
+            self.visible_answer_fragments.append(text_chunk)
             self.pending_chunks.append(_sse("response.output_text.delta", {
                 "response_id": self.response_id,
                 "item_id": self.message_id,
@@ -553,6 +564,29 @@ def drain_pending(self) -> list[str]:
         self.pending_chunks = []
         return chunks
 
+    def _flush_visible_sanitizers(self) -> None:
+        reasoning_tail = self.reasoning_sanitizer.flush()
+        if reasoning_tail:
+            self.reasoning_fragments.append(reasoning_tail)
+            self.pending_chunks.append(_sse("response.reasoning_text.delta", {
+                "response_id": self.response_id,
+                "item_id": self.message_id,
+                "output_index": self.output_index,
+                "content_index": self.content_index,
+                "delta": reasoning_tail,
+            }))
+        answer_tail = self.answer_sanitizer.flush()
+        if answer_tail:
+            self._ensure_text_item()
+            self.visible_answer_fragments.append(answer_tail)
+            self.pending_chunks.append(_sse("response.output_text.delta", {
+                "response_id": self.response_id,
+                "item_id": self.message_id,
+                "output_index": self.output_index,
+                "content_index": self.content_index,
+                "delta": answer_tail,
+            }))
+
     def emit_tool_calls(self, tool_calls: list[dict[str, Any]]) -> None:
         for tool_call in tool_calls:
             block = {
@@ -595,6 +629,7 @@ def emit_tool_calls(self, tool_calls: list[dict[str, Any]]) -> None:
             self.tool_calls_emitted = True
 
     def finalize(self, execution, directive) -> list[str]:
+        self._flush_visible_sanitizers()
         chunks = self.drain_pending()
         final_text = tool_directive_visible_text(
             directive,
@@ -621,9 +656,10 @@ def finalize(self, execution, directive) -> list[str]:
                     "delta": final_text,
                 }))
                 self.answer_fragments.append(final_text)
+                self.visible_answer_fragments.append(final_text)
             chunks.extend(self.drain_pending())
         elif directive.stop_reason != "tool_use":
-            streamed_text = "".join(self.answer_fragments)
+            streamed_text = "".join(self.visible_answer_fragments)
             if final_text.startswith(streamed_text):
                 missing_tail = final_text[len(streamed_text):]
                 if missing_tail:
@@ -635,6 +671,7 @@ def finalize(self, execution, directive) -> list[str]:
                         "delta": missing_tail,
                     }))
                     self.answer_fragments.append(missing_tail)
+                    self.visible_answer_fragments.append(missing_tail)
 
         if self.started_text:
             chunks.append(_sse("response.output_text.done", {