chore: import original project baseline
Import the pre-repair source tree as the history baseline. Runtime data (data/), virtualenvs, bytecode caches and logs are gitignored so local secrets and user state stay out of the repo.
This commit is contained in:
@@ -0,0 +1,58 @@
|
||||
"""
|
||||
core/agent —— pi agent 核心框架的 Python 1:1 重构
|
||||
==================================================
|
||||
模块对照（pi-main → 本包）：
|
||||
packages/agent/src/types.ts → types.py 数据模型（消息/事件/工具/配置）
|
||||
packages/agent/src/agent-loop.ts → loop.py 核心循环（runLoop 1:1）
|
||||
packages/agent/src/agent.ts → agent.py Agent 状态机（prompt/steer/followUp/abort）
|
||||
packages/ai/src/api/*.ts → stream_fn.py OpenAI 兼容流式（vLLM）+ 输出预算钳制
|
||||
packages/ai/src/api/simple-options→ context.py token 估算 + clampMaxTokensToContext
|
||||
packages/agent/src/compaction.ts → compaction.py 上下文压缩（切分/摘要/替换）
|
||||
agent-session.ts 后置恢复逻辑 → recovery.py 重试退避 + 溢出/截断压缩恢复
|
||||
packages/agent/src/tools/*.ts → tools.py 工具管线 + 内置 read/bash/write/edit
|
||||
|
||||
使用示例（最小闭环）：
|
||||
from core.agent import Agent, AgentConfig, ModelConfig
|
||||
from core.agent.stream_fn import openai_stream
|
||||
from core.agent.recovery import AgentRunner
|
||||
from core.agent.tools import default_tools
|
||||
|
||||
cfg = AgentConfig(model=ModelConfig(...), tools=default_tools())
|
||||
agent = Agent(cfg)
|
||||
agent.set_stream_fn(openai_stream)
|
||||
runner = AgentRunner(agent)
|
||||
result = runner.run("你好")
|
||||
"""
|
||||
from .types import (AgentConfig, AgentError, AgentEvent, AgentMessage,
|
||||
AgentState, AgentTool, AgentToolResult, AbortSignal,
|
||||
AssistantMessageEvent, ModelConfig, RetryConfig, RunResult,
|
||||
ToolCall, new_id)
|
||||
from .agent import Agent
|
||||
from .context import (CONTEXT_SAFETY_TOKENS, calculate_context_tokens,
|
||||
clamp_max_tokens_to_context,
|
||||
estimate_context_tokens, estimate_message_tokens,
|
||||
should_compact)
|
||||
from .compaction import (CompactionSettings,
|
||||
DEFAULT_COMPACTION_SETTINGS, compact_context,
|
||||
find_cut_point, prepare_compaction)
|
||||
from .recovery import (AgentRunner, compute_retry_delay_ms, find_last_assistant,
|
||||
is_context_overflow, is_recoverable_length,
|
||||
is_retryable_assistant_error)
|
||||
from .stream_fn import (classify_error, from_openai_messages, openai_stream,
|
||||
to_openai_messages)
|
||||
from .tools import default_tools
|
||||
|
||||
# 🆕 版本记录：核心框架为 pi (badlogic/pi-mono) 0.81.x 时代的 Python 移植，
|
||||
# 2026-09 对齐 pi 0.85.1 的估算/锚定/钳制/重试语义（P0/P1/P2/M1/M2/M3）。
|
||||
__version__ = "1.0.0 (pi 0.81.x port, aligned 0.85.1 semantics)"
|
||||
__all__ = [
|
||||
"Agent", "AgentConfig", "AgentError", "AgentEvent", "AgentMessage",
|
||||
"AgentRunner", "AgentState", "AgentTool", "AgentToolResult",
|
||||
"AbortSignal", "AssistantMessageEvent", "CONTEXT_SAFETY_TOKENS",
|
||||
"ModelConfig", "RetryConfig", "RunResult", "ToolCall",
|
||||
"clamp_max_tokens_to_context", "calculate_context_tokens", "classify_error", "compact_context",
|
||||
"compute_retry_delay_ms", "default_tools", "estimate_context_tokens",
|
||||
"find_last_assistant", "find_cut_point", "is_context_overflow",
|
||||
"is_recoverable_length", "is_retryable_assistant_error", "new_id",
|
||||
"openai_stream", "should_compact", "to_openai_messages",
|
||||
]
|
||||
@@ -0,0 +1,204 @@
|
||||
"""
|
||||
core/agent/agent.py
|
||||
===================
|
||||
🌟 pi Agent 状态机的 Python 1:1 移植
|
||||
|
||||
对照 pi-main 源码：packages/agent/src/agent.ts (592 行)
|
||||
class Agent {
|
||||
state: AgentState # messages / isStreaming / error
|
||||
config: AgentConfig
|
||||
private listeners: Set<cb>
|
||||
subscribe(cb) => unsubscribe # 事件订阅（TUI/Qt 都靠它）
|
||||
prompt(message | text) => Promise # 追加 user 消息并跑一轮 run-to-settle
|
||||
continue() => Promise # 从当前上下文继续（不能以 assistant 结尾）
|
||||
steer(text) # 中途注入（turn 边界消费）
|
||||
followUp(text) # 停止后注入（agent 停止时消费 → 续跑）
|
||||
abort() # 中止当前 run（下一边界生效）
|
||||
}
|
||||
|
||||
关键语义（与 pi 完全一致）：
|
||||
1. steering 在「循环开始前」和「每个 turn 结束后」被消费
|
||||
（one-at-a-time 每次取 1 条；"all" 一次取光）
|
||||
2. followUp 只在「run 即将结束」时被外层循环消费 → 触发续跑
|
||||
3. abort 不抛异常：信号置位 → 当前流/工具在下一检查点收尾 →
|
||||
助手消息 stop_reason="aborted" → agent_end
|
||||
4. 同一个 Agent 实例可反复 prompt（state.messages 持续累积 = 会话记忆）
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import Any, Callable, List, Optional, Union
|
||||
|
||||
from .loop import run_loop
|
||||
from .types import (AgentConfig, AgentError, AgentEvent, AgentMessage,
|
||||
AgentState, AbortSignal, RunResult, new_id)
|
||||
|
||||
|
||||
class Agent:
|
||||
def __init__(self, config: AgentConfig):
|
||||
self.config = config
|
||||
self.state = AgentState()
|
||||
self._listeners: List[Callable[[AgentEvent], None]] = []
|
||||
self._steering_queue: List[AgentMessage] = []
|
||||
self._follow_up_queue: List[AgentMessage] = []
|
||||
self._lock = threading.Lock()
|
||||
self._active_signal: Optional[AbortSignal] = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 事件订阅 —— 对照 agent.ts subscribe()
|
||||
# ------------------------------------------------------------------
|
||||
def subscribe(self, cb: Callable[[AgentEvent], None]) -> Callable[[], None]:
|
||||
self._listeners.append(cb)
|
||||
|
||||
def unsubscribe():
|
||||
try:
|
||||
self._listeners.remove(cb)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
return unsubscribe
|
||||
|
||||
def _emit(self, event: AgentEvent):
|
||||
for cb in list(self._listeners):
|
||||
try:
|
||||
cb(event)
|
||||
except Exception:
|
||||
# 监听器异常不打断循环（pi 同款宽容策略）
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 队列操作 —— 对照 agent.ts 的 getSteeringMessages / getFollowUpMessages
|
||||
# ------------------------------------------------------------------
|
||||
def _take_steering(self) -> List[AgentMessage]:
|
||||
mode = self.config.steering_mode
|
||||
with self._lock:
|
||||
if not self._steering_queue:
|
||||
return []
|
||||
if mode == "all":
|
||||
out = self._steering_queue[:]
|
||||
self._steering_queue.clear()
|
||||
else: # one-at-a-time（pi 默认）
|
||||
out = [self._steering_queue.pop(0)]
|
||||
return out
|
||||
|
||||
def _take_follow_ups(self) -> List[AgentMessage]:
|
||||
mode = self.config.follow_up_mode
|
||||
with self._lock:
|
||||
if not self._follow_up_queue:
|
||||
return []
|
||||
if mode == "all":
|
||||
out = self._follow_up_queue[:]
|
||||
self._follow_up_queue.clear()
|
||||
else:
|
||||
out = [self._follow_up_queue.pop(0)]
|
||||
return out
|
||||
|
||||
def has_queued(self) -> bool:
|
||||
with self._lock:
|
||||
return bool(self._steering_queue or self._follow_up_queue)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 用户入口 —— 对照 agent.ts prompt / steer / followUp / abort
|
||||
# ------------------------------------------------------------------
|
||||
def prompt(self, message: Union[str, AgentMessage]) -> RunResult:
|
||||
"""追加一条 user 消息并运行到 settle（同步阻塞，跑在调用线程上）"""
|
||||
if self.state.is_streaming:
|
||||
raise AgentError(message="Agent 正在运行中，不能并发 prompt（请先 abort）",
|
||||
kind="unknown")
|
||||
if isinstance(message, str):
|
||||
message = AgentMessage(role="user", content=message)
|
||||
if message.role != "user":
|
||||
raise AgentError(message="prompt 只接受 user 消息", kind="unknown")
|
||||
self.state.messages.append(message)
|
||||
return self._run(new_message=message)
|
||||
|
||||
def continue_(self) -> RunResult:
|
||||
"""
|
||||
对照 agent.ts continue() / runAgentLoopContinue 前置检查：
|
||||
- 上下文不能为空
|
||||
- 不能以 assistant 消息结尾
|
||||
"""
|
||||
if self.state.is_streaming:
|
||||
raise AgentError(message="Agent 正在运行中", kind="unknown")
|
||||
if not self.state.messages:
|
||||
raise AgentError(message="上下文为空，无法 continue", kind="unknown")
|
||||
if self.state.messages[-1].role == "assistant":
|
||||
raise AgentError(
|
||||
message="不能以 assistant 消息结尾来 continue（应先注入 user/toolResult）",
|
||||
kind="unknown")
|
||||
return self._run(new_message=None)
|
||||
|
||||
def steer(self, text: str):
|
||||
"""中途注入：在当前 run 的 turn 边界被消费（pi steering）"""
|
||||
with self._lock:
|
||||
self._steering_queue.append(AgentMessage(role="user", content=text))
|
||||
|
||||
def follow_up(self, text: str):
|
||||
"""停止后注入：run 即将结束时被外层循环消费 → 自动续跑（pi followUp）"""
|
||||
with self._lock:
|
||||
self._follow_up_queue.append(AgentMessage(role="user", content=text))
|
||||
|
||||
def abort(self):
|
||||
"""
|
||||
中止当前 run。
|
||||
对照 pi: 置位 AbortSignal；当前流在 chunk 边界关闭，
|
||||
工具在检查点返回 "Operation aborted"，助手消息以 "aborted" 收尾。
|
||||
"""
|
||||
sig = self._active_signal
|
||||
if sig is not None:
|
||||
sig.abort("aborted")
|
||||
with self._lock:
|
||||
# 与 pi 一致：未消费的队列消息保留，下次 run 生效
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 内部：一次 run（同步）
|
||||
# ------------------------------------------------------------------
|
||||
def _run(self, new_message: Optional[AgentMessage]) -> RunResult:
|
||||
signal = AbortSignal()
|
||||
self._active_signal = signal
|
||||
self.state.is_streaming = True
|
||||
self.state.error = None
|
||||
try:
|
||||
result = run_loop(self, new_message, signal, self._make_stream_fn())
|
||||
except AgentError:
|
||||
raise
|
||||
except Exception as e:
|
||||
self.state.error = AgentError(message=f"循环异常: {e}", kind="unknown")
|
||||
self._emit(AgentEvent(type="agent_end", stop_reason="error",
|
||||
error=self.state.error))
|
||||
result = RunResult(stop_reason="error", error=self.state.error)
|
||||
finally:
|
||||
self.state.is_streaming = False
|
||||
self.state.streaming_message = None
|
||||
self.state.streaming_delta = {}
|
||||
self._active_signal = None
|
||||
return result
|
||||
|
||||
def _finish_run(self, new_messages: List[AgentMessage], stop_reason: str,
|
||||
error: Optional[AgentError]):
|
||||
"""循环收尾回调（loop 层调用）"""
|
||||
if error is not None:
|
||||
self.state.error = error
|
||||
# 本轮新消息已在循环中逐条 append 进 state.messages；此处仅记录收尾
|
||||
self._last_run = (new_messages, stop_reason)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 流函数（由 AgentWorker 注入具体 provider 实现；默认报错）
|
||||
# ------------------------------------------------------------------
|
||||
def set_stream_fn(self, stream_fn: Callable):
|
||||
"""
|
||||
stream_fn(context, model, signal, max_tokens, tools=None) -> Iterator[(kind, payload)]
|
||||
对照 pi 的 streamFn 注入点（agentLoopConfig.streamFn）。
|
||||
tools 参数 = 本轮可用的 AgentTool 列表，需序列化进 provider 请求。
|
||||
"""
|
||||
self._stream_fn = stream_fn
|
||||
|
||||
def _make_stream_fn(self) -> Callable:
|
||||
fn = getattr(self, "_stream_fn", None)
|
||||
if fn is None:
|
||||
def _missing(context, model, signal, max_tokens, tools=None):
|
||||
raise AgentError(message="未配置 stream_fn", kind="unknown")
|
||||
yield # pragma: no cover
|
||||
return _missing
|
||||
return fn
|
||||
@@ -0,0 +1,637 @@
|
||||
"""
|
||||
core/agent/compaction.py
|
||||
========================
|
||||
🌟 pi 上下文压缩算法的 Python 1:1 移植（harness 原版算法）
|
||||
|
||||
对照 pi-main 源码（逐函数对应）：
|
||||
packages/agent/src/harness/compaction/compaction.ts
|
||||
DEFAULT_COMPACTION_SETTINGS -> CompactionSettings / DEFAULT_COMPACTION_SETTINGS
|
||||
shouldCompact -> 见 context.py: should_compact（同一公式）
|
||||
estimateTokens / estimateContextTokens -> context.py（usage 锚定，CJK 感知见说明）
|
||||
findValidCutPoints / findCutPoint -> find_valid_cut_points / find_cut_point
|
||||
findTurnStartIndex -> find_turn_start
|
||||
prepareCompaction -> prepare_compaction
|
||||
SUMMARIZATION_SYSTEM_PROMPT -> 同名（逐字移植）
|
||||
SUMMARIZATION_PROMPT -> 同名（逐字移植）
|
||||
UPDATE_SUMMARIZATION_PROMPT -> 同名（逐字移植，迭代更新用）
|
||||
TURN_PREFIX_SUMMARIZATION_PROMPT -> 同名（逐字移植，断轮前缀用）
|
||||
generateSummaryWithUsage -> generate_summary（maxTokens = 0.8×reserve）
|
||||
generateTurnPrefixSummary -> generate_turn_prefix_summary（0.5×reserve）
|
||||
compact -> compact_context（断轮双摘要 + 拼接格式 1:1）
|
||||
packages/agent/src/harness/compaction/utils.ts
|
||||
serializeConversation -> serialize_conversation（1:1，含 2000 字符截断）
|
||||
extractFileOpsFromMessage -> extract_file_ops_from_message
|
||||
computeFileLists -> compute_file_lists
|
||||
formatFileOperations -> format_file_operations
|
||||
TOOL_RESULT_MAX_CHARS = 2000 -> 同名常量
|
||||
|
||||
摘要 LLM 调用由上层注入：
|
||||
summarize_fn(prompt_text: str, system_prompt: str, max_tokens: int) -> str
|
||||
（pi 里是 models.completeSimple + retry；haocode 用 OpenAI 客户端非流式调用，
|
||||
由 llm_engine.AgentWorker 实现并注入。）
|
||||
|
||||
已声明的偏差（仅 2 处，见 context.py 头注）：
|
||||
1. 单条消息 token 估算用 CJK 感知启发式（pi 是 chars/4）——对中文会话更安全
|
||||
2. 摘要 LLM 调用失败时降级为机械摘录（pi 返回 CompactionError）——桌面应用优先不丢上下文
|
||||
其余全部 1:1：触发公式、usage 锚定、token 预算切点、有效切点规则、断轮双摘要、
|
||||
迭代式 previousSummary 更新、摘要提示词逐字、文件操作附录、拼接格式。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, List, Optional, Set, Tuple
|
||||
|
||||
from .types import AgentMessage, ModelConfig
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 压缩设置 —— 1:1 对照 DEFAULT_COMPACTION_SETTINGS
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class CompactionSettings:
|
||||
"""pi: interface CompactionSettings { enabled; reserveTokens; keepRecentTokens }"""
|
||||
enabled: bool = True
|
||||
reserve_tokens: int = 16384 # pi: 16384（摘要提示词与输出预留）
|
||||
keep_recent_tokens: int = 20000 # pi: 20000（压缩后保留的近期上下文预算）
|
||||
|
||||
|
||||
DEFAULT_COMPACTION_SETTINGS = CompactionSettings()
|
||||
|
||||
# pi utils.ts: const TOOL_RESULT_MAX_CHARS = 2000
|
||||
TOOL_RESULT_MAX_CHARS = 2000
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 摘要提示词 —— 从 pi compaction.ts 逐字移植（不得改写，摘要质量依赖它）
|
||||
# ======================================================================
|
||||
SUMMARIZATION_SYSTEM_PROMPT = \
|
||||
"You are a context summarization assistant. Your task is to read a conversation " \
|
||||
"between a user and an AI assistant, then produce a structured summary following " \
|
||||
"the exact format specified.\n\n" \
|
||||
"Do NOT continue the conversation. Do NOT respond to any questions in the " \
|
||||
"conversation. ONLY output the structured summary."
|
||||
|
||||
SUMMARIZATION_PROMPT = """The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Any constraints, preferences, or requirements mentioned by user]
|
||||
- [Or "(none)" if none were mentioned]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Completed tasks/changes]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Current work]
|
||||
|
||||
### Blocked
|
||||
- [Issues preventing progress, if any]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale]
|
||||
|
||||
## Next Steps
|
||||
1. [Ordered list of what should happen next]
|
||||
|
||||
## Critical Context
|
||||
- [Any data, examples, or references needed to continue]
|
||||
- [Or "(none)" if not applicable]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages."""
|
||||
|
||||
UPDATE_SUMMARIZATION_PROMPT = """The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.
|
||||
|
||||
Update the existing structured summary with new information. RULES:
|
||||
- PRESERVE all existing information from the previous summary
|
||||
- ADD new progress, decisions, and context from the new messages
|
||||
- UPDATE the Progress section: move items from "In Progress" to "Done" when completed
|
||||
- UPDATE "Next Steps" based on what was accomplished
|
||||
- PRESERVE exact file paths, function names, and error messages
|
||||
- If something is no longer relevant, you may remove it
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[Preserve existing goals, add new ones if the task expanded]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Preserve existing, add new ones discovered]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Include previously done items AND newly completed items]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Current work - update based on progress]
|
||||
|
||||
### Blocked
|
||||
- [Current blockers - remove if resolved]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale] (preserve all previous, add new)
|
||||
|
||||
## Next Steps
|
||||
1. [Update based on current state]
|
||||
|
||||
## Critical Context
|
||||
- [Preserve important context, add new if needed]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages."""
|
||||
|
||||
TURN_PREFIX_SUMMARIZATION_PROMPT = """This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained.
|
||||
|
||||
Summarize the prefix to provide context for the retained suffix:
|
||||
|
||||
## Original Request
|
||||
[What did the user ask for in this turn?]
|
||||
|
||||
## Early Progress
|
||||
- [Key decisions and work done in the prefix]
|
||||
|
||||
## Context for Suffix
|
||||
- [Information needed to understand the retained recent work]
|
||||
|
||||
Be concise. Focus on what's needed to understand the kept suffix."""
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 对话序列化 —— 1:1 对照 utils.ts serializeConversation
|
||||
# ======================================================================
|
||||
def _safe_json(value) -> str:
|
||||
try:
|
||||
s = json.dumps(value, ensure_ascii=False)
|
||||
return s if s is not None else "undefined"
|
||||
except Exception:
|
||||
return "[unserializable]"
|
||||
|
||||
|
||||
def _content_text(content) -> str:
|
||||
"""pi contentText: str 或 [{type:"text",text}] 列表取文本拼接"""
|
||||
if content is None:
|
||||
return ""
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts = []
|
||||
for b in content:
|
||||
if isinstance(b, dict) and b.get("type") == "text":
|
||||
parts.append(str(b.get("text", "")))
|
||||
return "\n".join(p for p in parts if p)
|
||||
return str(content)
|
||||
|
||||
|
||||
def _truncate_for_summary(text: str, max_chars: int) -> str:
|
||||
"""pi utils.ts truncateForSummary（逐字逻辑）"""
|
||||
if len(text) <= max_chars:
|
||||
return text
|
||||
truncated = len(text) - max_chars
|
||||
return f"{text[:max_chars]}\n\n[... {truncated} more characters truncated]"
|
||||
|
||||
|
||||
def serialize_conversation(messages: List[AgentMessage]) -> str:
|
||||
"""
|
||||
1:1 对照 utils.ts serializeConversation（输出格式逐字一致）：
|
||||
[User]: ...
|
||||
[Assistant thinking]: ...
|
||||
[Assistant]: ...
|
||||
[Assistant tool calls]: name(k=v, k2=v2); name2(...)
|
||||
[Tool result]: ...（超 2000 字符截断）
|
||||
"""
|
||||
parts: List[str] = []
|
||||
for msg in messages:
|
||||
if msg.role == "user":
|
||||
content = _content_text(msg.content)
|
||||
if content:
|
||||
parts.append(f"[User]: {content}")
|
||||
elif msg.role == "assistant":
|
||||
thinking_parts = []
|
||||
tool_calls = []
|
||||
if msg.reasoning:
|
||||
thinking_parts.append(msg.reasoning)
|
||||
text = _content_text(msg.content)
|
||||
for tc in (msg.tool_calls or []):
|
||||
args_str = ", ".join(f"{k}={_safe_json(v)}"
|
||||
for k, v in (tc.arguments or {}).items())
|
||||
tool_calls.append(f"{tc.name}({args_str})")
|
||||
if thinking_parts:
|
||||
parts.append(f"[Assistant thinking]: {chr(10).join(thinking_parts)}")
|
||||
if text:
|
||||
parts.append(f"[Assistant]: {text}")
|
||||
if tool_calls:
|
||||
parts.append(f"[Assistant tool calls]: {'; '.join(tool_calls)}")
|
||||
elif msg.role == "toolResult":
|
||||
content = _content_text(msg.content)
|
||||
if content:
|
||||
parts.append(f"[Tool result]: "
|
||||
f"{_truncate_for_summary(content, TOOL_RESULT_MAX_CHARS)}")
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 文件操作提取 —— 1:1 对照 utils.ts extractFileOps*/computeFileLists/formatFileOperations
|
||||
# ======================================================================
|
||||
class FileOperations:
|
||||
def __init__(self):
|
||||
self.read: Set[str] = set()
|
||||
self.written: Set[str] = set()
|
||||
self.edited: Set[str] = set()
|
||||
|
||||
|
||||
def extract_file_ops_from_message(message: AgentMessage, file_ops: FileOperations):
|
||||
"""pi: assistant 的 toolCall 参数里 path 字段 → read/write/edit 归类"""
|
||||
if message.role != "assistant":
|
||||
return
|
||||
for tc in (message.tool_calls or []):
|
||||
args = tc.arguments or {}
|
||||
path = args.get("path")
|
||||
if not isinstance(path, str) or not path:
|
||||
continue
|
||||
if tc.name == "read":
|
||||
file_ops.read.add(path)
|
||||
elif tc.name == "write":
|
||||
file_ops.written.add(path)
|
||||
elif tc.name == "edit":
|
||||
file_ops.edited.add(path)
|
||||
|
||||
|
||||
def compute_file_lists(file_ops: FileOperations) -> Tuple[List[str], List[str]]:
|
||||
"""pi computeFileLists: modified = edited|written；readOnly = read-modified；均排序"""
|
||||
modified = file_ops.edited | file_ops.written
|
||||
read_only = sorted(f for f in file_ops.read if f not in modified)
|
||||
return read_only, sorted(modified)
|
||||
|
||||
|
||||
def format_file_operations(read_files: List[str], modified_files: List[str]) -> str:
|
||||
"""pi formatFileOperations: <read-files>/<modified-files> 标签拼接"""
|
||||
sections = []
|
||||
if read_files:
|
||||
sections.append("<read-files>\n" + "\n".join(read_files) + "\n</read-files>")
|
||||
if modified_files:
|
||||
sections.append("<modified-files>\n" + "\n".join(modified_files) + "\n</modified-files>")
|
||||
if not sections:
|
||||
return ""
|
||||
return "\n\n" + "\n\n".join(sections)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 切分点 —— 1:1 对照 findValidCutPoints / findTurnStartIndex / findCutPoint
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class CutPointResult:
|
||||
"""pi: interface CutPointResult"""
|
||||
first_kept_index: int # 保留段首条在 compactable 列表里的下标
|
||||
turn_start_index: int = -1 # 断轮时：该轮起点（user 消息）下标；否则 -1
|
||||
is_split_turn: bool = False
|
||||
|
||||
|
||||
def find_valid_cut_points(messages: List[AgentMessage],
|
||||
start_index: int, end_index: int) -> List[int]:
|
||||
"""
|
||||
pi findValidCutPoints：消息角色为 user/assistant 的位置是有效切点
|
||||
（toolResult 不能做切点——它会与前面的 toolCall 分离）。
|
||||
pi 里的 bashExecution/branchSummary/compactionSummary 等角色
|
||||
在 haocode 消息模型中不存在，等价规则即 role in (user, assistant)。
|
||||
"""
|
||||
cut_points = []
|
||||
for i in range(start_index, end_index):
|
||||
if messages[i].role in ("user", "assistant"):
|
||||
cut_points.append(i)
|
||||
return cut_points
|
||||
|
||||
|
||||
def find_turn_start(messages: List[AgentMessage], entry_index: int,
|
||||
start_index: int) -> int:
|
||||
"""pi findTurnStartIndex：向前找本轮起点（user 消息 / branch_summary）"""
|
||||
for i in range(entry_index, start_index - 1, -1):
|
||||
if messages[i].role == "user":
|
||||
return i
|
||||
return -1
|
||||
|
||||
|
||||
def _estimate(msg: AgentMessage) -> int:
|
||||
from .context import estimate_message_tokens
|
||||
return estimate_message_tokens(msg)
|
||||
|
||||
|
||||
def find_cut_point(messages: List[AgentMessage], start_index: int,
|
||||
end_index: int, keep_recent_tokens: int) -> CutPointResult:
|
||||
"""
|
||||
1:1 对照 pi findCutPoint：
|
||||
1. 从尾部向前累计 token，直到累计 >= keep_recent_tokens
|
||||
2. 取该位置（含）之后的第一个有效切点
|
||||
3. 切点不是 user 消息 → 断轮：找本轮起点，前缀单独摘要
|
||||
（pi 的「回退跳过状态条目」循环针对 session 状态条目；
|
||||
haocode 消息列表没有状态条目，等价省略。）
|
||||
"""
|
||||
cut_points = find_valid_cut_points(messages, start_index, end_index)
|
||||
if not cut_points:
|
||||
return CutPointResult(first_kept_index=start_index)
|
||||
|
||||
accumulated = 0
|
||||
cut_index = cut_points[0]
|
||||
for i in range(end_index - 1, start_index - 1, -1):
|
||||
accumulated += _estimate(messages[i])
|
||||
if accumulated >= keep_recent_tokens:
|
||||
for c in cut_points:
|
||||
if c >= i:
|
||||
cut_index = c
|
||||
break
|
||||
break
|
||||
|
||||
is_user = messages[cut_index].role == "user"
|
||||
turn_start = -1 if is_user else find_turn_start(messages, cut_index, start_index)
|
||||
is_split = (not is_user) and turn_start != -1
|
||||
return CutPointResult(first_kept_index=cut_index,
|
||||
turn_start_index=turn_start,
|
||||
is_split_turn=is_split)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 压缩准备 —— 1:1 对照 prepareCompaction
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class CompactionPreparation:
|
||||
"""pi: interface CompactionPreparation"""
|
||||
messages_to_summarize: List[AgentMessage] = field(default_factory=list)
|
||||
turn_prefix_messages: List[AgentMessage] = field(default_factory=list)
|
||||
retained_tail: List[AgentMessage] = field(default_factory=list)
|
||||
is_split_turn: bool = False
|
||||
tokens_before: int = 0
|
||||
previous_summary: Optional[str] = None
|
||||
file_ops: FileOperations = field(default_factory=FileOperations)
|
||||
settings: CompactionSettings = DEFAULT_COMPACTION_SETTINGS
|
||||
|
||||
|
||||
def _compaction_diag(line: str) -> None:
|
||||
"""压缩自诊断日志（与 recovery.compact_diag_log 同一文件/规范：
|
||||
print + 追加 compaction_diag.log + 失败静默）。"""
|
||||
import os as _os
|
||||
import time as _time
|
||||
t = _time.time()
|
||||
stamp = (f"[{_time.strftime('%H:%M:%S', _time.localtime(t))}"
|
||||
f".{int(t * 1000) % 1000:03d}]")
|
||||
print(f"{stamp} {line}", flush=True)
|
||||
try:
|
||||
path = _os.path.join(
|
||||
_os.path.dirname(_os.path.dirname(_os.path.dirname(_os.path.abspath(__file__)))),
|
||||
"compaction_diag.log")
|
||||
with open(path, "a", encoding="utf-8") as f:
|
||||
f.write(f"{stamp} {line}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _nothing_to_summarize_diag(messages: List[AgentMessage],
|
||||
compactable: List[AgentMessage],
|
||||
cut: CutPointResult,
|
||||
settings: CompactionSettings) -> str:
|
||||
"""prepare_compaction 返回 None（无可摘要内容）时的自诊断行：
|
||||
记录分支子类型 + 关键量，下次失败可直接从日志定位原因。
|
||||
|
||||
子类型：
|
||||
no_valid_cut_points —— 可压缩范围内没有任何 user/assistant 条目
|
||||
total_below_keep_recent —— 可压缩范围总量 < keep_recent（通常意味着
|
||||
上下文主体是旧摘要本身，无新内容）
|
||||
cut_pinned_at_zero —— 预算从尾部累加只在 i=0 才达标，即首条
|
||||
条目独占 ≥ (总量-keep_recent) 的 token
|
||||
（单条超长消息/巨型工具输出主导上下文）
|
||||
"""
|
||||
from .context import estimate_context_tokens, estimate_message_tokens
|
||||
try:
|
||||
total = estimate_context_tokens(messages).tokens
|
||||
n_user = sum(1 for m in compactable if m.role == "user")
|
||||
n_asst = sum(1 for m in compactable if m.role == "assistant")
|
||||
ests = [(estimate_message_tokens(m), i) for i, m in enumerate(compactable)]
|
||||
e0 = ests[0][0] if ests else 0
|
||||
rest = sum(e for e, _ in ests[1:])
|
||||
top3 = sorted(ests, reverse=True)[:3]
|
||||
roles = " ".join(("T" if m.role == "toolResult" else m.role[0].upper())
|
||||
for m in compactable[:10])
|
||||
if not find_valid_cut_points(compactable, 0, len(compactable)):
|
||||
sub = "no_valid_cut_points"
|
||||
elif sum(e for e, _ in ests) < settings.keep_recent_tokens:
|
||||
sub = "total_below_keep_recent"
|
||||
else:
|
||||
sub = "cut_pinned_at_zero"
|
||||
big = "; ".join(f"idx{i}={e}" for e, i in top3)
|
||||
return (f"[COMPACT_NONE] branch=nothing_to_summarize sub={sub} "
|
||||
f"msgs={len(messages)} est_total={total} compactable={len(compactable)} "
|
||||
f"keep_recent={settings.keep_recent_tokens} "
|
||||
f"first_kept={cut.first_kept_index} turn_start={cut.turn_start_index} "
|
||||
f"split={cut.is_split_turn} user={n_user} asst={n_asst} "
|
||||
f"est_first={e0} est_rest={rest} top3=[{big}] head_roles=[{roles}]")
|
||||
except Exception as ex: # 诊断本身失败不能影响主流程
|
||||
return f"[COMPACT_NONE] branch=nothing_to_summarize (diag failed: {ex})"
|
||||
|
||||
|
||||
def prepare_compaction(messages: List[AgentMessage],
|
||||
settings: Optional[CompactionSettings] = None
|
||||
) -> Optional[CompactionPreparation]:
|
||||
"""
|
||||
1:1 对照 pi prepareCompaction：
|
||||
- 上一条压缩摘要（messages[0].kind == "compaction_summary"）不重复摘要，
|
||||
其内容作为 previousSummary 走迭代更新提示词
|
||||
- 可压缩范围 = 摘要之后的全部消息（pi 里等价于「上次保留尾 + 新消息」）
|
||||
- 切点在可压缩范围内选；tokens_before 按完整上下文（含摘要消息）估算
|
||||
不可压缩（空/无摘要对象）时返回 None（对照 pi 返回 ok(undefined)）。
|
||||
"""
|
||||
from .context import estimate_context_tokens
|
||||
settings = settings or DEFAULT_COMPACTION_SETTINGS
|
||||
if not messages:
|
||||
_compaction_diag("[COMPACT_NONE] branch=empty_messages")
|
||||
return None
|
||||
if messages[-1].kind == "compaction_summary":
|
||||
_compaction_diag(
|
||||
f"[COMPACT_NONE] branch=tail_is_summary msgs={len(messages)} "
|
||||
f"tail_head={_content_text(messages[-1].content)[:60]!r}")
|
||||
return None
|
||||
|
||||
previous_summary = None
|
||||
if messages[0].kind == "compaction_summary":
|
||||
previous_summary = _content_text(messages[0].content)
|
||||
compactable = messages[1:]
|
||||
else:
|
||||
compactable = messages
|
||||
if not compactable:
|
||||
_compaction_diag("[COMPACT_NONE] branch=compactable_empty (上下文只剩旧摘要)")
|
||||
return None
|
||||
|
||||
tokens_before = estimate_context_tokens(messages).tokens
|
||||
cut = find_cut_point(compactable, 0, len(compactable),
|
||||
settings.keep_recent_tokens)
|
||||
|
||||
history_end = cut.turn_start_index if cut.is_split_turn else cut.first_kept_index
|
||||
messages_to_summarize = compactable[:history_end]
|
||||
turn_prefix_messages = []
|
||||
if cut.is_split_turn:
|
||||
turn_prefix_messages = compactable[cut.turn_start_index:cut.first_kept_index]
|
||||
retained_tail = compactable[cut.first_kept_index:]
|
||||
|
||||
if not messages_to_summarize and not turn_prefix_messages:
|
||||
# 🆕 自诊断：记录是哪种子条件导致无东西可摘要（见 _nothing_to_summarize_diag）
|
||||
_compaction_diag(_nothing_to_summarize_diag(messages, compactable, cut,
|
||||
settings))
|
||||
return None # 没有可摘要内容
|
||||
|
||||
file_ops = FileOperations()
|
||||
for m in messages_to_summarize:
|
||||
extract_file_ops_from_message(m, file_ops)
|
||||
if cut.is_split_turn:
|
||||
for m in turn_prefix_messages:
|
||||
extract_file_ops_from_message(m, file_ops)
|
||||
|
||||
return CompactionPreparation(
|
||||
messages_to_summarize=messages_to_summarize,
|
||||
turn_prefix_messages=turn_prefix_messages,
|
||||
retained_tail=retained_tail,
|
||||
is_split_turn=cut.is_split_turn,
|
||||
tokens_before=tokens_before,
|
||||
previous_summary=previous_summary,
|
||||
file_ops=file_ops,
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 摘要生成 —— 1:1 对照 generateSummaryWithUsage / generateTurnPrefixSummary
|
||||
# ======================================================================
|
||||
def _build_summary_prompt(conversation_text: str, previous_summary: Optional[str]) -> str:
|
||||
"""pi generateSummaryWithUsage 的 prompt 组装（逐字结构）"""
|
||||
base = UPDATE_SUMMARIZATION_PROMPT if previous_summary else SUMMARIZATION_PROMPT
|
||||
prompt = f"<conversation>\n{conversation_text}\n</conversation>\n\n"
|
||||
if previous_summary:
|
||||
prompt += f"<previous-summary>\n{previous_summary}\n</previous-summary>\n\n"
|
||||
prompt += base
|
||||
return prompt
|
||||
|
||||
|
||||
def generate_summary(messages: List[AgentMessage],
|
||||
summarize_fn: Callable[[str, str, int], str],
|
||||
reserve_tokens: int,
|
||||
model_max_tokens: int,
|
||||
previous_summary: Optional[str] = None
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""
|
||||
返回 (summary_text, error)。
|
||||
maxTokens = min(0.8 × reserveTokens, model.maxTokens) —— 1:1 对照。
|
||||
"""
|
||||
max_tokens = min(
|
||||
int(0.8 * reserve_tokens),
|
||||
model_max_tokens if model_max_tokens > 0 else (1 << 30),
|
||||
)
|
||||
conversation = serialize_conversation(messages)
|
||||
prompt = _build_summary_prompt(conversation, previous_summary)
|
||||
try:
|
||||
text = summarize_fn(prompt, SUMMARIZATION_SYSTEM_PROMPT, max_tokens)
|
||||
except Exception as e:
|
||||
return None, f"Summarization failed: {e}"
|
||||
if not text or not text.strip():
|
||||
return None, "Summarization failed: empty response"
|
||||
return text.strip(), None
|
||||
|
||||
|
||||
def generate_turn_prefix_summary(messages: List[AgentMessage],
|
||||
summarize_fn: Callable[[str, str, int], str],
|
||||
reserve_tokens: int,
|
||||
model_max_tokens: int
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""maxTokens = min(0.5 × reserveTokens, model.maxTokens) —— 1:1 对照"""
|
||||
max_tokens = min(
|
||||
int(0.5 * reserve_tokens),
|
||||
model_max_tokens if model_max_tokens > 0 else (1 << 30),
|
||||
)
|
||||
conversation = serialize_conversation(messages)
|
||||
prompt = f"<conversation>\n{conversation}\n</conversation>\n\n{TURN_PREFIX_SUMMARIZATION_PROMPT}"
|
||||
try:
|
||||
text = summarize_fn(prompt, SUMMARIZATION_SYSTEM_PROMPT, max_tokens)
|
||||
except Exception as e:
|
||||
return None, f"Turn prefix summarization failed: {e}"
|
||||
if not text or not text.strip():
|
||||
return None, "Turn prefix summarization failed: empty response"
|
||||
return text.strip(), None
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 主入口 —— 1:1 对照 compact()
|
||||
# ======================================================================
|
||||
def compact_context(messages: List[AgentMessage],
|
||||
model: ModelConfig,
|
||||
summarize_fn: Callable[[str, str, int], str],
|
||||
settings: Optional[CompactionSettings] = None
|
||||
) -> Optional[List[AgentMessage]]:
|
||||
"""
|
||||
执行压缩。返回 [压缩摘要消息] + 保留尾巴；不可压缩时返回 None。
|
||||
|
||||
摘要消息: role="user", kind="compaction_summary"，content 为纯摘要文本
|
||||
（pi 的 compaction 条目；下次压缩时自动走迭代更新提示词）。
|
||||
|
||||
断轮（切点落在某轮中间）时 1:1 对照 pi compact()：
|
||||
历史摘要 与 轮前缀摘要 分两次 LLM 调用，拼接为
|
||||
{history}\n\n---\n\n**Turn Context (split turn):**\n\n{prefix}
|
||||
最后追加文件操作附录（<read-files>/<modified-files>）。
|
||||
"""
|
||||
settings = settings or DEFAULT_COMPACTION_SETTINGS
|
||||
prep = prepare_compaction(messages, settings)
|
||||
if prep is None:
|
||||
return None
|
||||
|
||||
history_text: Optional[str] = None
|
||||
prefix_error: Optional[str] = None
|
||||
history_error: Optional[str] = None
|
||||
|
||||
if prep.is_split_turn and prep.turn_prefix_messages:
|
||||
if prep.messages_to_summarize:
|
||||
history_text, history_error = generate_summary(
|
||||
prep.messages_to_summarize, summarize_fn,
|
||||
prep.settings.reserve_tokens, model.max_tokens,
|
||||
prep.previous_summary)
|
||||
if history_error:
|
||||
return _degraded_compact(prep, history_error)
|
||||
else:
|
||||
history_text = "No prior history."
|
||||
prefix_text, prefix_error = generate_turn_prefix_summary(
|
||||
prep.turn_prefix_messages, summarize_fn,
|
||||
prep.settings.reserve_tokens, model.max_tokens)
|
||||
if prefix_error:
|
||||
return _degraded_compact(prep, prefix_error)
|
||||
summary = (f"{history_text}\n\n---\n\n"
|
||||
f"**Turn Context (split turn):**\n\n{prefix_text}")
|
||||
else:
|
||||
if not prep.messages_to_summarize:
|
||||
return None
|
||||
summary, history_error = generate_summary(
|
||||
prep.messages_to_summarize, summarize_fn,
|
||||
prep.settings.reserve_tokens, model.max_tokens,
|
||||
prep.previous_summary)
|
||||
if history_error:
|
||||
return _degraded_compact(prep, history_error)
|
||||
|
||||
read_files, modified_files = compute_file_lists(prep.file_ops)
|
||||
summary += format_file_operations(read_files, modified_files)
|
||||
|
||||
summary_msg = AgentMessage(role="user", content=summary,
|
||||
kind="compaction_summary")
|
||||
return [summary_msg] + prep.retained_tail
|
||||
|
||||
|
||||
def _degraded_compact(prep: CompactionPreparation,
|
||||
error: str) -> Optional[List[AgentMessage]]:
|
||||
"""
|
||||
已声明偏差（对照 pi: 直接返回 CompactionError）：
|
||||
桌面应用优先「不丢上下文」——摘要失败时降级为机械摘录，
|
||||
保留尾巴原样不动。
|
||||
"""
|
||||
old = prep.messages_to_summarize + prep.turn_prefix_messages
|
||||
if not old:
|
||||
return None
|
||||
excerpt = serialize_conversation(old)[-500:]
|
||||
summary = (f"(自动压缩：摘要生成失败 [{error}]，以下为旧对话尾部摘录)\n\n{excerpt}")
|
||||
summary_msg = AgentMessage(role="user", content=summary,
|
||||
kind="compaction_summary")
|
||||
return [summary_msg] + prep.retained_tail
|
||||
@@ -0,0 +1,341 @@
|
||||
"""
|
||||
core/agent/context.py
|
||||
=====================
|
||||
🌟 pi 上下文管理的 Python 1:1 移植 —— 令牌估算 + 输出预算钳制
|
||||
|
||||
对照 pi-main 源码：
|
||||
packages/ai/src/api/simple-options.ts
|
||||
- CONTEXT_SAFETY_TOKENS = 4096 (行 12)
|
||||
- clampMaxTokensToContext() (行 12-34)
|
||||
maxOutput = contextWindow - inputTokens - 4096
|
||||
return min(model.maxTokens, maxOutput) (下限 1)
|
||||
- clampOutputsToContext() (行 36-60)
|
||||
多模型共享窗口时按比例分配剩余输出预算
|
||||
packages/ai/src/api/token-utils.ts
|
||||
- 每条消息估算: ceil(chars/4) + 4
|
||||
|
||||
🌟 已声明偏差（仅 1 处，其余全部 1:1）：
|
||||
pi 按 chars/4 估算（针对英文）。中文 1 字 ≈ 1 token，chars/4 会严重低估。
|
||||
这里采用 CJK 感知估算：CJK 字符按 1 token/字，其余按 4 字符/token。
|
||||
这直接影响「压缩触发时机」和「输出预算钳制」，必须更准才不会撞 vLLM 上限。
|
||||
|
||||
pi 的 usage 锚定机制 1:1 保留：若历史里存在有效 assistant usage，
|
||||
总估算 = 该 usage 的 totalTokens（服务商精确值）+ 其后消息的逐条估算。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .types import AgentMessage, ModelConfig
|
||||
|
||||
# ======================================================================
|
||||
# 对照 simple-options.ts: export const CONTEXT_SAFETY_TOKENS = 4096
|
||||
# 用途：上下文钳制时预留的安全余量（tokenizer 估算误差 + vLLM 精确计数的偏差）
|
||||
# ======================================================================
|
||||
CONTEXT_SAFETY_TOKENS = 4096
|
||||
|
||||
# 图片 token 估算（对照 pi 对 image part 的固定估算）
|
||||
IMAGE_TOKENS = 1600
|
||||
|
||||
_CJK_RE = re.compile(
|
||||
r"[\u2e80-\u2eff\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff"
|
||||
r"\uf900-\ufaff\uff00-\uffef]"
|
||||
)
|
||||
|
||||
|
||||
def _count_text_chars(value: Any) -> int:
|
||||
"""统计内容里的字符数（对照 pi token-utils 的字符统计口径）"""
|
||||
if isinstance(value, str):
|
||||
return len(value)
|
||||
if isinstance(value, list): # OpenAI 多模态 content 数组
|
||||
n = 0
|
||||
for part in value:
|
||||
if isinstance(part, dict):
|
||||
if part.get("type") == "text":
|
||||
n += len(part.get("text", "") or "")
|
||||
elif part.get("type") == "image_url":
|
||||
n += 0 # 图片单独按 IMAGE_TOKENS 计
|
||||
else:
|
||||
n += len(str(part))
|
||||
return n
|
||||
return len(str(value or ""))
|
||||
|
||||
|
||||
def _count_image_parts(content: Any) -> int:
|
||||
if isinstance(content, list):
|
||||
return sum(1 for p in content
|
||||
if isinstance(p, dict) and p.get("type") == "image_url")
|
||||
return 0
|
||||
|
||||
|
||||
def estimate_message_tokens(message: AgentMessage) -> int:
|
||||
"""
|
||||
单条消息 token 估算（CJK 感知版）。
|
||||
pi 口径: ceil(chars/4) + 4。本项目: CJK 字 ×1 + 其余 /4 + 4 + 图片×1600。
|
||||
"""
|
||||
raw = message.content
|
||||
if isinstance(raw, str):
|
||||
chars = len(raw)
|
||||
cjk = len(_CJK_RE.findall(raw))
|
||||
else:
|
||||
# 数组 content：拼接所有文本部分
|
||||
joined = []
|
||||
for part in raw if isinstance(raw, list) else []:
|
||||
if isinstance(part, dict) and part.get("type") == "text":
|
||||
joined.append(part.get("text", "") or "")
|
||||
text = "".join(joined)
|
||||
chars = len(text)
|
||||
cjk = len(_CJK_RE.findall(text))
|
||||
tokens = cjk + math.ceil((chars - cjk) / 4) + 4
|
||||
tokens += _count_image_parts(raw) * IMAGE_TOKENS
|
||||
# 思考内容与工具参数也占上下文（assistant 的 reasoning / tool_calls）
|
||||
if message.reasoning:
|
||||
rcjk = len(_CJK_RE.findall(message.reasoning))
|
||||
tokens += rcjk + math.ceil((len(message.reasoning) - rcjk) / 4)
|
||||
for tc in message.tool_calls:
|
||||
arg_text = _json_str(tc.arguments)
|
||||
tokens += math.ceil(len(arg_text) / 4)
|
||||
# toolResult 的 content 走 content 字段（str 或 list）
|
||||
return tokens
|
||||
|
||||
|
||||
def _json_str(obj: Any) -> str:
|
||||
import json
|
||||
try:
|
||||
return json.dumps(obj, ensure_ascii=False)
|
||||
except Exception:
|
||||
return str(obj)
|
||||
|
||||
|
||||
def _estimate_text_tokens(text: str) -> int:
|
||||
"""CJK 感知文本 token 估算（与 estimate_message_tokens 同公式，不含每条 +4 开销）"""
|
||||
if not text:
|
||||
return 0
|
||||
cjk = len(_CJK_RE.findall(text))
|
||||
return cjk + math.ceil((len(text) - cjk) / 4)
|
||||
|
||||
|
||||
def _estimate_tools_tokens(tools) -> int:
|
||||
"""
|
||||
🆕 P2: 工具 schema token 估算 —— name/description/parameters 的 JSON 字符数（CJK 感知）。
|
||||
对照 pi Context 分支: all tools = Σ estimate_text(json(tool))。
|
||||
"""
|
||||
total = 0
|
||||
for t in tools or []:
|
||||
schema = {
|
||||
"name": getattr(t, "name", "") or "",
|
||||
"description": getattr(t, "description", "") or "",
|
||||
"parameters": getattr(t, "parameters", None),
|
||||
}
|
||||
total += _estimate_text_tokens(_json_str(schema))
|
||||
return total
|
||||
|
||||
|
||||
def estimate_context_tokens(messages: List[AgentMessage],
|
||||
system_prompt: str = "",
|
||||
tools: Optional[List[Any]] = None
|
||||
) -> "ContextUsageEstimate":
|
||||
"""
|
||||
🌟 1:1 对照 pi estimateContextTokens（usage 锚定）：
|
||||
- 找到最后一条带有效 usage 的 assistant 消息（🆕 P0: 含 timestamp 时效校验）
|
||||
- 总 tokens = 该 usage 的精确值 + 其后消息的逐条估算
|
||||
- 无 usage 时 = 全部消息逐条估算
|
||||
🆕 P2: 可选 system_prompt / tools 参数（对照 pi Context 分支）：
|
||||
- 仅在无锚点分支计入 system 提示词 + 工具 schema 的估算
|
||||
（有锚点时 usage 已是含 system+tools 的精确总值，不重复计）
|
||||
逐条估算用 CJK 感知版 estimate_message_tokens（已声明偏差）。
|
||||
"""
|
||||
last_idx, last_usage = _find_last_usage(messages)
|
||||
|
||||
sys_tokens = _estimate_text_tokens(system_prompt) if system_prompt else 0
|
||||
tool_tokens = _estimate_tools_tokens(tools) if tools else 0
|
||||
|
||||
if last_idx is None:
|
||||
estimated = (sum(estimate_message_tokens(m) for m in messages)
|
||||
+ sys_tokens + tool_tokens)
|
||||
return ContextUsageEstimate(tokens=estimated, usage_tokens=0,
|
||||
trailing_tokens=estimated, last_usage_index=None)
|
||||
|
||||
usage_tokens = calculate_context_tokens(last_usage)
|
||||
trailing = sum(estimate_message_tokens(m) for m in messages[last_idx + 1:])
|
||||
return ContextUsageEstimate(tokens=usage_tokens + trailing,
|
||||
usage_tokens=usage_tokens,
|
||||
trailing_tokens=trailing,
|
||||
last_usage_index=last_idx)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ContextUsageEstimate:
|
||||
"""
|
||||
🌟 1:1 对照 pi estimateContextTokens 的返回值
|
||||
packages/agent/src/harness/compaction/compaction.ts ContextUsageEstimate:
|
||||
tokens 总估算
|
||||
usageTokens 最后一条有效 assistant usage 报告的精确 token 数
|
||||
trailingTokens 该 usage 之后消息的估算值
|
||||
lastUsageIndex 提供 usage 的消息下标（无则 None）
|
||||
"""
|
||||
tokens: int = 0
|
||||
usage_tokens: int = 0
|
||||
trailing_tokens: int = 0
|
||||
last_usage_index: Optional[int] = None
|
||||
|
||||
|
||||
def calculate_context_tokens(usage: Dict[str, Any]) -> int:
|
||||
"""pi calculateContextTokens: totalTokens || input+output+cacheRead+cacheWrite"""
|
||||
if not usage:
|
||||
return 0
|
||||
total = int(usage.get("totalTokens", 0) or 0)
|
||||
if total:
|
||||
return total
|
||||
return (int(usage.get("input", 0) or 0) + int(usage.get("output", 0) or 0)
|
||||
+ int(usage.get("cacheRead", 0) or 0)
|
||||
+ int(usage.get("cacheWrite", 0) or 0))
|
||||
|
||||
|
||||
def _get_assistant_usage(msg: AgentMessage) -> Dict[str, Any]:
|
||||
"""
|
||||
pi getAssistantUsage：只认「有效」的 assistant usage——
|
||||
stopReason 不是 aborted/error，且 calculateContextTokens > 0。
|
||||
（🆕 P0: 锚点扫描还需通过 timestamp 时效校验，见 _find_last_usage）
|
||||
"""
|
||||
if msg.role != "assistant":
|
||||
return {}
|
||||
if msg.stop_reason in ("aborted", "error"):
|
||||
return {}
|
||||
usage = msg.usage or {}
|
||||
if usage and calculate_context_tokens(usage) > 0:
|
||||
return usage
|
||||
return {}
|
||||
|
||||
|
||||
def _row_block_tool_flags(messages: List["AgentMessage"]) -> List[bool]:
|
||||
"""
|
||||
🆕 G1: 逐条标记「其 DB 行（同非空 db_msg_id 的连续块）是否含工具活动」。
|
||||
含工具输出的行，其入库 usage 是「本轮最后一个子请求」的内存快照
|
||||
（工具结果未受 4000 字回放截断），不等于下一请求（截断回放）→ 锚点失效。
|
||||
轮中内存条目（db_msg_id=""）一律 False：保持原规则
|
||||
（轮中条目的内存载荷 == 下一子请求载荷，锚点有效）。
|
||||
"""
|
||||
n = len(messages)
|
||||
flags = [False] * n
|
||||
i = 0
|
||||
while i < n:
|
||||
rid = getattr(messages[i], "db_msg_id", "") or ""
|
||||
if not rid:
|
||||
i += 1
|
||||
continue
|
||||
j = i
|
||||
has_tool = False
|
||||
while j < n and (getattr(messages[j], "db_msg_id", "") or "") == rid:
|
||||
mj = messages[j]
|
||||
if mj.role == "toolResult" or (mj.role == "assistant" and mj.tool_calls):
|
||||
has_tool = True
|
||||
j += 1
|
||||
for k in range(i, j):
|
||||
flags[k] = has_tool
|
||||
i = j
|
||||
return flags
|
||||
|
||||
|
||||
def _find_last_usage(messages: List["AgentMessage"]):
|
||||
"""
|
||||
🆕 P0: 1:1 对照 pi getLastAssistantUsageInfo（含锚点失效校验）：
|
||||
assistant 的 usage 要成为锚点，必须满足
|
||||
① timestamp >= 它前面所有消息的最大 timestamp（即它不比任何前消息「更旧」）。
|
||||
场景：压缩后 state = [新摘要消息(T_new)] + [保留的旧消息(T_old)]，
|
||||
保留窗里 assistant 的 usage 来自压缩前（如 60K 精确值）——不校验的话，
|
||||
下一次 should_compact 会锚定到过期值 → 误触发二次压缩（摘要套摘要）。
|
||||
摘要消息 timestamp 晚于全部保留消息 → 前缀游标抬高 → 旧 usage 自动失效；
|
||||
压缩后第一个新 assistant 回复到达 → 新 usage 锚点自动恢复。
|
||||
② 🆕 G1: 该 usage 所在 DB 行不含工具活动（toolResult / assistant.tool_calls）。
|
||||
含工具输出的行，入库 usage = 本轮内存快照（工具结果未截断）≠ 下一请求
|
||||
（4000 字截断回放）→ 系统性高估 → 误触发压缩（实测 113904 锚点 vs
|
||||
真实下一请求 ~90k）。纯文本行 / 轮中内存条目的 usage 仍是精确值，保留锚定。
|
||||
（消息无 timestamp 信息（全 0）时行为等同旧版「取最后一条有效」，完全向后兼容。）
|
||||
"""
|
||||
latest_prefix_ts = -1 # 对照 pi: Number.NEGATIVE_INFINITY
|
||||
last_idx = None
|
||||
last_usage: Dict[str, Any] = {}
|
||||
row_tool_flags = _row_block_tool_flags(messages)
|
||||
for i, msg in enumerate(messages):
|
||||
if msg.role == "assistant":
|
||||
applies = (msg.timestamp or 0) >= latest_prefix_ts
|
||||
if (applies
|
||||
and not row_tool_flags[i]
|
||||
and msg.stop_reason not in ("aborted", "error")
|
||||
and msg.usage
|
||||
and calculate_context_tokens(msg.usage) > 0):
|
||||
last_idx, last_usage = i, msg.usage
|
||||
ts = msg.timestamp or 0
|
||||
if ts > latest_prefix_ts:
|
||||
latest_prefix_ts = ts
|
||||
return last_idx, last_usage
|
||||
|
||||
|
||||
def clamp_max_tokens_to_context(model: ModelConfig,
|
||||
context: List[AgentMessage],
|
||||
system_prompt: str = "",
|
||||
tools: Optional[List[Any]] = None
|
||||
) -> Optional[Tuple[int, int]]:
|
||||
"""
|
||||
🌟 对照 simple-options.ts:12-34 clampMaxTokensToContext（1:1 公式）：
|
||||
|
||||
input_estimate = estimateContextTokens(context) # 🆕 P2: 含 system+tools（pi Context 实参）
|
||||
if (input_estimate > model.contextWindow) return null // 上下文已溢出
|
||||
maxOutput = contextWindow - inputTokens - CONTEXT_SAFETY_TOKENS
|
||||
return max(1, min(model.maxTokens, maxOutput))
|
||||
|
||||
这就是 pi 约束公式的本体：
|
||||
输入 + 输出 ≤ contextWindow − 4096 ≤ vLLM max_model_len
|
||||
🆕 P2: system_prompt / tools 仅在无 usage 锚点分支计入（首轮流式精确记账，
|
||||
与 pi 传完整 Context 的行为一致）；有锚点时 usage 已精确覆盖，不重复计。
|
||||
返回 (max_tokens, input_tokens)；输入已溢出窗口时返回 None（由恢复逻辑接管）。
|
||||
"""
|
||||
input_tokens = estimate_context_tokens(
|
||||
context, system_prompt=system_prompt, tools=tools).tokens
|
||||
if input_tokens > model.context_window:
|
||||
return None
|
||||
max_output = model.context_window - input_tokens - CONTEXT_SAFETY_TOKENS
|
||||
return max(1, min(model.max_tokens, max_output)), input_tokens
|
||||
|
||||
|
||||
def clamp_outputs_to_context(context: List[AgentMessage],
|
||||
models: List[ModelConfig],
|
||||
reserve_tokens: int = CONTEXT_SAFETY_TOKENS
|
||||
) -> List[int]:
|
||||
"""
|
||||
对照 simple-options.ts:36-60 clampOutputsToContext：
|
||||
多个模型共享同一窗口时，把剩余输出预算按比例分配给各模型。
|
||||
（haocode 单模型场景用不到，保留以求框架完整）
|
||||
"""
|
||||
input_tokens = estimate_context_tokens(context).tokens
|
||||
available = max(0, sum(m.context_window for m in models) - input_tokens - reserve_tokens)
|
||||
if not models:
|
||||
return []
|
||||
weights = [m.max_tokens for m in models]
|
||||
total_w = sum(weights) or 1
|
||||
alloc = [max(1, int(available * w / total_w)) for w in weights]
|
||||
# 各自不超过自身 max_tokens
|
||||
alloc = [min(a, m.max_tokens) for a, m in zip(alloc, models)]
|
||||
return alloc
|
||||
|
||||
|
||||
def should_compact(messages: List[AgentMessage], model: ModelConfig,
|
||||
reserve_tokens: int = 16384,
|
||||
system_prompt: str = "",
|
||||
tools: Optional[List[Any]] = None) -> Tuple[bool, int]:
|
||||
"""
|
||||
🌟 1:1 对照 pi compaction.ts shouldCompact（harness 原版公式）：
|
||||
contextTokens > contextWindow - reserveTokens
|
||||
reserveTokens 默认 16384（DEFAULT_COMPACTION_SETTINGS.reserveTokens），
|
||||
即「为摘要提示词与输出预留 16K」。估算值取 usage 锚定估算。
|
||||
🆕 G2: system_prompt/tools 透传 —— 无锚点分支（工具行快照锚点被 G1 失效后
|
||||
更常走到）也要计入 system + 工具 schema，度量「下一请求真实载荷」，
|
||||
与显示端（update_context_display）同口径。
|
||||
"""
|
||||
tokens = estimate_context_tokens(messages, system_prompt=system_prompt,
|
||||
tools=tools).tokens
|
||||
return (tokens > model.context_window - reserve_tokens), tokens
|
||||
@@ -0,0 +1,482 @@
|
||||
"""
|
||||
core/agent/loop.py
|
||||
==================
|
||||
🌟 pi 核心循环的 Python 1:1 移植 —— 整个框架的心脏
|
||||
|
||||
对照 pi-main 源码（packages/agent/src/agent-loop.ts, 796 行）：
|
||||
runAgentLoop() → run_loop(agent, new_message, ...)
|
||||
runLoop() → _run_loop() (行 163-278)
|
||||
streamAssistantResponse() → _stream_turn() (行 281-380)
|
||||
executeToolCalls() → execute_tool_calls() (行 413-427)
|
||||
executeToolCallsParallel/Sequential → _execute_parallel/_execute_sequential
|
||||
prepareToolCall() → tools.prepare_tool_calls
|
||||
failToolCallsFromTruncatedMessage → tools.fail_tool_calls_from_truncated_message
|
||||
shouldTerminateToolBatch → _should_terminate_batch (行 561-563)
|
||||
|
||||
pi runLoop 的结构（本文件逐行对应）：
|
||||
emit agent_start
|
||||
pending = getSteeringMessages() # 循环开始前先取一次
|
||||
outer: while True:
|
||||
inner: while hasMoreToolCalls or pending:
|
||||
turn_start（首轮不发，首轮 turn_start 在 runAgentLoop 入口发）
|
||||
注入 pending（message_start/end → 入 context + newMessages）
|
||||
assistant = streamAssistantResponse(...)
|
||||
if stopReason in (error, aborted): turn_end + agent_end + return
|
||||
toolCalls = assistant 的 toolCall 块
|
||||
if toolCalls:
|
||||
length → failToolCallsFromTruncatedMessage（不执行！）
|
||||
否则 → executeToolCalls（sequential/parallel 二选一）
|
||||
hasMoreToolCalls = !batch.terminate
|
||||
toolResults 入 context + newMessages
|
||||
turn_end
|
||||
prepareNextTurn 钩子（可换 context/model）
|
||||
shouldStopAfterTurn 钩子 → agent_end + return
|
||||
pending = getSteeringMessages() # 每轮结束取一次
|
||||
followUps = getFollowUpMessages()
|
||||
if followUps: pending = followUps; continue # 外层续跑
|
||||
break
|
||||
agent_end
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
from .context import clamp_max_tokens_to_context
|
||||
from .tools import (PreparedToolCall, execute_tool_call, fail_tool_calls_from_truncated_message,
|
||||
prepare_tool_calls)
|
||||
from .types import (AgentConfig, AgentError, AgentEvent, AgentMessage,
|
||||
AgentTool, AgentToolResult, AbortSignal, AssistantMessageEvent,
|
||||
RunResult, ToolCall, new_id)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 事件发射辅助
|
||||
# ======================================================================
|
||||
def _emit(agent, event: AgentEvent):
|
||||
agent._emit(event)
|
||||
|
||||
|
||||
def _message_events(agent, msg: AgentMessage):
|
||||
_emit(agent, AgentEvent(type="message_start", message=msg))
|
||||
_emit(agent, AgentEvent(type="message_end", message=msg))
|
||||
|
||||
|
||||
def _tool_result_message(finalized: Dict[str, Any]) -> AgentMessage:
|
||||
"""对照 createToolResultMessage"""
|
||||
tc: ToolCall = finalized["tool_call"]
|
||||
result: AgentToolResult = finalized["result"]
|
||||
return AgentMessage(
|
||||
role="toolResult",
|
||||
tool_call_id=tc.id,
|
||||
tool_name=tc.name,
|
||||
content=result.content,
|
||||
is_error=finalized.get("is_error", False) or result.is_error,
|
||||
)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 流式生成一轮助手消息 —— 对照 streamAssistantResponse (行 281-380)
|
||||
# ======================================================================
|
||||
def _stream_turn(agent, current_context: List[AgentMessage],
|
||||
config: AgentConfig, signal: AbortSignal,
|
||||
stream_fn: Callable) -> Tuple[AgentMessage, Optional[AgentError]]:
|
||||
"""
|
||||
返回 (assistant_message, error)。
|
||||
error 非 None 时 message.stop_reason == "error"。
|
||||
"""
|
||||
# 🌟 pi 同款：流开始前检查中止 → 立即产出 aborted 消息
|
||||
if signal.aborted:
|
||||
msg = AgentMessage(role="assistant", stop_reason="aborted")
|
||||
agent.state.messages.append(msg)
|
||||
_message_events(agent, msg)
|
||||
return msg, None
|
||||
|
||||
# 🌟 系统提示词注入（对照 pi：systemPrompt 放在每次 API 请求头部，
|
||||
# 不进 state.messages、不占压缩/历史）
|
||||
if config.system_prompt:
|
||||
api_context: List[AgentMessage] = [
|
||||
AgentMessage(role="system", content=config.system_prompt)
|
||||
] + current_context
|
||||
else:
|
||||
api_context = current_context
|
||||
|
||||
# 🌟 输出预算钳制（对照 simple-options clampMaxTokensToContext）
|
||||
# 🆕 P2: 把 system + 工具 schema 传给估算器（对照 pi 传完整 Context），
|
||||
# 仅在无 usage 锚点分支生效；否则新会话首轮会多预留 ~system+tools 的预算。
|
||||
clamped = clamp_max_tokens_to_context(
|
||||
config.model, current_context,
|
||||
system_prompt=config.system_prompt or "", tools=config.tools)
|
||||
if clamped is None:
|
||||
err = AgentError(
|
||||
message="上下文溢出：估算输入 token 已超过模型窗口（需要压缩）",
|
||||
kind="overflow", recoverable=True)
|
||||
msg = AgentMessage(role="assistant", stop_reason="error",
|
||||
error_message=err.message)
|
||||
agent.state.messages.append(msg)
|
||||
_message_events(agent, msg)
|
||||
return msg, err
|
||||
max_tokens, _input_tokens = clamped
|
||||
|
||||
msg = AgentMessage(role="assistant")
|
||||
agent.state.streaming_message = msg
|
||||
agent.state.streaming_delta = {}
|
||||
_emit(agent, AgentEvent(type="message_start", message=msg))
|
||||
|
||||
raw_tc: Dict[int, Dict[str, str]] = {} # 中止时用于保留部分工具调用（仅展示）
|
||||
got_final = False
|
||||
error: Optional[AgentError] = None
|
||||
|
||||
try:
|
||||
for kind, payload in stream_fn(api_context, config.model, signal,
|
||||
max_tokens, config.tools):
|
||||
if kind == "event":
|
||||
ev = payload
|
||||
if ev.type == "text_delta":
|
||||
msg.content = (msg.content or "") + ev.text
|
||||
agent.state.streaming_delta["text"] = ev.text
|
||||
_emit(agent, AgentEvent(type="message_update", message=msg,
|
||||
assistant_message_event=ev))
|
||||
elif ev.type == "thinking_delta":
|
||||
msg.reasoning += ev.text
|
||||
agent.state.streaming_delta["thinking"] = ev.text
|
||||
_emit(agent, AgentEvent(type="message_update", message=msg,
|
||||
assistant_message_event=ev))
|
||||
elif ev.type == "toolcall_delta":
|
||||
slot = raw_tc.setdefault(ev.tool_call_index,
|
||||
{"id": "", "name": "", "args": ""})
|
||||
if ev.tool_call_field == "id":
|
||||
slot["id"] = ev.tool_call_delta
|
||||
elif ev.tool_call_field == "name":
|
||||
slot["name"] += ev.tool_call_delta
|
||||
else:
|
||||
slot["args"] += ev.tool_call_delta
|
||||
_emit(agent, AgentEvent(type="message_update", message=msg,
|
||||
assistant_message_event=ev))
|
||||
else: # final
|
||||
final: AgentMessage = payload
|
||||
got_final = True
|
||||
# 采用权威 final：usage / stop_reason / 解析好的 tool_calls
|
||||
msg.usage = final.usage
|
||||
msg.stop_reason = final.stop_reason
|
||||
msg.tool_calls = final.tool_calls
|
||||
if not final.content and msg.content:
|
||||
pass # 以循环侧累积为准（二者一致）
|
||||
else:
|
||||
msg.content = final.content if final.content else msg.content
|
||||
msg.reasoning = final.reasoning or msg.reasoning
|
||||
except AgentError as e:
|
||||
error = e
|
||||
msg.stop_reason = "error"
|
||||
msg.error_message = e.message
|
||||
_emit(agent, AgentEvent(type="message_update", message=msg,
|
||||
assistant_message_event=
|
||||
AssistantMessageEvent.error(e.message)))
|
||||
except Exception as e:
|
||||
from .stream_fn import classify_error
|
||||
error = classify_error(e)
|
||||
msg.stop_reason = "error"
|
||||
msg.error_message = error.message
|
||||
|
||||
if not got_final and error is None:
|
||||
# 生成器中途结束且无异常 = 中止（流被关闭，pi 同款语义）
|
||||
msg.stop_reason = "aborted"
|
||||
# 保留部分工具调用（仅用于 UI 展示；中止后不会执行，也不会进入下次上下文）
|
||||
for idx in sorted(raw_tc.keys()):
|
||||
slot = raw_tc[idx]
|
||||
if slot.get("name"):
|
||||
msg.tool_calls.append(ToolCall(id=slot.get("id") or new_id("call"),
|
||||
name=slot["name"],
|
||||
raw_arguments=slot.get("args", "")))
|
||||
if error is None and msg.stop_reason not in ("stop", "length", "aborted"):
|
||||
msg.stop_reason = "stop"
|
||||
|
||||
# 🌟 兜底（haocode 扩展，pi 无此层）：不支持 tools API 的供应商/模型会把
|
||||
# 工具调用用文字"演"出来（如 <bash>ls</bash>）。识别单参工具 bash/read
|
||||
# 转成真 tool_call 继续执行；write/edit 多参歧义大不兜底。
|
||||
if error is None and not msg.tool_calls and msg.content:
|
||||
from .tools import parse_text_tool_calls
|
||||
cleaned, txt_calls = parse_text_tool_calls(msg.content)
|
||||
if txt_calls:
|
||||
msg.tool_calls = txt_calls
|
||||
|
||||
agent.state.streaming_message = None
|
||||
agent.state.streaming_delta = {}
|
||||
agent.state.messages.append(msg)
|
||||
_emit(agent, AgentEvent(type="message_end", message=msg))
|
||||
return msg, error
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 工具批量执行 —— 对照 executeToolCalls (行 413-427)
|
||||
# ======================================================================
|
||||
def _should_terminate_batch(finalized: List[Dict[str, Any]]) -> bool:
|
||||
"""对照 shouldTerminateToolBatch: 全部结果都请求 terminate 才终止"""
|
||||
return len(finalized) > 0 and all(
|
||||
f["result"].terminate for f in finalized)
|
||||
|
||||
|
||||
def _run_prepared_in_pool(prep: PreparedToolCall, assistant: AgentMessage,
|
||||
config: AgentConfig, signal: AbortSignal,
|
||||
agent) -> Dict[str, Any]:
|
||||
"""线程池内执行单个工具（对照 parallel 版的 async 闭包）"""
|
||||
on_update = None
|
||||
on_timer = None
|
||||
if agent is not None:
|
||||
def on_update(partial: str):
|
||||
_emit(agent, AgentEvent(type="tool_execution_update",
|
||||
tool_call=prep.tool_call, arg=partial))
|
||||
def on_timer(elapsed_i: int, timeout_i: int):
|
||||
# 🆕 bash 运行中每秒滴一次 → 前端气泡读秒 N/Ts
|
||||
_emit(agent, AgentEvent(type="tool_execution_timer",
|
||||
tool_call=prep.tool_call,
|
||||
arg=(int(elapsed_i), int(timeout_i))))
|
||||
result = execute_tool_call(prep, assistant, config, signal, on_update,
|
||||
on_timer)
|
||||
finalized = {"tool_call": prep.tool_call, "result": result,
|
||||
"is_error": result.is_error}
|
||||
_emit(agent, AgentEvent(type="tool_execution_end",
|
||||
tool_call=prep.tool_call, result=result,
|
||||
is_error=result.is_error))
|
||||
return finalized
|
||||
|
||||
|
||||
def _execute_parallel(current_context, assistant, config, signal, agent,
|
||||
prepared: List[PreparedToolCall]) -> Dict[str, Any]:
|
||||
"""对照 executeToolCallsParallel：start 事件串行发、准备串行、执行并发、
|
||||
结果消息按原始顺序产出"""
|
||||
finalized: List[Dict[str, Any]] = []
|
||||
pending_futures: List = []
|
||||
max_workers = max(1, min(8, len(prepared)))
|
||||
with ThreadPoolExecutor(max_workers=max_workers,
|
||||
thread_name_prefix="tool") as pool:
|
||||
for prep in prepared:
|
||||
_emit(agent, AgentEvent(type="tool_execution_start",
|
||||
tool_call=prep.tool_call,
|
||||
arg=str(prep.args) if prep.args else ""))
|
||||
if prep.error:
|
||||
# 准备失败（未知工具/参数非法）→ 立即结果（对照 kind:"immediate"）
|
||||
result = AgentToolResult.text(prep.error, is_error=True)
|
||||
fin = {"tool_call": prep.tool_call, "result": result,
|
||||
"is_error": True}
|
||||
_emit(agent, AgentEvent(type="tool_execution_end",
|
||||
tool_call=prep.tool_call, result=result,
|
||||
is_error=True))
|
||||
finalized.append(fin)
|
||||
if signal.aborted:
|
||||
break
|
||||
continue
|
||||
fut = pool.submit(_run_prepared_in_pool, prep, assistant, config,
|
||||
signal, agent)
|
||||
pending_futures.append(fut)
|
||||
if signal.aborted:
|
||||
# 已提交的任务仍会完成（它们内部检查 signal），不再追加
|
||||
pass
|
||||
for fut in pending_futures:
|
||||
finalized.append(fut.result())
|
||||
|
||||
messages: List[AgentMessage] = []
|
||||
for fin in finalized:
|
||||
tr = _tool_result_message(fin)
|
||||
_message_events(agent, tr)
|
||||
messages.append(tr)
|
||||
return {"messages": messages, "terminate": _should_terminate_batch(finalized)}
|
||||
|
||||
|
||||
def _execute_sequential(current_context, assistant, config, signal, agent,
|
||||
prepared: List[PreparedToolCall]) -> Dict[str, Any]:
|
||||
"""对照 executeToolCallsSequential：一次一个，完成一个再下一个"""
|
||||
finalized: List[Dict[str, Any]] = []
|
||||
messages: List[AgentMessage] = []
|
||||
for prep in prepared:
|
||||
_emit(agent, AgentEvent(type="tool_execution_start",
|
||||
tool_call=prep.tool_call,
|
||||
arg=str(prep.args) if prep.args else ""))
|
||||
if not prep.error:
|
||||
fin = _run_prepared_in_pool(prep, assistant, config, signal, agent)
|
||||
else:
|
||||
result = AgentToolResult.text(prep.error, is_error=True)
|
||||
fin = {"tool_call": prep.tool_call, "result": result, "is_error": True}
|
||||
_emit(agent, AgentEvent(type="tool_execution_end",
|
||||
tool_call=prep.tool_call, result=result,
|
||||
is_error=True))
|
||||
finalized.append(fin)
|
||||
tr = _tool_result_message(fin)
|
||||
_message_events(agent, tr)
|
||||
messages.append(tr)
|
||||
if signal.aborted:
|
||||
break
|
||||
return {"messages": messages, "terminate": _should_terminate_batch(finalized)}
|
||||
|
||||
|
||||
def execute_tool_calls(current_context, assistant: AgentMessage,
|
||||
config: AgentConfig, signal: AbortSignal, agent) -> Dict[str, Any]:
|
||||
"""对照 executeToolCalls: 批次里有 sequential 工具 → 整批串行"""
|
||||
prepared = prepare_tool_calls(assistant, config.tools)
|
||||
has_sequential = any(
|
||||
(p.tool and p.tool.execution_mode == "sequential") for p in prepared
|
||||
if not p.error)
|
||||
if config.tool_execution == "sequential" or has_sequential:
|
||||
return _execute_sequential(current_context, assistant, config, signal,
|
||||
agent, prepared)
|
||||
return _execute_parallel(current_context, assistant, config, signal, agent,
|
||||
prepared)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 主循环 —— 对照 runLoop (行 163-278)
|
||||
# ======================================================================
|
||||
def run_loop(agent, new_message: Optional[AgentMessage], signal: AbortSignal,
|
||||
stream_fn: Callable) -> RunResult:
|
||||
config = agent.config
|
||||
new_messages: List[AgentMessage] = []
|
||||
|
||||
# 入口校验（对照 runAgentLoopContinue 的前置检查由 Agent 层负责）
|
||||
_emit(agent, AgentEvent(type="agent_start"))
|
||||
_emit(agent, AgentEvent(type="turn_start")) # 首轮 turn_start（对照行 142/170）
|
||||
|
||||
# context 准备 + transformContext 钩子
|
||||
base = list(agent.state.messages)
|
||||
if config.transform_context:
|
||||
try:
|
||||
base = config.transform_context(base) or base
|
||||
except Exception:
|
||||
pass
|
||||
current_context = list(base)
|
||||
|
||||
first_turn = True
|
||||
pending: List[AgentMessage] = agent._take_steering()
|
||||
|
||||
while True: # outer loop
|
||||
has_more_tool_calls = True
|
||||
|
||||
while has_more_tool_calls or pending: # inner loop
|
||||
if not first_turn:
|
||||
_emit(agent, AgentEvent(type="turn_start"))
|
||||
else:
|
||||
first_turn = False
|
||||
|
||||
# ---- 注入 pending 消息（steering / followUp）----
|
||||
if pending:
|
||||
for m in pending:
|
||||
_message_events(agent, m)
|
||||
current_context.append(m)
|
||||
new_messages.append(m)
|
||||
agent.state.messages.append(m)
|
||||
pending = []
|
||||
|
||||
# ---- 🆕 轮中主动压缩检查（haocode 增强，偏离 pi 1:1）----
|
||||
# 单条工具输出可能把上下文顶出窗口；发下一次请求前主动检查
|
||||
#（与轮首 should_compact 同公式）。compact_fn 返回新列表
|
||||
#（发生了压缩）→ 同步循环局部上下文。
|
||||
if config.compact_fn is not None and agent.state.messages:
|
||||
try:
|
||||
compacted = config.compact_fn(list(agent.state.messages))
|
||||
if compacted is not None:
|
||||
agent.state.messages = compacted
|
||||
current_context = list(compacted)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---- 流式生成助手消息 ----
|
||||
assistant, error = _stream_turn(agent, current_context, config,
|
||||
signal, stream_fn)
|
||||
new_messages.append(assistant)
|
||||
current_context.append(assistant)
|
||||
|
||||
if assistant.stop_reason in ("error", "aborted"):
|
||||
agent.state.error = error
|
||||
_emit(agent, AgentEvent(type="turn_end", message=assistant))
|
||||
_emit(agent, AgentEvent(type="agent_end",
|
||||
stop_reason=assistant.stop_reason,
|
||||
error=error,
|
||||
messages=new_messages))
|
||||
agent._finish_run(new_messages, assistant.stop_reason, error)
|
||||
return RunResult(stop_reason=assistant.stop_reason, error=error,
|
||||
message_count=len(new_messages))
|
||||
|
||||
# ---- 工具调用 ----
|
||||
tool_results: List[AgentMessage] = []
|
||||
has_more_tool_calls = False
|
||||
if assistant.tool_calls:
|
||||
if assistant.stop_reason == "length":
|
||||
# 🌟 截断保护（对照行 212-213）：参数可能残缺，一律失败不执行
|
||||
fail_msgs = fail_tool_calls_from_truncated_message(
|
||||
assistant, "length")
|
||||
# 事件流与正常执行对齐
|
||||
for tc, tr in zip(assistant.tool_calls, fail_msgs):
|
||||
_emit(agent, AgentEvent(type="tool_execution_start",
|
||||
tool_call=tc))
|
||||
res = AgentToolResult.text(
|
||||
f'工具调用 "{tc.name}" 未执行：输出达到 token 上限，'
|
||||
f'参数可能被截断。请用完整参数重新发起。', is_error=True)
|
||||
_emit(agent, AgentEvent(type="tool_execution_end",
|
||||
tool_call=tc, result=res,
|
||||
is_error=True))
|
||||
_message_events(agent, tr)
|
||||
batch = {"messages": fail_msgs, "terminate": False}
|
||||
else:
|
||||
batch = execute_tool_calls(current_context, assistant, config,
|
||||
signal, agent)
|
||||
tool_results = batch["messages"]
|
||||
has_more_tool_calls = not batch["terminate"]
|
||||
for r in tool_results:
|
||||
current_context.append(r)
|
||||
new_messages.append(r)
|
||||
agent.state.messages.append(r)
|
||||
|
||||
_emit(agent, AgentEvent(type="turn_end", message=assistant))
|
||||
|
||||
# ---- prepareNextTurn 钩子 ----
|
||||
if config.prepare_next_turn:
|
||||
try:
|
||||
snap = config.prepare_next_turn({
|
||||
"message": assistant, "tool_results": tool_results,
|
||||
"context": current_context, "new_messages": new_messages,
|
||||
})
|
||||
if snap:
|
||||
current_context = snap.get("context") or current_context
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---- shouldStopAfterTurn 钩子 ----
|
||||
if config.should_stop_after_turn:
|
||||
try:
|
||||
if config.should_stop_after_turn({
|
||||
"message": assistant, "tool_results": tool_results,
|
||||
"context": current_context,
|
||||
"new_messages": new_messages}):
|
||||
stop = _last_assistant_stop(new_messages)
|
||||
_emit(agent, AgentEvent(type="agent_end",
|
||||
stop_reason=stop,
|
||||
messages=new_messages))
|
||||
agent._finish_run(new_messages, stop, None)
|
||||
return RunResult(stop_reason=stop,
|
||||
message_count=len(new_messages))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---- 每轮结束取 steering ----
|
||||
pending = agent._take_steering()
|
||||
|
||||
# ---- 外层：followUp ----
|
||||
follow_ups = agent._take_follow_ups()
|
||||
if follow_ups:
|
||||
pending = follow_ups
|
||||
continue
|
||||
break
|
||||
|
||||
stop = _last_assistant_stop(new_messages)
|
||||
_emit(agent, AgentEvent(type="agent_end", stop_reason=stop,
|
||||
messages=new_messages))
|
||||
agent._finish_run(new_messages, stop, None)
|
||||
return RunResult(stop_reason=stop, message_count=len(new_messages))
|
||||
|
||||
|
||||
def _last_assistant_stop(new_messages: List[AgentMessage]) -> str:
|
||||
"""对照 pi: agent_end 不携带 stopReason，会话层从最后一条 assistant 消息读取"""
|
||||
for m in reversed(new_messages):
|
||||
if m.role == "assistant":
|
||||
return m.stop_reason or "stop"
|
||||
return "stop"
|
||||
@@ -0,0 +1,519 @@
|
||||
"""
|
||||
core/agent/recovery.py
|
||||
======================
|
||||
🌟 pi 会话层恢复逻辑的 Python 1:1 移植
|
||||
|
||||
对照 pi-main 源码：
|
||||
packages/coding-agent/src/core/agent-session.ts
|
||||
_runAgentPrompt (行 1074): prompt → while _handlePostAgentRun(): continue()
|
||||
_handlePostAgentRun (行 1088): 基于「最后一条 assistant 消息」的三路决策
|
||||
_prepareRetry (行 2811): 计数 → 移除坏消息 → 退避睡眠
|
||||
_isRetryableError (行 2770): 溢出交给压缩；其余走重试判定
|
||||
_checkCompaction (行 ~2034): 溢出压缩恢复（只试一次）
|
||||
packages/ai/src/utils/retry.ts
|
||||
isRetryableAssistantError (行 223) + 可重试/不可重试正则模式表
|
||||
退避: baseDelayMs × 2^(attempt-1)（默认 2s, 4s, 8s；maxRetries 默认 3）
|
||||
packages/ai/src/utils/overflow.ts
|
||||
isContextOverflow (行 134)：三种溢出识别
|
||||
isRecoverableLength：length 且 output < 原始输出上限
|
||||
|
||||
🌟 决策表（与 pi 完全一致，基于最后一条 assistant 消息）：
|
||||
(1) stop_reason=="error"：
|
||||
a. 是上下文溢出 → 删除坏消息 + 压缩 + continue（只试一次）
|
||||
b. 命中可重试模式（429/5xx/超时/断连…且非配额耗尽）
|
||||
且未超 maxRetries → 删除坏消息 + 退避(2s×2^(n-1)) + continue
|
||||
c. 其他 → 结束
|
||||
(2) isRecoverableLength（length 且实际输出 < 原始 maxTokens，
|
||||
说明是被上下文窗口挤断而非输出上限截断）且未尝试过
|
||||
→ 删除该助手消息 + 压缩 + continue（只试一次）
|
||||
(3) 队列里还有 followUp/steering → continue（续跑）
|
||||
(4) 都不满足 → settle
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Callable, Optional
|
||||
|
||||
from .compaction import CompactionSettings, _content_text, compact_context
|
||||
from .context import estimate_context_tokens, should_compact
|
||||
from .types import (AgentMessage, AgentConfig, RetryConfig, RunResult)
|
||||
|
||||
|
||||
# ========== 压缩诊断日志（与 main_window.diag_log 同一规范：
|
||||
# [HH:MM:SS.mmm] 时间戳 + 项目根目录追加 + UTF-8 + 失败静默） ==========
|
||||
_COMPACT_LOG_PATH = os.path.join(
|
||||
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
|
||||
"compaction_diag.log")
|
||||
|
||||
|
||||
def compact_diag_log(line: str):
|
||||
"""压缩日志只记录：同一行内容 print（控制台）+ 追加写 compaction_diag.log"""
|
||||
t = time.time()
|
||||
stamp = (f"[{time.strftime('%H:%M:%S', time.localtime(t))}"
|
||||
f".{int(t * 1000) % 1000:03d}]")
|
||||
print(f"{stamp} {line}", flush=True)
|
||||
try:
|
||||
with open(_COMPACT_LOG_PATH, "a", encoding="utf-8") as f:
|
||||
f.write(f"{stamp} {line}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _settings_of(cfg: AgentConfig) -> CompactionSettings:
|
||||
"""AgentConfig 压缩参数 → CompactionSettings（1:1 映射 pi settings 字段）"""
|
||||
return CompactionSettings(reserve_tokens=cfg.compaction_reserve,
|
||||
keep_recent_tokens=cfg.compaction_keep_recent)
|
||||
|
||||
|
||||
def _cut_ids_of(msgs: list, new_msgs: list) -> tuple:
|
||||
"""🆕 压缩持久化切点：new_msgs = [摘要] + 保留尾巴（旧 msgs 的连续后缀）。
|
||||
按 identity 定位尾巴起点 → (cut_before_id, first_retained_id)（DB 行 id）。
|
||||
|
||||
⚠️ 行粒度：同一 DB 行（assistant timeline）会被 build_api_context 回放成
|
||||
多条 API 消息（asst/tool 条目共享同一 _db_msg_id）。若尾巴起点落在某行
|
||||
回放序列中间，必须回退到该行第一条，否则 cut_before 与 first_retained
|
||||
会是同一行 → DB 自环（链死循环卡死）。
|
||||
回退后下一轮会从该行完整重放（保守方向，宁多勿漏）。
|
||||
任一为空 = 无法在 DB 链上切（如尾巴全是轮中内存消息）→ UI 不插标记。"""
|
||||
if not new_msgs or getattr(new_msgs[0], "kind", None) != "compaction_summary":
|
||||
return "", ""
|
||||
tail = new_msgs[1:]
|
||||
if not tail:
|
||||
return "", ""
|
||||
first = tail[0]
|
||||
for i, m in enumerate(msgs):
|
||||
if m is first:
|
||||
# 回退到同一 DB 行回放序列的头部
|
||||
fid = getattr(first, "db_msg_id", "") or ""
|
||||
j = i
|
||||
while j > 0 and fid and getattr(msgs[j - 1], "db_msg_id", "") == fid:
|
||||
j -= 1
|
||||
if j == 0:
|
||||
return "", "" # 没有“之前”的行 → 无切点
|
||||
return (getattr(msgs[j - 1], "db_msg_id", "") or "",
|
||||
getattr(first, "db_msg_id", "") or "")
|
||||
return "", "" # 尾巴首条不在旧列表（不应发生）→ 不插标记
|
||||
|
||||
# ======================================================================
|
||||
# 对照 retry.ts 的模式表（原样移植）
|
||||
# ======================================================================
|
||||
_NON_RETRYABLE_LIMIT_RE = re.compile(
|
||||
r"GoUsageLimitError|FreeUsageLimitError|"
|
||||
r"Monthly usage limit reached|available balance|"
|
||||
r"insufficient_quota|out of budget|quota exceeded|billing",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
_RETRYABLE_RE = re.compile(
|
||||
r"overloaded|rate.?limit|too many requests|429|500|502|503|504|524|"
|
||||
r"service.?unavailable|server.?error|internal.?error|"
|
||||
r"provider.?returned.?error|exceeded request buffer limit while retrying upstream|"
|
||||
r"network.?error|connection.?error|connection.?refused|connection.?lost|"
|
||||
r"other side closed|fetch failed|getaddrinfo|ENOTFOUND|EAI_AGAIN|"
|
||||
r"upstream.?connect|reset before headers|socket hang up|timeout|timed?\s?out",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
# 对照 overflow.ts OVERFLOW_PATTERNS（各服务商话术）
|
||||
_OVERFLOW_RE = re.compile(
|
||||
r"context\s*window|prompt is too long|maximum context|too many tokens|"
|
||||
r"context length|exceeds the (model|maximum) context|"
|
||||
r"maximum context length is \d+ tokens|exceeds the maximum allowed input length|"
|
||||
r"longer than the model's context length|exceeds the available context size|"
|
||||
r"greater than the context length|exceeded model token limit|"
|
||||
r"Range of input length should be|configured context size|"
|
||||
r"长于|超出.*上下文|上下文.*(超出|超过|不足)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
# 限速类不是溢出（对照 overflow.ts NON_OVERFLOW_PATTERNS 的判定意图）
|
||||
_NON_OVERFLOW_RE = re.compile(
|
||||
r"rate.?limit|too many requests|429|overloaded", re.IGNORECASE)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 判定函数（对照 pi overflow.ts / retry.ts，基于 assistant 消息）
|
||||
# ======================================================================
|
||||
def is_context_overflow(message: AgentMessage,
|
||||
context_window: int = 0) -> bool:
|
||||
"""
|
||||
对照 isContextOverflow 的三种情形：
|
||||
1. error + 消息文本命中溢出模式（且不是限速类）
|
||||
2. 静默溢出：stop 但 usage.input > contextWindow（z.ai 风格）
|
||||
3. 服务端截断式溢出：length 且 output==0 且 input ≥ 0.99×窗口
|
||||
"""
|
||||
if message.role != "assistant":
|
||||
return False
|
||||
if message.stop_reason == "error" and message.error_message:
|
||||
if not _NON_OVERFLOW_RE.search(message.error_message) \
|
||||
and _OVERFLOW_RE.search(message.error_message):
|
||||
return True
|
||||
usage = message.usage or {}
|
||||
input_tokens = int(usage.get("input", 0) or 0) + int(usage.get("cacheRead", 0) or 0)
|
||||
if context_window and message.stop_reason == "stop":
|
||||
if input_tokens > context_window:
|
||||
return True
|
||||
if context_window and message.stop_reason == "length" \
|
||||
and int(usage.get("output", 0) or 0) == 0:
|
||||
if input_tokens >= context_window * 0.99:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_retryable_assistant_error(message: AgentMessage) -> bool:
|
||||
"""对照 isRetryableAssistantError：基于 errorMessage 文本模式"""
|
||||
if message.role != "assistant" or message.stop_reason != "error":
|
||||
return False
|
||||
if not message.error_message:
|
||||
return False
|
||||
if _NON_RETRYABLE_LIMIT_RE.search(message.error_message):
|
||||
return False
|
||||
return bool(_RETRYABLE_RE.search(message.error_message))
|
||||
|
||||
|
||||
def is_recoverable_length(message: AgentMessage,
|
||||
desired_max_output: int) -> bool:
|
||||
"""
|
||||
对照 pi isRecoverableLength：
|
||||
length 停止 且 实际输出 < 「原始输出上限」（钳制前的 maxTokens）。
|
||||
"""
|
||||
if message.role != "assistant" or message.stop_reason != "length":
|
||||
return False
|
||||
if desired_max_output <= 0:
|
||||
return False
|
||||
out = int((message.usage or {}).get("output", 0) or 0)
|
||||
return out < desired_max_output
|
||||
|
||||
|
||||
def compute_retry_delay_ms(attempt: int, cfg: RetryConfig) -> float:
|
||||
"""对照 pi _prepareRetry: baseDelayMs × 2^(attempt-1)（无抖动、默认无封顶）"""
|
||||
delay = cfg.base_delay_ms * (cfg.factor ** (attempt - 1))
|
||||
if cfg.max_delay_ms and delay > cfg.max_delay_ms:
|
||||
delay = cfg.max_delay_ms
|
||||
return float(delay)
|
||||
|
||||
|
||||
def find_last_assistant(messages) -> Optional[AgentMessage]:
|
||||
"""对照 agent-session _findLastAssistantMessage"""
|
||||
for m in reversed(messages):
|
||||
if m.role == "assistant":
|
||||
return m
|
||||
return None
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 编排器 —— 对照 agent-session 的 run 循环
|
||||
# ======================================================================
|
||||
class AgentRunner:
|
||||
"""
|
||||
对照 agent-session.ts 的 run 编排（agent 之外的「会话层」）。
|
||||
haocode 里由 AgentWorker(QThread) 持有一个 AgentRunner。
|
||||
|
||||
summarize_fn(transcript) -> str：压缩用的非流式 LLM 调用（缺省 → 禁压缩）。
|
||||
"""
|
||||
|
||||
def __init__(self, agent: Agent,
|
||||
summarize_fn: Optional[Callable[[str], str]] = None,
|
||||
on_retry_scheduled: Optional[Callable] = None,
|
||||
on_retry_finished: Optional[Callable] = None,
|
||||
on_compaction_started: Optional[Callable] = None,
|
||||
on_compaction_finished: Optional[Callable] = None):
|
||||
self.agent = agent
|
||||
self.summarize_fn = summarize_fn
|
||||
# 🆕 M3: 重试回调（对照 pi retryAssistantCall 的
|
||||
# onRetryScheduled / onRetryFinished）——UI 可提示「正在重试」
|
||||
self.on_retry_scheduled = on_retry_scheduled
|
||||
self.on_retry_finished = on_retry_finished
|
||||
# 🆕 压缩开始回调（摘要 LLM 调用阻塞前触发）——UI 显示「执行中」动态气泡
|
||||
self.on_compaction_started = on_compaction_started
|
||||
# 🆕 压缩完成回调（压缩结束瞬间触发，不等整轮 run 结束）
|
||||
# ——前端气泡即时定格「已完成上下文压缩」（与思考块定格同款时机）
|
||||
self.on_compaction_finished = on_compaction_finished
|
||||
self._overflow_recovered = False # 对照 _overflowRecoveryAttempted
|
||||
self._retry_attempt = 0 # 对照 _retryAttempt
|
||||
self.last_action = "none" # 诊断: none/retry/overflow_compact/length_compact/queued
|
||||
self.compactions_performed = 0 # 成功压缩次数（UI 提示用）
|
||||
# 🆕 压缩事件队列（UI 气泡用：摘要全文/前后 token/耗时/路径）
|
||||
self.compaction_events: list = []
|
||||
# 🆕 轮中主动压缩护栏（同一 run 内；AgentRunner 每轮新建 → 自动复位）
|
||||
self._mid_turn_fail_streak = 0
|
||||
self._mid_turn_exhausted = False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 发送前压缩（🌟 1:1 对照 pi shouldCompact：tokens > window - reserveTokens）
|
||||
# ------------------------------------------------------------------
|
||||
def pre_prompt_compaction(self) -> bool:
|
||||
if self.summarize_fn is None:
|
||||
return False
|
||||
cfg = self.agent.config
|
||||
msgs = self.agent.state.messages
|
||||
m = cfg.model
|
||||
# 🆕 G2: 透传 system_prompt/tools —— 无锚点时度量「下一请求真实载荷」
|
||||
should, tokens = should_compact(msgs, m, cfg.compaction_reserve,
|
||||
system_prompt=cfg.system_prompt,
|
||||
tools=cfg.tools)
|
||||
# 🆕 每次请求前的校验：print + 写 compaction_diag.log（触发与否都记）
|
||||
compact_diag_log(
|
||||
f"[PRE_CHECK] model={m.name} window={m.context_window} "
|
||||
f"reserve={cfg.compaction_reserve} "
|
||||
f"threshold={m.context_window - cfg.compaction_reserve} "
|
||||
f"est={tokens} triggered={str(should).lower()}")
|
||||
if not should:
|
||||
return False
|
||||
self._notify_compaction_started("pre_prompt", tokens)
|
||||
t0 = time.time()
|
||||
new_msgs = compact_context(msgs, cfg.model, self.summarize_fn,
|
||||
settings=_settings_of(cfg))
|
||||
if new_msgs is not None:
|
||||
cut_before, first_retained = _cut_ids_of(msgs, new_msgs)
|
||||
self.agent.state.messages = new_msgs
|
||||
self.compactions_performed += 1
|
||||
self._record_compaction("pre_prompt", tokens, new_msgs, t0,
|
||||
cut_before_id=cut_before,
|
||||
first_retained_id=first_retained)
|
||||
return True
|
||||
self._record_compaction_failed("pre_prompt", tokens, t0)
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 🆕 轮中主动压缩（haocode 增强，偏离 pi 1:1）
|
||||
# ------------------------------------------------------------------
|
||||
def compact_if_needed(self, msgs: list) -> Optional[list]:
|
||||
"""接线到 AgentConfig.compact_fn：内层循环每次 LLM 请求前调用。
|
||||
|
||||
与轮首 pre_prompt_compaction 同公式（tokens > window − reserve）。
|
||||
轮中单条工具输出可把上下文顶出窗口（轮中内存工具结果没有
|
||||
4000 字回放上限），发请求前主动压缩，避免「一次失败往返 +
|
||||
响应式兜底」。返回新消息列表（发生了压缩）或 None（不需要/
|
||||
不可压缩 → 原样发请求，响应式 overflow 安全网仍在）。
|
||||
|
||||
防刷屏护栏：同一 run 内连续失败 2 次后不再尝试（形状无有效
|
||||
切点，如单条巨型条目，硬试无意义）；成功则复位计数。
|
||||
"""
|
||||
if self.summarize_fn is None or self._mid_turn_exhausted:
|
||||
return None
|
||||
cfg = self.agent.config
|
||||
m = cfg.model
|
||||
# 🆕 G2: 透传 system_prompt/tools（与轮首同口径）
|
||||
should, tokens = should_compact(msgs, m, cfg.compaction_reserve,
|
||||
system_prompt=cfg.system_prompt,
|
||||
tools=cfg.tools)
|
||||
compact_diag_log(
|
||||
f"[PRE_CHECK] path=mid_turn est={tokens} "
|
||||
f"threshold={m.context_window - cfg.compaction_reserve} "
|
||||
f"triggered={str(should).lower()}")
|
||||
if not should:
|
||||
return None
|
||||
if self._do_compaction("mid_turn", msgs=msgs):
|
||||
self._mid_turn_fail_streak = 0
|
||||
return list(self.agent.state.messages)
|
||||
self._mid_turn_fail_streak += 1
|
||||
if self._mid_turn_fail_streak >= 2:
|
||||
self._mid_turn_exhausted = True
|
||||
compact_diag_log(
|
||||
"[PRE_CHECK] path=mid_turn exhausted "
|
||||
"（连续 2 次失败，本轮不再尝试）")
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 主入口 —— 对照 _runAgentPrompt
|
||||
# ------------------------------------------------------------------
|
||||
def run(self, message) -> RunResult:
|
||||
self.pre_prompt_compaction()
|
||||
self._overflow_recovered = False
|
||||
self._retry_attempt = 0
|
||||
result = self.agent.prompt(message)
|
||||
return self._post_loop(result)
|
||||
|
||||
def _post_loop(self, result: RunResult) -> RunResult:
|
||||
# 对照: while (await this._handlePostAgentRun()) await this.agent.continue()
|
||||
while True:
|
||||
action = self._handle_post_agent_run(result)
|
||||
if action == "stop":
|
||||
return result
|
||||
if action == "retry":
|
||||
result = self.agent.continue_()
|
||||
# 🆕 M3: 重试结果回调（成功 = 下一轮未以 error/aborted 收尾）
|
||||
if self.on_retry_finished is not None:
|
||||
la = find_last_assistant(self.agent.state.messages)
|
||||
ok = (la is not None
|
||||
and la.stop_reason not in ("error", "aborted"))
|
||||
try:
|
||||
self.on_retry_finished(bool(ok))
|
||||
except Exception:
|
||||
pass
|
||||
continue
|
||||
if action == "compact_retry":
|
||||
if not (self._remove_last_bad_assistant()
|
||||
and self._do_compaction(self.last_action)):
|
||||
return result # 压缩不可用 → 放弃，交 UI 报错
|
||||
result = self.agent.continue_()
|
||||
continue
|
||||
if action == "queued":
|
||||
result = self.agent.continue_()
|
||||
continue
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 三路决策 —— 对照 _handlePostAgentRun
|
||||
# ------------------------------------------------------------------
|
||||
def _handle_post_agent_run(self, result: RunResult) -> str:
|
||||
msg = find_last_assistant(self.agent.state.messages)
|
||||
if msg is None:
|
||||
return "stop"
|
||||
cfg = self.agent.config
|
||||
|
||||
# (1) 错误路径
|
||||
if msg.stop_reason == "error":
|
||||
# a. 溢出 → 压缩恢复（只试一次）
|
||||
if is_context_overflow(msg, cfg.model.context_window) \
|
||||
and not self._overflow_recovered \
|
||||
and self.summarize_fn is not None:
|
||||
self._overflow_recovered = True
|
||||
self.last_action = "overflow_compact"
|
||||
return "compact_retry"
|
||||
# b. 可重试错误 → 移除坏消息 + 退避 + continue
|
||||
if is_retryable_assistant_error(msg) \
|
||||
and self._prepare_retry(cfg.retry,
|
||||
reason=msg.error_message or ""):
|
||||
self.last_action = "retry"
|
||||
return "retry"
|
||||
return "stop"
|
||||
|
||||
# (2) 可恢复的 length 截断
|
||||
if is_recoverable_length(msg, cfg.model.max_tokens) \
|
||||
and not self._overflow_recovered \
|
||||
and self.summarize_fn is not None:
|
||||
self._overflow_recovered = True
|
||||
self.last_action = "length_compact"
|
||||
return "compact_retry"
|
||||
|
||||
# (3) 队列里还有消息 → 续跑
|
||||
if self.agent.has_queued():
|
||||
self.last_action = "queued"
|
||||
return "queued"
|
||||
|
||||
# (4) settle
|
||||
return "stop"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 对照 _prepareRetry：计数 → 移除坏消息 → 退避睡眠
|
||||
# ------------------------------------------------------------------
|
||||
def _prepare_retry(self, cfg: RetryConfig, reason: str = "") -> bool:
|
||||
if self._retry_attempt >= cfg.max_attempts:
|
||||
return False
|
||||
self._retry_attempt += 1
|
||||
delay_ms = compute_retry_delay_ms(self._retry_attempt, cfg)
|
||||
# 移除错误助手消息（对照: messages.slice(0,-1)，会话历史保留由 UI 层负责）
|
||||
messages = self.agent.state.messages
|
||||
if messages and messages[-1].role == "assistant":
|
||||
self.agent.state.messages = messages[:-1]
|
||||
# 🆕 M3: 通知重试已调度（参数: attempt, max_attempts, delay_ms, reason）
|
||||
if self.on_retry_scheduled is not None:
|
||||
try:
|
||||
self.on_retry_scheduled(self._retry_attempt, cfg.max_attempts,
|
||||
delay_ms, reason)
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(delay_ms / 1000.0)
|
||||
return True
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 压缩恢复辅助
|
||||
# ------------------------------------------------------------------
|
||||
def _remove_last_bad_assistant(self) -> bool:
|
||||
msgs = self.agent.state.messages
|
||||
if msgs and msgs[-1].role == "assistant" \
|
||||
and msgs[-1].stop_reason in ("length", "error"):
|
||||
self.agent.state.messages = msgs[:-1]
|
||||
return True
|
||||
return False
|
||||
|
||||
def _do_compaction(self, path: str = "reactive",
|
||||
msgs: Optional[list] = None) -> bool:
|
||||
if self.summarize_fn is None:
|
||||
return False
|
||||
cfg = self.agent.config
|
||||
if msgs is None:
|
||||
msgs = self.agent.state.messages
|
||||
tokens_before = estimate_context_tokens(msgs).tokens
|
||||
self._notify_compaction_started(path, tokens_before)
|
||||
t0 = time.time()
|
||||
new_msgs = compact_context(msgs, cfg.model, self.summarize_fn,
|
||||
settings=_settings_of(cfg))
|
||||
if new_msgs is not None:
|
||||
cut_before, first_retained = _cut_ids_of(msgs, new_msgs)
|
||||
self.agent.state.messages = new_msgs
|
||||
self.compactions_performed += 1
|
||||
self._record_compaction(path, tokens_before, new_msgs, t0,
|
||||
cut_before_id=cut_before,
|
||||
first_retained_id=first_retained)
|
||||
return True
|
||||
self._record_compaction_failed(path, tokens_before, t0)
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 🆕 压缩可视化辅助：开始/完成 两阶段（print + 日志 + 事件供 UI 气泡）
|
||||
# ------------------------------------------------------------------
|
||||
def _notify_compaction_started(self, path: str, tokens_before: int) -> None:
|
||||
compact_diag_log(f"[COMPACT_START] path={path} before={tokens_before}")
|
||||
if self.on_compaction_started is not None:
|
||||
try:
|
||||
self.on_compaction_started(path)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _record_compaction_failed(self, path: str, tokens_before: int,
|
||||
t0: float) -> None:
|
||||
"""压缩不可行（如历史无法切分）→ 气泡显示失败原因，不能永远卡在「执行中」"""
|
||||
duration_ms = int((time.time() - t0) * 1000)
|
||||
compact_diag_log(
|
||||
f"[COMPACT_FAIL] path={path} before={tokens_before} "
|
||||
f"duration_ms={duration_ms}")
|
||||
ev = {
|
||||
"summary": "本次压缩未能执行（历史无法切分），继续以原有历史对话。",
|
||||
"before": tokens_before,
|
||||
"after": tokens_before,
|
||||
"duration_ms": duration_ms,
|
||||
"path": path,
|
||||
"failed": True,
|
||||
}
|
||||
self.compaction_events.append(ev)
|
||||
if self.on_compaction_finished is not None:
|
||||
try:
|
||||
self.on_compaction_finished(dict(ev))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _record_compaction(self, path: str, tokens_before: int,
|
||||
new_msgs: list, t0: float,
|
||||
cut_before_id: str = "",
|
||||
first_retained_id: str = "") -> None:
|
||||
summary = ""
|
||||
if new_msgs and getattr(new_msgs[0], "kind", None) == "compaction_summary":
|
||||
summary = _content_text(new_msgs[0].content) or ""
|
||||
tokens_after = estimate_context_tokens(new_msgs).tokens
|
||||
duration_ms = int((time.time() - t0) * 1000)
|
||||
compact_diag_log(
|
||||
f"[COMPACT_DONE] path={path} before={tokens_before} "
|
||||
f"after={tokens_after} summary_chars={len(summary)} "
|
||||
f"duration_ms={duration_ms} "
|
||||
f"cut_before={cut_before_id or '-'} first_retained={first_retained_id or '-'}")
|
||||
ev = {
|
||||
"summary": summary,
|
||||
"before": tokens_before,
|
||||
"after": tokens_after,
|
||||
"duration_ms": duration_ms,
|
||||
"path": path,
|
||||
# 🆕 压缩持久化切点：UI 收到后在 DB 链上插标记行（两者齐备才插）
|
||||
"cut_before_id": cut_before_id,
|
||||
"first_retained_id": first_retained_id,
|
||||
}
|
||||
self.compaction_events.append(ev)
|
||||
if self.on_compaction_finished is not None:
|
||||
try:
|
||||
self.on_compaction_finished(dict(ev))
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,438 @@
|
||||
"""
|
||||
core/agent/stream_fn.py
|
||||
=======================
|
||||
🌟 pi-ai 流式接口的 Python 1:1 移植 —— OpenAI 兼容（含 vLLM）
|
||||
|
||||
对照 pi-main 源码：
|
||||
packages/ai/src/api/openai-completions.ts
|
||||
- streamSimple() (行 469 起): 消费 SSE chunk → 产出 AssistantMessage
|
||||
- buildParams() (行 536 起): 组装请求体
|
||||
* vLLM/OpenAI 默认字段名 = max_completion_tokens
|
||||
* 少数服务商(chutes/deepseek/moonshot 等)用 max_tokens
|
||||
- mapStopReason(): finish_reason "length" → stopReason "length"
|
||||
packages/ai/src/api/simple-options.ts
|
||||
- buildBaseOptions(): maxTokens = options?.maxTokens ?? model.maxTokens
|
||||
（即模型定义里的 maxTokens 一定会进入请求体，这就是它被发给 vLLM 的原因）
|
||||
|
||||
本文件职责：
|
||||
1. to_openai_messages() —— pi 消息格式 → OpenAI API 格式（每轮调用前转换）
|
||||
2. openai_stream() —— 流式请求，逐 chunk 产出 AssistantMessageEvent，
|
||||
结束产出最终 AgentMessage（含 usage / stop_reason / tool_calls）
|
||||
3. classify_error() —— 异常分类（对照 pi isRetryable 的输入）
|
||||
|
||||
🌟 修复 haocode 原有 P1 bug：openai 客户端带 timeout（原来无超时，
|
||||
挂死的流会让线程永久阻塞）。max_retries=0（重试统一交给 recovery 层，
|
||||
与 pi 一致：网络层不重试，会话层按 1600ms×1.6^n 退避重试）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, Iterator, List, Optional, Tuple
|
||||
|
||||
from openai import OpenAI
|
||||
from openai import APIConnectionError, APIStatusError, APITimeoutError, RateLimitError
|
||||
|
||||
from .types import (AgentError, AgentMessage, AssistantMessageEvent,
|
||||
AbortSignal, ModelConfig, ToolCall, new_id)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# pi 消息 → OpenAI 消息（每轮发送前调用）
|
||||
# ======================================================================
|
||||
def to_openai_messages(context: List[AgentMessage],
|
||||
pass_reasoning: bool = True) -> List[Dict[str, Any]]:
|
||||
"""pi 消息 → OpenAI chat 格式。
|
||||
🆕 thinking 模式（DeepSeek 等）：assistant 消息必须把 reasoning_content
|
||||
原样回传，否则服务端 400
|
||||
（'The reasoning_content in the thinking mode must be passed back to the API'）。
|
||||
默认开启：轮中续问/工具回灌场景，内存里本轮刚流式产出的 assistant 带
|
||||
reasoning → 回传；DB 回放的旧行 reasoning=""（不入库）→ 无该字段，
|
||||
跨轮全新请求不受影响。pass_reasoning=False 显式关闭（特殊网关逃生阀）。
|
||||
"""
|
||||
out: List[Dict[str, Any]] = []
|
||||
for m in context:
|
||||
if m.role == "system":
|
||||
# 只放行请求头部的系统提示词（loop 注入）；
|
||||
# 历史中间出现的 system（DB 防御）仍然跳过
|
||||
if not out:
|
||||
out.append({"role": "system", "content": m.content})
|
||||
continue
|
||||
if m.role == "toolResult":
|
||||
out.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": m.tool_call_id,
|
||||
"content": m.content if isinstance(m.content, str)
|
||||
else "".join(c.get("text", "") for c in m.content
|
||||
if isinstance(c, dict)),
|
||||
})
|
||||
continue
|
||||
if m.role == "assistant":
|
||||
entry: Dict[str, Any] = {"role": "assistant"}
|
||||
# 纯工具调用轮：content 可能为空 → 用 ""（vLLM 接受，避免 null 报错）
|
||||
entry["content"] = m.content if m.content else ""
|
||||
if m.tool_calls:
|
||||
entry["tool_calls"] = [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.name,
|
||||
"arguments": json.dumps(tc.arguments, ensure_ascii=False)
|
||||
if tc.arguments else tc.raw_arguments or "{}",
|
||||
},
|
||||
}
|
||||
for tc in m.tool_calls
|
||||
]
|
||||
if pass_reasoning and m.reasoning:
|
||||
entry["reasoning_content"] = m.reasoning
|
||||
out.append(entry)
|
||||
continue
|
||||
# user（content 可以是 str 或 OpenAI 多模态 list，原样透传）
|
||||
out.append({"role": "user", "content": m.content})
|
||||
return out
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 异常分类 —— 对照 pi isRetryable / retryable error 判定
|
||||
# ======================================================================
|
||||
def classify_error(e: Exception) -> AgentError:
|
||||
"""把 openai SDK 异常分类，恢复层据此决定重试/放弃。"""
|
||||
if isinstance(e, RateLimitError):
|
||||
return AgentError(message=str(e), kind="rate_limit",
|
||||
status_code=429, recoverable=True)
|
||||
if isinstance(e, APITimeoutError):
|
||||
return AgentError(message=f"请求超时: {e}", kind="timeout", recoverable=True)
|
||||
if isinstance(e, APIConnectionError):
|
||||
return AgentError(message=f"连接失败: {e}", kind="connection", recoverable=True)
|
||||
if isinstance(e, APIStatusError):
|
||||
code = e.status_code
|
||||
msg = str(e)
|
||||
if code in (502, 503, 504):
|
||||
return AgentError(message=msg, kind="server_error",
|
||||
status_code=code, recoverable=True)
|
||||
if code == 529:
|
||||
return AgentError(message=msg, kind="overload",
|
||||
status_code=code, recoverable=True)
|
||||
if code == 408:
|
||||
return AgentError(message=msg, kind="timeout",
|
||||
status_code=code, recoverable=True)
|
||||
if code in (401, 403):
|
||||
return AgentError(message=msg, kind="auth",
|
||||
status_code=code, recoverable=False)
|
||||
# 4xx（参数错误/模型不存在/上下文超长等）不可重试，交恢复层细分
|
||||
return AgentError(message=msg, kind="unknown",
|
||||
status_code=code, recoverable=False)
|
||||
if isinstance(e, (ConnectionError, TimeoutError)):
|
||||
return AgentError(message=str(e), kind="connection", recoverable=True)
|
||||
return AgentError(message=str(e), kind="unknown", recoverable=False)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 🆕 M1/M2: 流式字段提取（对照 pi openai-completions.ts 的字段优先级）
|
||||
# ======================================================================
|
||||
# 🆕 M1: 1:1 对照 pi chunk 级 reasoning 字段优先级（llama.cpp/vLLM/OpenRouter 各不同）
|
||||
_REASONING_FIELDS = ("reasoning_content", "reasoning", "reasoning_text")
|
||||
|
||||
|
||||
def _pick_reasoning(delta) -> str:
|
||||
"""从 delta 上取第一个非空思考字段（只取一个，防同内容多字段重复计）"""
|
||||
if delta is None:
|
||||
return ""
|
||||
dump = None
|
||||
for name in _REASONING_FIELDS:
|
||||
v = getattr(delta, name, None)
|
||||
if v is None and dump is None and hasattr(delta, "model_dump"):
|
||||
try:
|
||||
dump = delta.model_dump()
|
||||
except Exception:
|
||||
dump = {}
|
||||
if v is None and isinstance(dump, dict):
|
||||
v = dump.get(name)
|
||||
if v:
|
||||
return str(v)
|
||||
return ""
|
||||
|
||||
|
||||
def _pick_usage(chunk):
|
||||
"""🆕 M2: 对照 pi：先 chunk.usage，再 choices[0].usage（Moonshot 系只放 choice 里）"""
|
||||
u = getattr(chunk, "usage", None)
|
||||
if u is not None:
|
||||
return u
|
||||
try:
|
||||
choices = getattr(chunk, "choices", None)
|
||||
if choices:
|
||||
return getattr(choices[0], "usage", None)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 流式请求核心 —— 对照 openai-completions.ts streamSimple
|
||||
# ======================================================================
|
||||
def openai_stream(context: List[AgentMessage],
|
||||
model: ModelConfig,
|
||||
signal: AbortSignal,
|
||||
max_tokens: int,
|
||||
tools: Optional[List[Any]] = None,
|
||||
) -> Iterator[Tuple[str, Any]]:
|
||||
"""
|
||||
生成器协议（对照 pi streamSimple 的 for-await 消费方式）：
|
||||
yield ("event", AssistantMessageEvent) # 增量事件（text/thinking/toolcall/done）
|
||||
yield ("final", AgentMessage) # 最终助手消息（含 stop_reason/usage）
|
||||
异常统一转 AgentError 抛出（由循环层捕获 → stop_reason="error"）。
|
||||
|
||||
🌟 中止语义与 pi 一致：每个 chunk 边界检查 signal.aborted，
|
||||
命中即关闭流、以 stop_reason="aborted" 收尾（不是 error）。
|
||||
"""
|
||||
client = OpenAI(
|
||||
api_key=model.api_key or "EMPTY",
|
||||
base_url=model.base_url,
|
||||
# 🌟 P1 修复：显式超时（pi 侧由 fetch timeout 保证）
|
||||
timeout=model.timeout_seconds,
|
||||
max_retries=0, # 重试统一由 recovery 层负责
|
||||
)
|
||||
|
||||
params: Dict[str, Any] = {
|
||||
"model": model.name,
|
||||
"messages": to_openai_messages(context, model.pass_reasoning),
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True}, # 拿真实 usage（pi 同款）
|
||||
"temperature": model.temperature,
|
||||
}
|
||||
# 对照 pi buildParams：默认 max_completion_tokens（vLLM 兼容）
|
||||
if getattr(model, "use_max_tokens_field", False):
|
||||
params["max_tokens"] = max_tokens
|
||||
else:
|
||||
params["max_completion_tokens"] = max_tokens
|
||||
|
||||
# 🌟 工具定义（对照 pi buildParams 的 tools 段）——不发给模型，
|
||||
# 模型就不可能发出真正的 tool_call，只会用文字"演"工具调用！
|
||||
if tools:
|
||||
params["tools"] = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": t.name,
|
||||
"description": t.description,
|
||||
"parameters": t.parameters,
|
||||
},
|
||||
}
|
||||
for t in tools
|
||||
]
|
||||
|
||||
# 累积状态
|
||||
content_parts: List[str] = []
|
||||
reasoning_parts: List[str] = []
|
||||
tc_acc: Dict[int, Dict[str, Any]] = {} # index -> {id, name, args}
|
||||
finish_reason: Optional[str] = None
|
||||
usage: Dict[str, Any] = {}
|
||||
stream = None
|
||||
|
||||
try:
|
||||
stream = client.chat.completions.create(**params)
|
||||
for chunk in stream:
|
||||
# 🌟 pi 同款：chunk 边界检查中止
|
||||
if signal.aborted:
|
||||
finish_reason = "aborted"
|
||||
try:
|
||||
stream.close()
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
|
||||
# 🆕 M2: usage 读取顺序对照 pi —— chunk.usage 优先，choice.usage 兜底
|
||||
usage_obj = _pick_usage(chunk)
|
||||
if usage_obj is not None:
|
||||
usage = {
|
||||
"input": getattr(usage_obj, "prompt_tokens", 0) or 0,
|
||||
"output": getattr(usage_obj, "completion_tokens", 0) or 0,
|
||||
"cacheRead": 0, "cacheWrite": 0,
|
||||
}
|
||||
if not getattr(chunk, "choices", None):
|
||||
continue
|
||||
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
|
||||
# 🆕 M1: 思考增量（对照 pi reasoningFields 优先级，含多字段去重）
|
||||
reasoning = _pick_reasoning(delta)
|
||||
if reasoning:
|
||||
yield ("event", AssistantMessageEvent.thinking_delta(reasoning))
|
||||
reasoning_parts.append(reasoning)
|
||||
|
||||
if getattr(delta, "content", None):
|
||||
yield ("event", AssistantMessageEvent.text_delta(delta.content))
|
||||
content_parts.append(delta.content)
|
||||
|
||||
# 工具调用增量（按 index 累积 id/name/arguments 三段）
|
||||
for tc_delta in (getattr(delta, "tool_calls", None) or []):
|
||||
idx = tc_delta.index or 0
|
||||
slot = tc_acc.setdefault(idx, {"id": "", "name": "", "args": ""})
|
||||
if tc_delta.id:
|
||||
slot["id"] = tc_delta.id
|
||||
yield ("event", AssistantMessageEvent.toolcall_delta(
|
||||
idx, tc_delta.id, field="id"))
|
||||
fn = tc_delta.function
|
||||
if fn is not None:
|
||||
if fn.name:
|
||||
slot["name"] += fn.name
|
||||
yield ("event", AssistantMessageEvent.toolcall_delta(
|
||||
idx, fn.name, field="name"))
|
||||
if fn.arguments:
|
||||
slot["args"] += fn.arguments
|
||||
yield ("event", AssistantMessageEvent.toolcall_delta(
|
||||
idx, fn.arguments, field="arguments"))
|
||||
|
||||
if choice.finish_reason:
|
||||
finish_reason = choice.finish_reason
|
||||
except Exception as e:
|
||||
# 中止过程中断网/断流不当错误处理（pi 同款语义）
|
||||
if signal.aborted:
|
||||
finish_reason = "aborted"
|
||||
else:
|
||||
raise classify_error(e)
|
||||
finally:
|
||||
try:
|
||||
if stream is not None:
|
||||
stream.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---- 组装最终助手消息（对照 mapStopReason + AssistantMessage 构造）----
|
||||
# pi mapStopReason: "length"→"length", "stop"/"tool_calls"→"stop", 其余→"stop"
|
||||
if finish_reason == "aborted":
|
||||
stop_reason = "aborted"
|
||||
elif finish_reason == "length":
|
||||
stop_reason = "length"
|
||||
else:
|
||||
stop_reason = "stop"
|
||||
|
||||
tool_calls: List[ToolCall] = []
|
||||
for idx in sorted(tc_acc.keys()):
|
||||
slot = tc_acc[idx]
|
||||
try:
|
||||
args = json.loads(slot["args"]) if slot["args"] else {}
|
||||
except json.JSONDecodeError:
|
||||
# 参数 JSON 被截断/损坏：保留 raw，arguments 置空，
|
||||
# 由循环层按「截断保护」路径处理（不执行残缺调用）
|
||||
args = {}
|
||||
tool_calls.append(ToolCall(
|
||||
id=slot["id"] or new_id("call"),
|
||||
name=slot["name"],
|
||||
arguments=args,
|
||||
raw_arguments=slot["args"],
|
||||
))
|
||||
|
||||
final = AgentMessage(
|
||||
role="assistant",
|
||||
content="".join(content_parts),
|
||||
reasoning="".join(reasoning_parts),
|
||||
tool_calls=tool_calls,
|
||||
stop_reason=stop_reason,
|
||||
usage=usage,
|
||||
)
|
||||
yield ("event", AssistantMessageEvent.done())
|
||||
yield ("final", final)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# OpenAI 格式 → pi 格式（from_openai_messages 的逆转换，用于从 DB 装载历史）
|
||||
# ======================================================================
|
||||
def _parse_tool_call(tc: Dict[str, Any]) -> Optional[ToolCall]:
|
||||
"""
|
||||
解析一条 OpenAI tool_call → ToolCall。
|
||||
🌟 参数 JSON 可能残缺（历史里存了被截断的调用）：
|
||||
解析失败 → arguments={} 且 raw_arguments 兜底成合法 "{}"，
|
||||
保证 to_openai_messages 回发 API 时参数始终是合法 JSON
|
||||
（API 只校验 JSON 合法性 + tool_call_id 配对，不校验参数内容）。
|
||||
"""
|
||||
if not isinstance(tc, dict):
|
||||
return None
|
||||
fn = tc.get("function") or {}
|
||||
raw_args = fn.get("arguments")
|
||||
if not isinstance(raw_args, str):
|
||||
raw_args = "" if raw_args is None else str(raw_args)
|
||||
try:
|
||||
args = json.loads(raw_args) if raw_args else {}
|
||||
if not isinstance(args, dict):
|
||||
args = {"_": args}
|
||||
except (json.JSONDecodeError, TypeError, ValueError):
|
||||
args = {}
|
||||
raw_args = "{}" # 残缺 JSON 兜底，防回发 API 时 400
|
||||
return ToolCall(
|
||||
id=tc.get("id") or new_id("call"),
|
||||
name=fn.get("name") or "",
|
||||
arguments=args,
|
||||
raw_arguments=raw_args or "{}",
|
||||
)
|
||||
|
||||
|
||||
def from_openai_messages(messages: List[Dict[str, Any]]) -> List[AgentMessage]:
|
||||
"""
|
||||
把 OpenAI chat 消息数组转成 pi 的 AgentMessage 列表（全保真）。
|
||||
|
||||
🌟 修复：worker 模式此前在此丢失全部工具历史（tool 消息被跳过、
|
||||
assistant.tool_calls 不解析）→ 任务完成/中断后再次提问，模型看不到
|
||||
之前读过什么文件、执行过什么命令。现改为无损转换：
|
||||
- user: content 保持 str 或 OpenAI 多模态 list（原样透传）
|
||||
- assistant: content + reasoning + tool_calls（纯工具轮保留 tool_calls，
|
||||
content 置 ""，不再变成空 assistant 污染上下文）
|
||||
- tool: → AgentMessage(role="toolResult", tool_call_id, content)
|
||||
- system: 防御性跳过（DB 链表里不会出现）
|
||||
|
||||
下游 to_openai_messages 原生支持 toolResult / assistant.tool_calls，
|
||||
与 pi 语义一致，往返（to(from(x))）保真。
|
||||
"""
|
||||
out: List[AgentMessage] = []
|
||||
for m in messages:
|
||||
if not isinstance(m, dict):
|
||||
continue
|
||||
role = m.get("role")
|
||||
content = m.get("content")
|
||||
|
||||
if role == "user":
|
||||
if content is None:
|
||||
content = ""
|
||||
out.append(AgentMessage(role="user", content=content,
|
||||
# 🆕 压缩持久化：跨轮携带 DB 行 id + 摘要 kind
|
||||
# （kind=compaction_summary → 二次压缩走迭代更新）
|
||||
kind=m.get("_kind") or "",
|
||||
db_msg_id=m.get("_db_msg_id") or "",
|
||||
timestamp=int(m.get("timestamp") or 0)))
|
||||
|
||||
elif role == "assistant":
|
||||
# tool_calls 解析（容忍缺失/损坏）
|
||||
tool_calls: List[ToolCall] = []
|
||||
for tc in (m.get("tool_calls") or []):
|
||||
parsed = _parse_tool_call(tc)
|
||||
if parsed is not None:
|
||||
tool_calls.append(parsed)
|
||||
out.append(AgentMessage(
|
||||
role="assistant",
|
||||
content=content if isinstance(content, str) else "",
|
||||
reasoning=m.get("reasoning") or "",
|
||||
tool_calls=tool_calls,
|
||||
stop_reason="stop",
|
||||
# 🆕 P1: 回放入库的 usage（锚点）+ timestamp（P0 时效校验）
|
||||
usage=m.get("usage") or {},
|
||||
timestamp=int(m.get("timestamp") or 0),
|
||||
db_msg_id=m.get("_db_msg_id") or "",
|
||||
))
|
||||
|
||||
elif role == "tool":
|
||||
# tool 结果 → toolResult（保留 tool_call_id 供 API 配对）
|
||||
if content is None:
|
||||
content = ""
|
||||
elif not isinstance(content, str):
|
||||
content = str(content)
|
||||
out.append(AgentMessage(
|
||||
role="toolResult",
|
||||
content=content,
|
||||
tool_call_id=m.get("tool_call_id") or "",
|
||||
timestamp=int(m.get("timestamp") or 0),
|
||||
db_msg_id=m.get("_db_msg_id") or "",
|
||||
))
|
||||
|
||||
# 其他 role（system 等）防御性跳过
|
||||
return out
|
||||
@@ -0,0 +1,873 @@
|
||||
"""
|
||||
core/agent/tools.py
|
||||
===================
|
||||
🌟 pi 工具执行管线的 Python 1:1 移植 + 内置工具
|
||||
|
||||
对照 pi-main 源码：
|
||||
packages/agent/src/agent-loop.ts
|
||||
- prepareToolCalls() (行 ~470): 校验参数 → 标记错误（不执行）
|
||||
- executeTool() (行 ~520): before 钩子 → 执行 → after 钩子 → 结果定型
|
||||
- 并行语义: prepare 串行 → 执行并发(Promise.all) → 结果按原始顺序回写
|
||||
packages/agent/src/tools/*.ts (coding-agent 内置工具 read/bash/edit/write)
|
||||
|
||||
JSON Schema 校验：pi 用 validate-json-schema + ai/src/utils/validation.ts
|
||||
（structuredClone → normalizeOptionalNulls → Value.Convert → coerceWithJsonSchema →
|
||||
全量错误上报 + 回显收到的参数）；这里实现核心子集
|
||||
（type/required/properties/enum/items + 可选字段 null 归一化 + 数字/布尔轻量转换），
|
||||
零外部依赖。工具自定义预处理对照 pi 的 tool.prepareArguments（见 edit 的 legacy 兼容）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from .types import (AgentMessage, AgentTool, AgentToolResult, AbortSignal,
|
||||
ToolCall, new_id)
|
||||
|
||||
try:
|
||||
from core.debug_log import debug_log as _dbg_log # 🆕 计时观察日志（线程安全/静默）
|
||||
except Exception: # 导入失败也不影响工具执行
|
||||
def _dbg_log(msg, tag="APP"):
|
||||
pass
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# JSON Schema 校验（子集）
|
||||
# ======================================================================
|
||||
# ======================================================================
|
||||
# 参数归一化 + JSON Schema 校验（子集）
|
||||
# 对照 pi ai/src/utils/validation.ts:317-347 validateToolArguments：
|
||||
# structuredClone → normalizeOptionalNulls → Value.Convert
|
||||
# → coerceWithJsonSchema → Check → 【报全部错误 + 回显收到的参数】
|
||||
# 这里实现核心子集（type/required/properties/enum/items），零外部依赖。
|
||||
# ======================================================================
|
||||
_TYPE_MAP = {
|
||||
"string": str,
|
||||
"integer": int,
|
||||
"number": (int, float),
|
||||
"boolean": bool,
|
||||
"array": list,
|
||||
"object": dict,
|
||||
}
|
||||
|
||||
|
||||
def _norm_types(schema: Dict[str, Any]) -> List[str]:
|
||||
"""schema 声明的类型（兼容 TypeBox 的 type 数组形式）"""
|
||||
t = schema.get("type")
|
||||
if isinstance(t, list):
|
||||
return [x for x in t if isinstance(x, str)]
|
||||
return [t] if isinstance(t, str) else []
|
||||
|
||||
|
||||
def _type_ok(value: Any, t: str) -> bool:
|
||||
"""类型匹配。🌟 bool 不算 integer/number（python 里 bool 是 int 子类）"""
|
||||
if t == "integer":
|
||||
return isinstance(value, int) and not isinstance(value, bool)
|
||||
if t == "number":
|
||||
return isinstance(value, (int, float)) and not isinstance(value, bool)
|
||||
if t == "boolean":
|
||||
return isinstance(value, bool)
|
||||
py = _TYPE_MAP.get(t)
|
||||
return isinstance(value, py) if py else True
|
||||
|
||||
|
||||
def _join_path(path: str, key: str) -> str:
|
||||
return f"{path}.{key}" if path else key
|
||||
|
||||
|
||||
def normalize_and_coerce(args: Any, schema: Dict[str, Any]) -> Any:
|
||||
"""归一化 + 轻量强制转换（对照 pi normalizeOptionalNulls + coerceWithJsonSchema）
|
||||
|
||||
- 可选字段的显式 null → 删除该键（模型常对「未用到的参数」发 null）
|
||||
- 数字字段收到纯数字字符串 → integer/number（"30" → 30）
|
||||
- boolean 字段收到 "true"/"false" → 布尔
|
||||
返回新对象，不修改入参。
|
||||
"""
|
||||
if not isinstance(schema, dict) or not isinstance(args, dict):
|
||||
return args
|
||||
props = schema.get("properties") or {}
|
||||
required = set(schema.get("required") or [])
|
||||
out: Dict[str, Any] = dict(args)
|
||||
for key, val in list(out.items()):
|
||||
sub = props.get(key)
|
||||
if not isinstance(sub, dict):
|
||||
continue
|
||||
if val is None:
|
||||
if key not in required:
|
||||
del out[key]
|
||||
continue
|
||||
out[key] = _coerce_value(val, sub)
|
||||
return out
|
||||
|
||||
|
||||
def _coerce_value(val: Any, sub: Dict[str, Any]) -> Any:
|
||||
types = _norm_types(sub)
|
||||
if isinstance(val, bool) or val is None:
|
||||
return val
|
||||
if isinstance(val, str):
|
||||
s = val.strip()
|
||||
if s and ("integer" in types or "number" in types):
|
||||
try:
|
||||
f = float(s)
|
||||
if f.is_integer():
|
||||
return int(f)
|
||||
if "number" in types:
|
||||
return f
|
||||
return val # integer 字段收到 1.5 → 保持原值（交由校验报错）
|
||||
except ValueError:
|
||||
return val
|
||||
if "boolean" in types and s.lower() in ("true", "false"):
|
||||
return s.lower() == "true"
|
||||
return val
|
||||
if isinstance(val, dict):
|
||||
return normalize_and_coerce(val, sub)
|
||||
if isinstance(val, list):
|
||||
items = sub.get("items")
|
||||
if isinstance(items, dict):
|
||||
return [_coerce_value(v, items) if not isinstance(v, dict)
|
||||
else normalize_and_coerce(v, items) for v in val]
|
||||
return val
|
||||
|
||||
|
||||
def _type_error(path: str, types: List[str], value: Any) -> str:
|
||||
t = types[0] if types else "object"
|
||||
if t == "string":
|
||||
return f"参数 {path} 应为 string"
|
||||
if t in ("integer", "number") and isinstance(value, bool):
|
||||
return f"参数 {path} 应为 {t}"
|
||||
return f"参数 {path} 类型错误: 期望 {t}"
|
||||
|
||||
|
||||
def _check_value(value: Any, schema: Dict[str, Any], path: str,
|
||||
errs: List[str], root: bool = False) -> None:
|
||||
"""递归收集【全部】校验错误（对照 pi Errors() 全量上报）"""
|
||||
if not isinstance(schema, dict):
|
||||
return
|
||||
if root and not isinstance(value, dict):
|
||||
errs.append(f"参数必须是对象，实际是 {type(value).__name__}")
|
||||
return
|
||||
types = _norm_types(schema) or (["object"] if root else [])
|
||||
if types and not any(_type_ok(value, t) for t in types):
|
||||
errs.append(_type_error(path, types, value))
|
||||
return # 类型不符 → 后续检查无意义
|
||||
if "enum" in schema and value not in schema["enum"]:
|
||||
errs.append(f"参数 {path} 取值必须是 {schema['enum']} 之一")
|
||||
if isinstance(value, dict):
|
||||
props = schema.get("properties") or {}
|
||||
for req in schema.get("required") or []:
|
||||
if req not in value:
|
||||
errs.append(f"缺少必填参数: {req}" if root
|
||||
else f"参数 {path} 缺少必填字段 {req}")
|
||||
for k, v in value.items():
|
||||
if k in props:
|
||||
_check_value(v, props[k], _join_path(path, k), errs)
|
||||
elif isinstance(value, list):
|
||||
items = schema.get("items")
|
||||
if isinstance(items, dict):
|
||||
for i, item in enumerate(value):
|
||||
_check_value(item, items, f"{path}[{i}]", errs)
|
||||
|
||||
|
||||
def validate_json_schema(args: Any, schema: Dict[str, Any]) -> Optional[str]:
|
||||
"""返回错误描述（多条以「; 」连接）；合法返回 None。"""
|
||||
errs: List[str] = []
|
||||
_check_value(args, schema, "", errs, root=True)
|
||||
return "; ".join(errs) if errs else None
|
||||
|
||||
|
||||
def _validate_value(value: Any, schema: Dict[str, Any], path: str) -> Optional[str]:
|
||||
"""（保留旧签名：返回该节点的首个错误）"""
|
||||
errs: List[str] = []
|
||||
_check_value(value, schema, path, errs)
|
||||
return errs[0] if errs else None
|
||||
|
||||
|
||||
def _brief_json(obj: Any, limit: int = 600) -> str:
|
||||
"""参数回显（长内容截断，避免 write 的大 content 撑爆错误消息）"""
|
||||
try:
|
||||
s = json.dumps(obj, ensure_ascii=False)
|
||||
except Exception:
|
||||
s = repr(obj)
|
||||
return s if len(s) <= limit else s[:limit] + f"…(共 {len(s)} 字符)"
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 准备阶段 —— 对照 agent-loop.ts prepareToolCalls
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class PreparedToolCall:
|
||||
"""一次工具调用的完整准备结果（执行前定型）"""
|
||||
tool_call: ToolCall
|
||||
tool: Optional[AgentTool]
|
||||
args: Dict[str, Any]
|
||||
error: str = "" # 准备阶段失败原因（未知工具/参数非法）→ 直接返回错误 toolResult
|
||||
|
||||
|
||||
def prepare_tool_calls(assistant: AgentMessage,
|
||||
tools: List[AgentTool]) -> List[PreparedToolCall]:
|
||||
tool_map = {t.name: t for t in tools}
|
||||
prepared: List[PreparedToolCall] = []
|
||||
for tc in assistant.tool_calls:
|
||||
tool = tool_map.get(tc.name)
|
||||
if tool is None:
|
||||
prepared.append(PreparedToolCall(
|
||||
tool_call=tc, tool=None, args={},
|
||||
error=f"未知工具: {tc.name}",
|
||||
))
|
||||
continue
|
||||
if not isinstance(tc.arguments, dict):
|
||||
prepared.append(PreparedToolCall(
|
||||
tool_call=tc, tool=tool, args={},
|
||||
error="工具参数解析失败（JSON 不完整）",
|
||||
))
|
||||
continue
|
||||
# 🌟 工具自定义参数预处理（对照 pi tool.prepareArguments）
|
||||
raw_args = tc.arguments
|
||||
if tool.prepare_arguments is not None:
|
||||
try:
|
||||
raw_args = tool.prepare_arguments(raw_args)
|
||||
except Exception as e:
|
||||
prepared.append(PreparedToolCall(
|
||||
tool_call=tc, tool=tool, args=tc.arguments,
|
||||
error=f"参数预处理失败: {e}",
|
||||
))
|
||||
continue
|
||||
if not isinstance(raw_args, dict):
|
||||
prepared.append(PreparedToolCall(
|
||||
tool_call=tc, tool=tool, args=tc.arguments,
|
||||
error="参数预处理返回的不是对象",
|
||||
))
|
||||
continue
|
||||
# 🌟 归一化 + 轻量强制转换（null 可选字段删除 / "30" → 30）
|
||||
args = normalize_and_coerce(raw_args, tool.parameters)
|
||||
err = validate_json_schema(args, tool.parameters)
|
||||
if err:
|
||||
prepared.append(PreparedToolCall(
|
||||
tool_call=tc, tool=tool, args=args,
|
||||
error=f"参数校验失败: {err}；收到的参数: {_brief_json(raw_args)}",
|
||||
))
|
||||
continue
|
||||
prepared.append(PreparedToolCall(tool_call=tc, tool=tool,
|
||||
args=args))
|
||||
return prepared
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 执行阶段 —— 对照 agent-loop.ts executeTool
|
||||
# before 钩子 → 执行 → after 钩子 → 异常兜底
|
||||
# ======================================================================
|
||||
def execute_tool_call(prepared: PreparedToolCall,
|
||||
assistant: AgentMessage,
|
||||
config,
|
||||
signal: AbortSignal,
|
||||
on_update: Optional[Callable[[str], None]],
|
||||
on_timer: Optional[Callable[[int, int], None]] = None,
|
||||
) -> AgentToolResult:
|
||||
"""对照 agent-loop.ts executePreparedToolCall + prepareToolCall 的钩子语义
|
||||
🆕 on_timer(elapsed_s, timeout_s)：bash 运行中每秒滴一次（气泡读秒）"""
|
||||
tc = prepared.tool_call
|
||||
|
||||
# 准备阶段已失败 → 直接错误结果（pi 同款：不进入执行）
|
||||
if prepared.error:
|
||||
return AgentToolResult.text(prepared.error, is_error=True)
|
||||
if signal.aborted:
|
||||
return AgentToolResult.text("操作已中止 (Operation aborted)", is_error=True)
|
||||
|
||||
# before 钩子（对照 pi: 可修改 args / 拒绝 block / 请求 terminate）
|
||||
if config.before_tool_call:
|
||||
try:
|
||||
decision = config.before_tool_call(
|
||||
{"assistant_message": assistant, "tool_call": tc,
|
||||
"args": prepared.args, "context": config.tool_context},
|
||||
signal)
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"before 钩子异常: {e}", is_error=True)
|
||||
if signal.aborted:
|
||||
return AgentToolResult.text("操作已中止 (Operation aborted)",
|
||||
is_error=True)
|
||||
if decision:
|
||||
if decision.get("block"):
|
||||
# 🐛 修复：旧代码向 AgentToolResult.text() 传了不存在的 terminate 参数
|
||||
# （钩子一旦 block 就抛 TypeError）—— 改为直接构造，保留 terminate 语义
|
||||
return AgentToolResult(
|
||||
content=[{"type": "text",
|
||||
"text": decision.get("reason") or "工具执行被拦截"}],
|
||||
is_error=True,
|
||||
terminate=bool(decision.get("terminate")))
|
||||
if decision.get("args") is not None:
|
||||
# 🌟 对照 pi applyBeforeToolDecision：钩子改参后【重新校验】
|
||||
new_args = decision["args"]
|
||||
if not isinstance(new_args, dict):
|
||||
return AgentToolResult.text(
|
||||
"before 钩子返回的参数不是对象", is_error=True)
|
||||
new_args = normalize_and_coerce(new_args, prepared.tool.parameters)
|
||||
_verr = validate_json_schema(new_args, prepared.tool.parameters)
|
||||
if _verr:
|
||||
return AgentToolResult.text(
|
||||
f"before 钩子修改后的参数校验失败: {_verr}；"
|
||||
f"收到的参数: {_brief_json(decision['args'])}",
|
||||
is_error=True)
|
||||
prepared.args = new_args
|
||||
|
||||
t0 = time.time()
|
||||
try:
|
||||
# 🆕 秒级滴答回调注入 ctx（不污染调用方的 tool_context 字典）
|
||||
_ctx = dict(config.tool_context or {})
|
||||
if on_timer is not None:
|
||||
_ctx["on_timer"] = on_timer
|
||||
result = prepared.tool.execute(
|
||||
tc.id, prepared.args, signal, on_update, _ctx,
|
||||
)
|
||||
if not isinstance(result, AgentToolResult):
|
||||
# 宽容处理：工具返回 str 也接受
|
||||
result = AgentToolResult.text(str(result))
|
||||
except Exception as e:
|
||||
result = AgentToolResult.text(f"工具执行异常: {e}", is_error=True)
|
||||
|
||||
# after 钩子
|
||||
if config.after_tool_call:
|
||||
try:
|
||||
config.after_tool_call(tc, result, result.is_error)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 结果定型（对照 pi finalizeToolResult：确保 content 结构合法）
|
||||
if not result.content:
|
||||
result = AgentToolResult.text("(无输出)")
|
||||
result.details = {"duration_ms": int((time.time() - t0) * 1000),
|
||||
**(result.details or {} if isinstance(result.details, dict) else {})}
|
||||
return result
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 截断保护 —— 对照 agent-loop.ts failToolCallsFromTruncatedMessage
|
||||
# 助手消息被 length/aborted 截断时，其工具调用参数可能残缺：
|
||||
# 一律替换为错误 toolResult，绝不执行残缺调用。
|
||||
# ======================================================================
|
||||
def fail_tool_calls_from_truncated_message(assistant: AgentMessage,
|
||||
reason: str = "aborted") -> List[AgentMessage]:
|
||||
"""为每个工具调用生成错误 toolResult 消息（对照 pi 返回值的组装）"""
|
||||
msgs: List[AgentMessage] = []
|
||||
for tc in assistant.tool_calls:
|
||||
# 对照 pi: "Tool call {name} was not executed: the response hit the
|
||||
# output token limit before the arguments were complete..."
|
||||
text = (f"工具调用 {tc.name} 未执行（{reason}）："
|
||||
f"响应在参数完整前达到输出长度上限，部分参数已被丢弃。"
|
||||
f"请用完整参数重试该操作。")
|
||||
msgs.append(AgentMessage(
|
||||
role="toolResult",
|
||||
tool_call_id=tc.id,
|
||||
tool_name=tc.name,
|
||||
content=text,
|
||||
is_error=True,
|
||||
))
|
||||
return msgs
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 内置工具 —— 对照 pi coding-agent 的 read / bash / write / edit
|
||||
# ======================================================================
|
||||
_MAX_READ_LINES = 2000
|
||||
_MAX_OUTPUT_BYTES = 50 * 1024
|
||||
|
||||
|
||||
def _resolve_path(path: str, ctx: Dict[str, Any]) -> str:
|
||||
"""相对路径基于 tool_context 的 cwd（默认项目根）"""
|
||||
if os.path.isabs(path):
|
||||
return os.path.abspath(path)
|
||||
cwd = ctx.get("cwd", os.getcwd())
|
||||
return os.path.abspath(os.path.join(cwd, path))
|
||||
|
||||
|
||||
def tool_read(tool_call_id: str, args: Dict[str, Any], signal: AbortSignal,
|
||||
on_update, ctx: Dict[str, Any]) -> AgentToolResult:
|
||||
"""read: 读取文本文件（支持 offset/limit 行窗口），带行号输出"""
|
||||
if not args.get("path"):
|
||||
return AgentToolResult.text("path 不能为空", is_error=True)
|
||||
path = _resolve_path(args["path"], ctx)
|
||||
if not os.path.isfile(path):
|
||||
return AgentToolResult.text(f"文件不存在: {path}", is_error=True)
|
||||
try:
|
||||
offset = max(1, int(args.get("offset", 1)))
|
||||
except (TypeError, ValueError):
|
||||
return AgentToolResult.text("offset 必须是整数", is_error=True)
|
||||
try:
|
||||
# 🌟 limit 必须有下限：负数会被 python 负索引语义吃掉
|
||||
# （旧版 limit=-5 → lines[0:-5] → 除末尾 5 行外全部返回，与直觉完全相反）
|
||||
limit = max(1, min(int(args.get("limit", 2000)), _MAX_READ_LINES))
|
||||
except (TypeError, ValueError):
|
||||
return AgentToolResult.text("limit 必须是整数", is_error=True)
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8", errors="replace") as f:
|
||||
lines = f.readlines()
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"读取失败: {e}", is_error=True)
|
||||
total = len(lines)
|
||||
if total == 0:
|
||||
return AgentToolResult.text("[文件为空（0 行）]")
|
||||
if offset > total:
|
||||
return AgentToolResult.text(
|
||||
f"[起始行 offset={offset} 超出文件范围，该文件共 {total} 行]")
|
||||
chunk = lines[offset - 1: offset - 1 + limit]
|
||||
numbered = "".join(
|
||||
f"{i + offset:6d}\t{line}" for i, line in enumerate(chunk)
|
||||
)
|
||||
shown_hi = offset - 1 + len(chunk)
|
||||
footer = f"\n[已显示 {offset}–{shown_hi} 行，共 {total} 行]"
|
||||
if shown_hi < total:
|
||||
footer += f"（还有 {total - shown_hi} 行未显示，用 offset={shown_hi + 1} 继续）"
|
||||
out = numbered + footer
|
||||
if len(out.encode("utf-8")) > _MAX_OUTPUT_BYTES:
|
||||
out = out.encode("utf-8")[:_MAX_OUTPUT_BYTES].decode("utf-8", "ignore")
|
||||
out += "\n[输出超过 50KB 已截断]"
|
||||
return AgentToolResult.text(out)
|
||||
|
||||
|
||||
def tool_bash(tool_call_id: str, args: Dict[str, Any], signal: AbortSignal,
|
||||
on_update, ctx: Dict[str, Any]) -> AgentToolResult:
|
||||
"""bash: 执行 shell 命令（默认 120s 超时）。对照 pi: bash 是 sequential 工具
|
||||
|
||||
🆕 秒级滴答：Popen + communicate(timeout=1) 循环，每秒：
|
||||
① 日志 [timer] bash 读秒 N/Ts
|
||||
② on_timer(N, T) 推前端气泡读秒
|
||||
③ 到期杀进程树（Windows taskkill /T，连孤儿子进程一起杀）→ 真超时
|
||||
"""
|
||||
if signal.aborted:
|
||||
return AgentToolResult.text("操作已中止 (Operation aborted)", is_error=True)
|
||||
command = args.get("command", "")
|
||||
if not command.strip():
|
||||
return AgentToolResult.text("command 不能为空", is_error=True)
|
||||
timeout = min(float(args.get("timeout", 120)), 600)
|
||||
cwd = ctx.get("cwd") or os.getcwd()
|
||||
# 对照 pi bash 工具：命令始终经 shell 解释（支持管道/别名/内置命令）
|
||||
use_shell = ctx.get("shell", True)
|
||||
on_timer = ctx.get("on_timer")
|
||||
t0 = time.time()
|
||||
# 🆕 计时观察①：计时器启动时刻 + 模型实际传的 timeout 值
|
||||
_dbg_log(f"[timer] bash 开始 timeout={timeout:.0f}s "
|
||||
f"(显式={args.get('timeout')}) cmd={command[:80]!r}")
|
||||
|
||||
def _tick(elapsed_i: int):
|
||||
"""每秒一次：日志 + 推前端（静默吞异常，绝不影响执行）"""
|
||||
_dbg_log(f"[timer] bash 读秒 {elapsed_i}/{int(timeout)}s")
|
||||
try:
|
||||
if on_timer:
|
||||
on_timer(elapsed_i, int(timeout))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _kill_tree(proc):
|
||||
"""杀整个进程树：Windows 上默认 kill 只杀 cmd 壳，孤儿子进程继续
|
||||
持管道 → 假超时（设定10s 实际20s）。taskkill /T 整树杀。"""
|
||||
try:
|
||||
if os.name == "nt":
|
||||
subprocess.run(["taskkill", "/F", "/T", "/PID", str(proc.pid)],
|
||||
capture_output=True, timeout=10)
|
||||
else:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
try:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 🐛 修复：text=True 不带 encoding 时按系统码页（中文 Windows=GBK）解码，
|
||||
# 子进程输出 UTF-8（python/git/中文 echo）→ _readerthread UnicodeDecodeError。
|
||||
# 强制 UTF-8 + 容错替换；PYTHONIOENCODING 让 python 子进程也按 UTF-8 输出。
|
||||
_env = dict(os.environ, PYTHONIOENCODING="utf-8")
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
command, shell=use_shell,
|
||||
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||||
text=True, encoding="utf-8", errors="replace",
|
||||
cwd=cwd, env=_env,
|
||||
)
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"执行失败: {e}", is_error=True)
|
||||
|
||||
timed_out = False
|
||||
out_parts: List[str] = []
|
||||
err_parts: List[str] = []
|
||||
err_header_fed = False
|
||||
|
||||
# 🆕 实时输出：stdout/stderr 各起一个 reader 线程 → 队列，
|
||||
# 主循环（仍在 worker 线程内）抽干队列并回调 on_update。
|
||||
# 绝不从 reader 线程直接回调 —— UI 侧 _on_tool_updated 会改 timeline，跨线程不安全。
|
||||
_q: "queue.Queue" = queue.Queue()
|
||||
|
||||
def _reader(stream, tag):
|
||||
try:
|
||||
for line in iter(stream.readline, ""):
|
||||
_q.put((tag, line))
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
stream.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
_t_out = threading.Thread(target=_reader, args=(proc.stdout, "out"), daemon=True)
|
||||
_t_err = threading.Thread(target=_reader, args=(proc.stderr, "err"), daemon=True)
|
||||
_t_out.start()
|
||||
_t_err.start()
|
||||
|
||||
def _feed(tag, chunk):
|
||||
"""抽到一块输出：累积 + 推实时流（只在 worker 线程内调用）"""
|
||||
nonlocal err_header_fed
|
||||
piece = chunk
|
||||
if tag == "err":
|
||||
err_parts.append(chunk)
|
||||
if not err_header_fed:
|
||||
err_header_fed = True
|
||||
piece = "[stderr]\n" + chunk
|
||||
else:
|
||||
out_parts.append(chunk)
|
||||
try:
|
||||
if on_update and piece:
|
||||
on_update(piece)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _drain():
|
||||
"""抽干队列里已到达的输出（单轮上限防极端刷屏卡死）"""
|
||||
for _ in range(500):
|
||||
try:
|
||||
tag, chunk = _q.get_nowait()
|
||||
except queue.Empty:
|
||||
return
|
||||
_feed(tag, chunk)
|
||||
|
||||
last_tick = 0 # 首次滴答仍在 elapsed=1（与旧 communicate 实现一致，不在 0 秒多滴一次）
|
||||
try:
|
||||
while True:
|
||||
# ① 最多 0.2s 醒一次 → 把新输出实时推给前端
|
||||
try:
|
||||
tag, chunk = _q.get(timeout=0.2)
|
||||
_feed(tag, chunk)
|
||||
except queue.Empty:
|
||||
pass
|
||||
_drain()
|
||||
|
||||
# ② 进程已退出：等 reader 读完管道残余 → 收尾
|
||||
if proc.poll() is not None:
|
||||
_t_out.join(timeout=1.0)
|
||||
_t_err.join(timeout=1.0)
|
||||
_drain()
|
||||
break
|
||||
|
||||
# ③ 每秒读秒 + 中止/超时判定（原语义不变）
|
||||
elapsed_i = int(time.time() - t0)
|
||||
if elapsed_i > last_tick:
|
||||
last_tick = elapsed_i
|
||||
if signal.aborted:
|
||||
_kill_tree(proc)
|
||||
_t_out.join(timeout=1.0)
|
||||
_t_err.join(timeout=1.0)
|
||||
_drain()
|
||||
return AgentToolResult.text(
|
||||
"操作已中止 (Operation aborted)", is_error=True)
|
||||
_tick(elapsed_i)
|
||||
if time.time() - t0 >= timeout:
|
||||
# 到期 → 真杀（进程树）
|
||||
timed_out = True
|
||||
_kill_tree(proc)
|
||||
try:
|
||||
proc.wait(timeout=15)
|
||||
except Exception:
|
||||
pass
|
||||
_t_out.join(timeout=5.0)
|
||||
_t_err.join(timeout=5.0)
|
||||
_drain()
|
||||
break
|
||||
except Exception as e:
|
||||
try:
|
||||
_kill_tree(proc)
|
||||
except Exception:
|
||||
pass
|
||||
return AgentToolResult.text(f"执行失败: {e}", is_error=True)
|
||||
|
||||
out = "".join(out_parts)
|
||||
err = "".join(err_parts)
|
||||
code = proc.returncode if proc.returncode is not None else -1
|
||||
|
||||
dur = time.time() - t0
|
||||
if timed_out:
|
||||
# 🆕 计时观察②：计时器到期
|
||||
_dbg_log(f"[timer] bash 超时触发 设定={timeout:.0f}s "
|
||||
f"实际={dur:.1f}s cmd={command[:80]!r}")
|
||||
return AgentToolResult.text(f"命令超时（>{timeout:.0f}s）已终止", is_error=True)
|
||||
# 🆕 计时观察③：正常结束 + 实际耗时
|
||||
_dbg_log(f"[timer] bash 正常结束 dur={dur:.1f}s "
|
||||
f"exit={code} 设定timeout={timeout:.0f}s")
|
||||
result = f"$ {command}\n"
|
||||
if out:
|
||||
result += out if out.endswith("\n") else out + "\n"
|
||||
if err:
|
||||
result += f"[stderr]\n{err}"
|
||||
result += f"\n[exit {code}] ({dur:.1f}s)"
|
||||
if len(result.encode("utf-8")) > _MAX_OUTPUT_BYTES:
|
||||
result = result.encode("utf-8")[:_MAX_OUTPUT_BYTES].decode("utf-8", "ignore")
|
||||
result += "\n[输出超过 50KB 已截断]"
|
||||
return AgentToolResult.text(result, is_error=code != 0,
|
||||
details={"exit_code": code})
|
||||
|
||||
|
||||
def _atomic_write(path: str, content: str) -> None:
|
||||
"""原子写：同目录临时文件 + os.replace（避免半截文件）
|
||||
|
||||
保留原有 newline 语义（默认 None → 平台换行翻译），仅增加原子性。
|
||||
"""
|
||||
d = os.path.dirname(os.path.abspath(path)) or "."
|
||||
fd, tmp = tempfile.mkstemp(dir=d, prefix=".hocode_w_", suffix=".tmp")
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
os.replace(tmp, path) # Windows/POSIX 均为原子替换
|
||||
except BaseException:
|
||||
try:
|
||||
os.remove(tmp)
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
def tool_write(tool_call_id: str, args: Dict[str, Any], signal: AbortSignal,
|
||||
on_update, ctx: Dict[str, Any]) -> AgentToolResult:
|
||||
"""write: 创建/覆盖文件（自动建父目录）"""
|
||||
if not args.get("path"):
|
||||
return AgentToolResult.text("path 不能为空", is_error=True)
|
||||
path = _resolve_path(args["path"], ctx)
|
||||
content = args.get("content", "")
|
||||
if content is None:
|
||||
content = ""
|
||||
if not isinstance(content, str):
|
||||
content = str(content)
|
||||
try:
|
||||
parent = os.path.dirname(path)
|
||||
if parent:
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
_atomic_write(path, content)
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"写入失败: {e}", is_error=True)
|
||||
return AgentToolResult.text(f"已写入 {len(content)} 字符 → {path}")
|
||||
|
||||
|
||||
def _prepare_edit_args(args: Any) -> Any:
|
||||
"""归一化 edit 参数（对照 pi edit.ts:56-70 prepareEditArguments）
|
||||
|
||||
- edits 为 JSON 字符串 → 解析为数组
|
||||
- edits 内条目为 JSON 字符串 → 解析为对象
|
||||
- 兼容 legacy 单条形式 {path, oldText, newText} → 包装为 edits:[{...}]
|
||||
"""
|
||||
if not isinstance(args, dict):
|
||||
return args
|
||||
out = dict(args)
|
||||
edits = out.get("edits")
|
||||
if isinstance(edits, str):
|
||||
try:
|
||||
parsed = json.loads(edits)
|
||||
if isinstance(parsed, list):
|
||||
out["edits"] = parsed
|
||||
elif isinstance(parsed, dict):
|
||||
out["edits"] = [parsed]
|
||||
except Exception:
|
||||
pass # 解析不了就交给校验层报错
|
||||
if out.get("edits") is None and ("oldText" in out or "newText" in out):
|
||||
ed: Dict[str, Any] = {}
|
||||
if "oldText" in out:
|
||||
ed["oldText"] = out["oldText"]
|
||||
if "newText" in out:
|
||||
ed["newText"] = out["newText"]
|
||||
out = {"path": out.get("path"), "edits": [ed]}
|
||||
if isinstance(out.get("edits"), list):
|
||||
norm = []
|
||||
for it in out["edits"]:
|
||||
if isinstance(it, str):
|
||||
try:
|
||||
it = json.loads(it)
|
||||
except Exception:
|
||||
pass
|
||||
norm.append(it)
|
||||
out["edits"] = norm
|
||||
return out
|
||||
|
||||
|
||||
def _plan_edits(content: str, edits: List[Any]):
|
||||
"""对【原始内容】定位每条 oldText（要求恰好 1 次）+ 区间重叠检测
|
||||
|
||||
对照 pi edit-diff.ts:348「edits[i] and edits[j] overlap … Merge them into one edit」。
|
||||
返回 (错误信息, [(start, end, index, new_text), ...] 已按 start 排序)
|
||||
"""
|
||||
spans: List[tuple] = []
|
||||
for i, ed in enumerate(edits):
|
||||
if not isinstance(ed, dict):
|
||||
return f"第 {i + 1} 条 edit 不是对象", None
|
||||
old = ed.get("oldText", "")
|
||||
if not isinstance(old, str):
|
||||
return f"第 {i + 1} 条 edit 的 oldText 必须是字符串", None
|
||||
if old == "":
|
||||
return f"第 {i + 1} 条 edit 的 oldText 不能为空", None
|
||||
c = content.count(old)
|
||||
if c == 0:
|
||||
return f"第 {i + 1} 条 edit 未找到匹配文本（oldText 不存在或已变化）", None
|
||||
if c > 1:
|
||||
return f"第 {i + 1} 条 edit 匹配到 {c} 处（要求唯一），请提供更长的上下文", None
|
||||
start = content.index(old)
|
||||
spans.append((start, start + len(old), i, ed.get("newText", "")))
|
||||
ordered = sorted(spans)
|
||||
for a, b in zip(ordered, ordered[1:]):
|
||||
if b[0] < a[1]:
|
||||
return (f"edits[{a[2]}] 与 edits[{b[2]}] 区域重叠，"
|
||||
f"请合并为一条 edit 或改为互不相交的修改"), None
|
||||
return None, ordered
|
||||
|
||||
|
||||
def tool_edit(tool_call_id: str, args: Dict[str, Any], signal: AbortSignal,
|
||||
on_update, ctx: Dict[str, Any]) -> AgentToolResult:
|
||||
"""edit: 精确文本替换（edits: [{oldText, newText}]）
|
||||
|
||||
🌟 语义对齐 pi：全部 edit 都对【原始文件内容】定位，要求各自唯一且区间互不重叠，
|
||||
然后按偏移一次性重建（不是逐条 replace 的增量语义）。
|
||||
"""
|
||||
if not args.get("path"):
|
||||
return AgentToolResult.text("path 不能为空", is_error=True)
|
||||
path = _resolve_path(args["path"], ctx)
|
||||
if not os.path.isfile(path):
|
||||
return AgentToolResult.text(f"文件不存在: {path}", is_error=True)
|
||||
edits = args.get("edits", [])
|
||||
if not isinstance(edits, list) or not edits:
|
||||
return AgentToolResult.text("edits 不能为空", is_error=True)
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"读取失败: {e}", is_error=True)
|
||||
err, ordered = _plan_edits(content, edits)
|
||||
if err:
|
||||
return AgentToolResult.text(err, is_error=True)
|
||||
# 按偏移从后往前替换（前面的偏移不受影响）
|
||||
out = content
|
||||
for start, end, _i, new_text in reversed(ordered):
|
||||
out = out[:start] + (new_text if isinstance(new_text, str) else str(new_text)) + out[end:]
|
||||
try:
|
||||
_atomic_write(path, out)
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"编辑失败: {e}", is_error=True)
|
||||
return AgentToolResult.text(f"已应用 {len(ordered)} 处编辑 → {path}")
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 默认工具集 —— 对照 pi 默认启用 read/bash/edit/write
|
||||
# ======================================================================
|
||||
|
||||
_TEXT_TOOL_RE = re.compile(r"<(bash|read)>(.*?)</\1>", re.DOTALL)
|
||||
|
||||
|
||||
def parse_text_tool_calls(content: str):
|
||||
"""
|
||||
兜底解析（haocode 扩展，pi 无此层）：
|
||||
对不支持 tools API 的供应商/模型——它们会把工具调用用纯文字"演"出来
|
||||
（例如 <bash>ls</bash>、<read>路径</read>），本函数识别单参数工具
|
||||
bash / read 并转成真 ToolCall 供循环执行。
|
||||
write / edit 参数多、文本歧义大，不做兜底（保持安全）。
|
||||
返回 (原文, [ToolCall])；未命中时 calls 为空列表。
|
||||
"""
|
||||
if not content or "<" not in content:
|
||||
return content, []
|
||||
calls: List[ToolCall] = []
|
||||
|
||||
def _sub(m):
|
||||
name, payload = m.group(1), m.group(2).strip()
|
||||
if not payload:
|
||||
return m.group(0)
|
||||
if name == "bash":
|
||||
calls.append(ToolCall(id=new_id("txtcall"), name="bash",
|
||||
arguments={"command": payload}))
|
||||
elif name == "read":
|
||||
calls.append(ToolCall(id=new_id("txtcall"), name="read",
|
||||
arguments={"path": payload}))
|
||||
return m.group(0) # 保留原文（UI 已渲染，不回改；执行由 tool_calls 驱动）
|
||||
|
||||
cleaned = _TEXT_TOOL_RE.sub(_sub, content)
|
||||
return cleaned, calls
|
||||
|
||||
def default_tools() -> List[AgentTool]:
|
||||
return [
|
||||
AgentTool(
|
||||
name="read", label="读取文件",
|
||||
description="读取文本文件内容（带行号）。支持 offset/limit 按行窗口读取大文件。",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "文件路径"},
|
||||
"offset": {"type": "integer", "description": "起始行号（从 1 开始）"},
|
||||
"limit": {"type": "integer", "description": "最多读取行数（默认 2000）"},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
execute=tool_read,
|
||||
),
|
||||
AgentTool(
|
||||
name="bash", label="执行命令",
|
||||
description="执行 shell 命令并返回 stdout/stderr/退出码。默认 120 秒超时。",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {"type": "string", "description": "要执行的命令"},
|
||||
"timeout": {"type": "number", "description": "超时秒数（默认 120，最大 600）"},
|
||||
},
|
||||
"required": ["command"],
|
||||
},
|
||||
execute=tool_bash,
|
||||
execution_mode="sequential", # 对照 pi: bash 整批串行
|
||||
),
|
||||
AgentTool(
|
||||
name="write", label="写入文件",
|
||||
description="创建或覆盖写入文件（自动创建父目录）。",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "文件路径"},
|
||||
"content": {"type": "string", "description": "文件内容"},
|
||||
},
|
||||
"required": ["path", "content"],
|
||||
},
|
||||
execute=tool_write,
|
||||
),
|
||||
AgentTool(
|
||||
name="edit", label="编辑文件",
|
||||
description=("对文件做精确文本替换。edits 中每条 oldText 必须在原文件中唯一，"
|
||||
"且各条区间不得重叠（重叠请合并为一条）。"),
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "文件路径"},
|
||||
"edits": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"oldText": {"type": "string"},
|
||||
"newText": {"type": "string"},
|
||||
},
|
||||
"required": ["oldText"],
|
||||
},
|
||||
"description": "替换操作列表",
|
||||
},
|
||||
},
|
||||
"required": ["path", "edits"],
|
||||
},
|
||||
execute=tool_edit,
|
||||
prepare_arguments=_prepare_edit_args,
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,331 @@
|
||||
"""
|
||||
core/agent/types.py
|
||||
====================
|
||||
🌟 pi agent-core 的 Python 1:1 移植 —— 数据模型层
|
||||
|
||||
对照 pi-main 源码：
|
||||
packages/agent/src/types.ts (443 行)
|
||||
- AgentMessage 联合类型 (user / assistant / toolResult)
|
||||
- AgentEvent 11 种事件
|
||||
- AgentTool / AgentToolCall / AgentToolResult
|
||||
- AgentState / AgentConfig / AgentError
|
||||
|
||||
设计原则（与 pi 完全一致）：
|
||||
1. 消息是不可变事实（append-only），循环只追加、不修改
|
||||
2. 事件是唯一对外输出（TUI/Qt 都只是事件订阅者）
|
||||
3. 工具 = (name, description, JSON-Schema 参数, execute 函数)
|
||||
4. 中止用 AbortSignal 标志位，不用异常控制流
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
|
||||
def _now_ms() -> int:
|
||||
return int(time.time() * 1000)
|
||||
|
||||
|
||||
def new_id(prefix: str) -> str:
|
||||
return f"{prefix}_{uuid.uuid4().hex[:16]}"
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 中止信号 —— 对照 pi: AbortController/AbortSignal（Web 标准）
|
||||
# pi 里 streamSimple 每消费一个 chunk 检查 signal.aborted
|
||||
# ======================================================================
|
||||
class AbortSignal:
|
||||
def __init__(self):
|
||||
self.aborted: bool = False
|
||||
self.reason: str = ""
|
||||
|
||||
def abort(self, reason: str = "aborted"):
|
||||
self.aborted = True
|
||||
self.reason = reason
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 消息模型 —— 对照 types.ts 的 AgentMessage 联合类型
|
||||
# pi: type AgentMessage = AgentUserMessage | AgentAssistantMessage | AgentToolResultMessage
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""pi: AssistantMessage 内容块里的 {type:"toolCall", id, name, arguments}"""
|
||||
id: str
|
||||
name: str
|
||||
arguments: Dict[str, Any] = field(default_factory=dict)
|
||||
# 流式阶段累积的原始 JSON 字符串（参数被截断时可能是残缺 JSON）
|
||||
raw_arguments: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentMessage:
|
||||
"""
|
||||
统一消息容器。role 决定哪些字段有效：
|
||||
user -> content: str 或 OpenAI 多模态 list
|
||||
assistant -> content(str 正文) + reasoning(思考) + tool_calls + stop_reason + usage
|
||||
toolResult -> tool_call_id + tool_name + content + is_error
|
||||
与 pi 的三种 Message 类型字段一一对应。
|
||||
"""
|
||||
role: str # "user" | "assistant" | "toolResult"
|
||||
content: Any = "" # str | List[dict]
|
||||
reasoning: str = "" # pi: thinking block（vLLM: reasoning_content）
|
||||
tool_calls: List[ToolCall] = field(default_factory=list)
|
||||
tool_call_id: str = "" # toolResult 专用
|
||||
tool_name: str = "" # toolResult 专用
|
||||
is_error: bool = False # toolResult 专用
|
||||
stop_reason: str = "" # assistant 专用: stop|length|aborted|error
|
||||
error_message: str = "" # assistant 出错时的说明
|
||||
usage: Dict[str, Any] = field(default_factory=dict) # {input, output, cacheRead...}
|
||||
kind: str = "" # "" | "compaction_summary"（对照 pi compaction 条目）
|
||||
db_msg_id: str = "" # 🆕 压缩持久化：对应 DB 行的 id（build_api_context 注入，切点计算用）
|
||||
timestamp: int = field(default_factory=_now_ms)
|
||||
id: str = field(default_factory=lambda: new_id("msg"))
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"role": self.role, "content": self.content, "reasoning": self.reasoning,
|
||||
"tool_calls": [
|
||||
{"id": tc.id, "name": tc.name, "arguments": tc.arguments,
|
||||
"raw_arguments": tc.raw_arguments}
|
||||
for tc in self.tool_calls
|
||||
],
|
||||
"tool_call_id": self.tool_call_id, "tool_name": self.tool_name,
|
||||
"is_error": self.is_error, "stop_reason": self.stop_reason,
|
||||
"error_message": self.error_message, "usage": self.usage,
|
||||
"kind": self.kind,
|
||||
"db_msg_id": self.db_msg_id,
|
||||
"timestamp": self.timestamp, "id": self.id,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: Dict[str, Any]) -> "AgentMessage":
|
||||
return AgentMessage(
|
||||
role=d.get("role", "user"), content=d.get("content", ""),
|
||||
reasoning=d.get("reasoning", ""),
|
||||
tool_calls=[ToolCall(id=t.get("id", ""), name=t.get("name", ""),
|
||||
arguments=t.get("arguments", {}),
|
||||
raw_arguments=t.get("raw_arguments", ""))
|
||||
for t in d.get("tool_calls", [])],
|
||||
tool_call_id=d.get("tool_call_id", ""), tool_name=d.get("tool_name", ""),
|
||||
is_error=d.get("is_error", False), stop_reason=d.get("stop_reason", ""),
|
||||
error_message=d.get("error_message", ""), usage=d.get("usage", {}),
|
||||
kind=d.get("kind", ""),
|
||||
db_msg_id=d.get("db_msg_id", ""),
|
||||
timestamp=d.get("timestamp", 0), id=d.get("id") or new_id("msg"),
|
||||
)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 助手消息的流式增量事件 —— 对照 pi: AssistantMessageEvent
|
||||
# {type:"text_delta", textDelta} | {type:"thinking_delta", ...}
|
||||
# | {type:"toolcall_delta", ...} | {type:"done"} | {type:"error"}
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class AssistantMessageEvent:
|
||||
type: str # "text_delta" | "thinking_delta" | "toolcall_delta" | "done" | "error"
|
||||
text: str = "" # text_delta / thinking_delta 的增量
|
||||
tool_call_index: int = -1 # toolcall_delta 属于第几个 toolCall
|
||||
tool_call_delta: str = "" # toolcall_delta 的原始增量片段
|
||||
tool_call_field: str = "arguments" # 片段属于 id | name | arguments 哪一段
|
||||
error_message: str = "" # error 时的说明
|
||||
|
||||
@staticmethod
|
||||
def text_delta(text: str) -> "AssistantMessageEvent":
|
||||
return AssistantMessageEvent(type="text_delta", text=text)
|
||||
|
||||
@staticmethod
|
||||
def thinking_delta(text: str) -> "AssistantMessageEvent":
|
||||
return AssistantMessageEvent(type="thinking_delta", text=text)
|
||||
|
||||
@staticmethod
|
||||
def toolcall_delta(index: int, delta: str,
|
||||
field: str = "arguments") -> "AssistantMessageEvent":
|
||||
return AssistantMessageEvent(type="toolcall_delta", tool_call_index=index,
|
||||
tool_call_delta=delta, tool_call_field=field)
|
||||
|
||||
@staticmethod
|
||||
def done() -> "AssistantMessageEvent":
|
||||
return AssistantMessageEvent(type="done")
|
||||
|
||||
@staticmethod
|
||||
def error(message: str) -> "AssistantMessageEvent":
|
||||
return AssistantMessageEvent(type="error", error_message=message)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# Agent 事件（11 种）—— 对照 types.ts 的 AgentEvent 联合类型
|
||||
# ======================================================================
|
||||
EVENT_TYPES = (
|
||||
"agent_start", "agent_end",
|
||||
"turn_start", "turn_end",
|
||||
"message_start", "message_update", "message_end",
|
||||
"tool_execution_start", "tool_execution_update", "tool_execution_end",
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentEvent:
|
||||
"""
|
||||
事件容器。字段随 type 变化：
|
||||
agent_end -> stop_reason, error
|
||||
message_start/end -> message (AgentMessage)
|
||||
message_update -> message(当前累积快照), assistant_message_event
|
||||
tool_execution_start -> tool_call, tool_name
|
||||
tool_execution_update -> tool_call, arg(增量)
|
||||
tool_execution_timer -> tool_call, arg=(已耗秒, 超时秒) 🆕 bash 读秒
|
||||
tool_execution_end -> tool_call, result, is_error
|
||||
"""
|
||||
type: str
|
||||
message: Optional[AgentMessage] = None
|
||||
assistant_message_event: Optional[AssistantMessageEvent] = None
|
||||
tool_call: Optional[ToolCall] = None
|
||||
arg: str = ""
|
||||
result: Optional["AgentToolResult"] = None
|
||||
is_error: bool = False
|
||||
stop_reason: str = ""
|
||||
error: Optional["AgentError"] = None
|
||||
messages: List[AgentMessage] = field(default_factory=list) # agent_end 携带的本轮新消息
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 工具 —— 对照 types.ts: AgentTool / AgentToolCall / AgentToolResult
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class AgentToolResult:
|
||||
"""pi: {content: OutputContent[], details?, isError, terminate?}"""
|
||||
content: List[Dict[str, Any]] = field(default_factory=list) # [{type:"text", text}]
|
||||
details: Any = None
|
||||
is_error: bool = False
|
||||
terminate: bool = False # 对照 pi: 工具可请求终止整个批次/循环
|
||||
|
||||
@staticmethod
|
||||
def text(text: str, is_error: bool = False, details: Any = None) -> "AgentToolResult":
|
||||
return AgentToolResult(content=[{"type": "text", "text": text}],
|
||||
details=details, is_error=is_error)
|
||||
|
||||
def as_text(self) -> str:
|
||||
return "".join(c.get("text", "") for c in self.content if c.get("type") == "text")
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentTool:
|
||||
"""
|
||||
pi: interface AgentTool<T> { name; description; parameters(JSON Schema);
|
||||
execute(toolCallId, args, signal, onUpdate, context) => Promise<ToolResult> }
|
||||
Python 版 execute 签名完全一致（同步执行，循环里用线程池并发）。
|
||||
"""
|
||||
name: str
|
||||
description: str
|
||||
parameters: Dict[str, Any] # JSON Schema
|
||||
execute: Callable # (tool_call_id, args, signal, on_update, context) -> AgentToolResult
|
||||
label: str = "" # UI 展示用短标签
|
||||
# 对照 pi types.ts: executionMode —— "sequential" 工具会把整个批次拉回串行执行
|
||||
execution_mode: str = "parallel"
|
||||
# 🌟 对照 pi AgentHarnessTool.prepareArguments：校验前的确定性参数预处理
|
||||
# （如 edit 兼容 legacy 单条形式 / edits 为 JSON 字符串），签名 (args) -> args
|
||||
prepare_arguments: Optional[Callable] = None
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 运行时状态 / 配置 —— 对照 types.ts: AgentState / AgentConfig
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class AgentState:
|
||||
messages: List[AgentMessage] = field(default_factory=list)
|
||||
is_streaming: bool = False
|
||||
streaming_message: Optional[AgentMessage] = None # 正在流式生成的助手消息
|
||||
streaming_delta: Dict[str, str] = field(default_factory=dict) # 各通道当前增量缓冲
|
||||
error: Optional["AgentError"] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetryConfig:
|
||||
"""对照 pi settings-manager.ts retry 默认值（行 32-33）：
|
||||
maxRetries=3, baseDelayMs=2000，指数退避 2s → 4s → 8s（base × 2^(n-1)）
|
||||
max_delay_ms=0 表示不设上限（pi _prepareRetry 无封顶）"""
|
||||
max_attempts: int = 3
|
||||
base_delay_ms: int = 2000
|
||||
factor: float = 2.0
|
||||
max_delay_ms: int = 0 # 0 = 不封顶
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelConfig:
|
||||
"""
|
||||
对照 pi models.json 里的模型定义。
|
||||
context_window = pi 的 contextWindow；max_tokens = pi 的 maxTokens。
|
||||
"""
|
||||
provider: str = ""
|
||||
name: str = ""
|
||||
context_window: int = 128000
|
||||
max_tokens: int = 8192
|
||||
# 请求参数
|
||||
temperature: float = 0.7
|
||||
timeout_seconds: float = 180.0
|
||||
api_key: str = ""
|
||||
base_url: str = "https://api.openai.com/v1"
|
||||
# 🆕 thinking 模式（DeepSeek 等）：assistant 消息必须把 reasoning_content
|
||||
# 原样回传，否则 400。默认开启（只影响有 reasoning 的轮中消息；
|
||||
# DB 回放行 reasoning="" 不带字段 → 全新请求体不变）。
|
||||
pass_reasoning: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentConfig:
|
||||
"""
|
||||
对照 types.ts: AgentConfig
|
||||
model / tools / systemPrompt / maxTurns / steeringMode / followUpMode
|
||||
+ 6 个循环钩子（transformContext / beforeToolCall / afterToolCall /
|
||||
shouldStopAfterTurn / prepareNextTurn）
|
||||
"""
|
||||
model: ModelConfig
|
||||
system_prompt: str = ""
|
||||
tools: List[AgentTool] = field(default_factory=list)
|
||||
max_turns: int = 50 # pi 默认 50
|
||||
steering_mode: str = "one-at-a-time" # pi 默认 "one-at-a-time"（"all" 一次性全注入）
|
||||
follow_up_mode: str = "one-at-a-time"
|
||||
# 对照 pi types.ts ToolExecutionMode: 批次级执行模式（有 sequential 工具自动转串行）
|
||||
tool_execution: str = "parallel" # "sequential" | "parallel"
|
||||
retry: RetryConfig = field(default_factory=RetryConfig)
|
||||
# 压缩（🌟 1:1 对照 pi DEFAULT_COMPACTION_SETTINGS）
|
||||
compaction_reserve: int = 16384 # reserveTokens: 摘要提示词与输出预留
|
||||
compaction_keep_recent: int = 20000 # keepRecentTokens: 压缩后保留的近期上下文预算
|
||||
# 钩子（None = 无钩子，对照 pi 的可选字段）
|
||||
transform_context: Optional[Callable] = None # (messages) -> messages
|
||||
before_tool_call: Optional[Callable] = None # (tool_call, args) -> 可修改/拒绝
|
||||
after_tool_call: Optional[Callable] = None # (tool_call, result, is_error) -> None
|
||||
should_stop_after_turn: Optional[Callable] = None # (messages) -> bool
|
||||
prepare_next_turn: Optional[Callable] = None # (context) -> context
|
||||
# 🆕 haocode 增强（偏离 pi 1:1）：轮中请求前的主动压缩检查
|
||||
# (messages) -> messages|None：返回新列表 = 发生了压缩（循环需同步上下文）；
|
||||
# None = 不需要/不可压缩（原样发请求，响应式安全网仍在）
|
||||
compact_fn: Optional[Callable] = None # (messages) -> messages|None
|
||||
# 运行上下文（透传给 tool.execute，如工作目录）
|
||||
tool_context: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentError(Exception):
|
||||
"""对照 pi 的 AgentError —— 携带分类信息供恢复逻辑判定"""
|
||||
message: str = ""
|
||||
# 分类: "rate_limit" | "overload" | "server_error" | "timeout" |
|
||||
# "connection" | "auth" | "unknown"
|
||||
kind: str = "unknown"
|
||||
status_code: Optional[int] = None
|
||||
recoverable: bool = False # 恢复逻辑判定后的标记
|
||||
|
||||
def __str__(self):
|
||||
return self.message or "AgentError"
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 运行结果 —— 对照 agent-session.ts 里对 agent_end 事件的读取
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class RunResult:
|
||||
stop_reason: str = "stop" # stop | length | aborted | error
|
||||
error: Optional[AgentError] = None
|
||||
message_count: int = 0
|
||||
@@ -0,0 +1,715 @@
|
||||
import sqlite3
|
||||
import os
|
||||
import sys
|
||||
import re
|
||||
import json
|
||||
import uuid
|
||||
import time
|
||||
import shutil
|
||||
from typing import List, Dict, Optional
|
||||
|
||||
# 🝙 P0 fix: DB path based on the file's location (no longer depends on the process's working directory)
|
||||
# 🌟 打包（PyInstaller onedir）时：优先与源码树共用 data/chat_history.db（历史不丢失）；
|
||||
# 若 exe 被复制到源码树之外 → 回退为 exe 旁的 data/chat_history.db
|
||||
if getattr(sys, "frozen", False):
|
||||
_exe_dir = os.path.dirname(os.path.abspath(sys.executable))
|
||||
_shared = os.path.abspath(os.path.join(_exe_dir, "..", "..", "data", "chat_history.db"))
|
||||
if os.path.isdir(os.path.dirname(_shared)):
|
||||
_DEFAULT_DB = _shared
|
||||
else:
|
||||
_DEFAULT_DB = os.path.join(_exe_dir, "data", "chat_history.db")
|
||||
else:
|
||||
_DEFAULT_DB = os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)), "..", "data", "chat_history.db")
|
||||
|
||||
|
||||
class _AutoCloseConn:
|
||||
"""🝙 P0 fix: wraps sqlite3.Connection.
|
||||
|
||||
sqlite3.Connection's with only handles commit/rollback and does NOT close.
|
||||
All 19 call sites use `with self.get_connection() as conn:`, so we auto-close on with exit.
|
||||
"""
|
||||
|
||||
def __init__(self, conn):
|
||||
self._conn = conn
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._conn, name)
|
||||
|
||||
def __enter__(self):
|
||||
self._conn.__enter__()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
try:
|
||||
return self._conn.__exit__(exc_type, exc, tb)
|
||||
finally:
|
||||
try:
|
||||
self._conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class DBManager:
|
||||
def __init__(self, db_path: str = None):
|
||||
# 🝙 Default path is anchored to the core/ directory, not affected by launch CWD
|
||||
self.db_path = db_path or os.path.abspath(_DEFAULT_DB)
|
||||
|
||||
# 🆕 附件/媒体文件根目录（copy_session 深拷贝磁盘文件用）。默认=项目根；测试可覆盖。
|
||||
self.files_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
self.is_first_run = not os.path.exists(self.db_path)
|
||||
d = os.path.dirname(self.db_path)
|
||||
if d:
|
||||
os.makedirs(d, exist_ok=True)
|
||||
|
||||
self._init_db()
|
||||
|
||||
def get_connection(self):
|
||||
conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
return _AutoCloseConn(conn)
|
||||
|
||||
def _init_db(self):
|
||||
"""初始化表结构并注入默认数据"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 1. 创建 sessions 表 (新增 current_leaf_msg_id)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS sessions (
|
||||
id TEXT PRIMARY KEY,
|
||||
title TEXT,
|
||||
created_at INTEGER,
|
||||
updated_at INTEGER,
|
||||
has_messages BOOLEAN DEFAULT 0,
|
||||
sort_order INTEGER DEFAULT 0,
|
||||
is_starred BOOLEAN DEFAULT 0,
|
||||
current_leaf_msg_id TEXT
|
||||
)
|
||||
""")
|
||||
|
||||
# 2. 创建 messages 表 (新增 parent_id)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id TEXT PRIMARY KEY,
|
||||
session_id TEXT,
|
||||
role TEXT,
|
||||
content TEXT,
|
||||
reasoning TEXT,
|
||||
is_ignored BOOLEAN,
|
||||
created_at INTEGER,
|
||||
attachment_metadata TEXT,
|
||||
parent_id TEXT,
|
||||
timeline TEXT,
|
||||
usage TEXT,
|
||||
stop_reason TEXT,
|
||||
error_message TEXT,
|
||||
FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE
|
||||
)
|
||||
""")
|
||||
|
||||
# ==========================================
|
||||
# 🌟 核心性能优化：为高频查询的字段建立索引
|
||||
# ==========================================
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_session_id ON messages(session_id)")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_parent_id ON messages(parent_id)")
|
||||
|
||||
# 3. 🌟 自动化热升级：检测并兼容旧数据库
|
||||
self._upgrade_schema(cursor)
|
||||
|
||||
# 4. 检查是否需要插入初始默认对话
|
||||
cursor.execute("SELECT COUNT(*) FROM sessions")
|
||||
if cursor.fetchone()[0] == 0:
|
||||
self._seed_default_chat(cursor)
|
||||
|
||||
conn.commit()
|
||||
|
||||
|
||||
def _upgrade_schema(self, cursor):
|
||||
"""检测缺少的新字段并自动补齐，如果是刚升级，则自动将旧线性数据串联成链表"""
|
||||
upgraded = False
|
||||
|
||||
# 兼容 sessions 字段
|
||||
for col in ["has_messages", "sort_order", "is_starred"]:
|
||||
try:
|
||||
cursor.execute(f"SELECT {col} FROM sessions LIMIT 1")
|
||||
except sqlite3.OperationalError:
|
||||
if col == "has_messages": cursor.execute("ALTER TABLE sessions ADD COLUMN has_messages BOOLEAN DEFAULT 0")
|
||||
if col == "sort_order": cursor.execute("ALTER TABLE sessions ADD COLUMN sort_order INTEGER DEFAULT 0")
|
||||
if col == "is_starred": cursor.execute("ALTER TABLE sessions ADD COLUMN is_starred BOOLEAN DEFAULT 0")
|
||||
|
||||
# 🌟 核心：兼容链表树架构
|
||||
try:
|
||||
# 兼容 sessions.mode 列（chat/worker 模式锁定，NULL=未发送过）
|
||||
try:
|
||||
cursor.execute("SELECT mode FROM sessions LIMIT 1")
|
||||
except Exception:
|
||||
cursor.execute("ALTER TABLE sessions ADD COLUMN mode TEXT")
|
||||
|
||||
cursor.execute("SELECT current_leaf_msg_id FROM sessions LIMIT 1")
|
||||
except sqlite3.OperationalError:
|
||||
print("[DB]: 正在升级 sessions 表 (添加 current_leaf_msg_id)...")
|
||||
cursor.execute("ALTER TABLE sessions ADD COLUMN current_leaf_msg_id TEXT")
|
||||
upgraded = True
|
||||
|
||||
try:
|
||||
cursor.execute("SELECT attachment_metadata, parent_id FROM messages LIMIT 1")
|
||||
except sqlite3.OperationalError:
|
||||
print("[DB]: 正在升级 messages 表 (添加 attachment_metadata, parent_id)...")
|
||||
try: cursor.execute("ALTER TABLE messages ADD COLUMN attachment_metadata TEXT")
|
||||
except: pass
|
||||
try: cursor.execute("ALTER TABLE messages ADD COLUMN parent_id TEXT")
|
||||
except: pass
|
||||
upgraded = True
|
||||
|
||||
# 🌟 messages.timeline 列（agent 时间线 JSON: 思考/文本/工具 按事件顺序）
|
||||
try:
|
||||
cursor.execute("SELECT timeline FROM messages LIMIT 1")
|
||||
except sqlite3.OperationalError:
|
||||
print("[DB]: 正在升级 messages 表 (添加 timeline)...")
|
||||
try: cursor.execute("ALTER TABLE messages ADD COLUMN timeline TEXT")
|
||||
except: pass
|
||||
upgraded = True
|
||||
|
||||
# 🆕 P1: messages.usage 列（assistant 回复的精确 usage JSON，
|
||||
# 供显示/压缩估算做 usage 锚定，对照 pi 内存态 usage 回放）
|
||||
try:
|
||||
cursor.execute("SELECT usage FROM messages LIMIT 1")
|
||||
except sqlite3.OperationalError:
|
||||
print("[DB]: 正在升级 messages 表 (添加 usage)...")
|
||||
try: cursor.execute("ALTER TABLE messages ADD COLUMN usage TEXT")
|
||||
except: pass
|
||||
upgraded = True
|
||||
|
||||
# 🆕 失败轮次持久化（对照 pi: message_end 无条件入库）：
|
||||
# messages.stop_reason / error_message —— 区分「正常行 / 出错行」，
|
||||
# 供 UI 展示与 build_api_context 回放决策（NULL = 正常）
|
||||
#
|
||||
# ⚠️⚠️ 绝不能置 upgraded=True：该标志会触发下方的「旧数据链表化重构」，
|
||||
# 把用户的**树状分支拍平成线性链**（数据破坏）！
|
||||
# 纯追加列对本迁移自身而言是安全的，与旧库结构修复无关。
|
||||
try:
|
||||
cursor.execute("SELECT stop_reason, error_message FROM messages LIMIT 1")
|
||||
except sqlite3.OperationalError:
|
||||
print("[DB]: 正在升级 messages 表 (添加 stop_reason, error_message)...")
|
||||
try: cursor.execute("ALTER TABLE messages ADD COLUMN stop_reason TEXT")
|
||||
except: pass
|
||||
try: cursor.execute("ALTER TABLE messages ADD COLUMN error_message TEXT")
|
||||
except: pass
|
||||
|
||||
# 如果刚才执行了树状结构升级，立即对旧数据进行“时间线串联”修复
|
||||
if upgraded:
|
||||
print("[DB]: 🚀 正在执行旧数据链表化重构...")
|
||||
cursor.execute("SELECT id FROM sessions")
|
||||
sessions = cursor.fetchall()
|
||||
for s in sessions:
|
||||
sid = s['id']
|
||||
cursor.execute("SELECT id FROM messages WHERE session_id = ? ORDER BY created_at ASC", (sid,))
|
||||
msgs = cursor.fetchall()
|
||||
if not msgs: continue
|
||||
|
||||
# 遍历消息，将后一条的 parent_id 指向上一条
|
||||
prev_id = None
|
||||
for m in msgs:
|
||||
mid = m['id']
|
||||
if prev_id:
|
||||
cursor.execute("UPDATE messages SET parent_id = ? WHERE id = ?", (prev_id, mid))
|
||||
prev_id = mid
|
||||
|
||||
# 最后一个 msg_id 就是这棵树的末端叶子节点
|
||||
cursor.execute("UPDATE sessions SET current_leaf_msg_id = ? WHERE id = ?", (prev_id, sid))
|
||||
print("[DB]: ✅ 数据结构完美升级为链表树！")
|
||||
|
||||
|
||||
def _seed_default_chat(self, cursor):
|
||||
session_id = f"sess_{uuid.uuid4().hex[:12]}"
|
||||
now = int(time.time())
|
||||
|
||||
cursor.execute(
|
||||
"INSERT INTO sessions (id, title, created_at, updated_at, has_messages) VALUES (?, ?, ?, ?, ?)",
|
||||
(session_id, "让我们从这里开始吧", now, now, 1)
|
||||
)
|
||||
|
||||
sys_id = f"msg_sys_init"
|
||||
cursor.execute("""
|
||||
INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, parent_id)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (sys_id, session_id, "system", "你是一个优秀的助手！", "", 0, now, None))
|
||||
|
||||
default_messages = [
|
||||
{"role": "user", "content": "你好呀，你是谁？"},
|
||||
{"role": "assistant", "content": "嗨！我是你的 AI 助手 ✨ ..."},
|
||||
{"role": "user", "content": "那你到底能帮我做什么？"},
|
||||
{"role": "assistant", "content": "简单来说，能打字问的我都聊..."}
|
||||
]
|
||||
|
||||
prev_id = sys_id
|
||||
for msg in default_messages:
|
||||
msg_id = f"msg_{uuid.uuid4().hex[:16]}"
|
||||
now += 1
|
||||
cursor.execute("""
|
||||
INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, parent_id)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (msg_id, session_id, msg["role"], msg["content"], "", 0, now, prev_id))
|
||||
prev_id = msg_id
|
||||
|
||||
# 设置默认会话的叶子节点
|
||||
cursor.execute("UPDATE sessions SET current_leaf_msg_id = ? WHERE id = ?", (prev_id, session_id))
|
||||
|
||||
# ==================== 会话 (Session) 操作 ====================
|
||||
def get_session_mode(self, session_id: str) -> Optional[str]:
|
||||
"""读取会话锁定的模式（chat/worker），未发送过返回 None"""
|
||||
with self.get_connection() as conn:
|
||||
row = conn.execute("SELECT mode FROM sessions WHERE id = ?",
|
||||
(session_id,)).fetchone()
|
||||
return row["mode"] if row else None
|
||||
|
||||
def set_session_mode(self, session_id: str, mode: str):
|
||||
"""锁定会话模式（首条消息发送时调用，之后不可变）"""
|
||||
with self.get_connection() as conn:
|
||||
conn.execute("UPDATE sessions SET mode = ? WHERE id = ?",
|
||||
(mode, session_id))
|
||||
|
||||
def get_all_sessions(self) -> List[Dict]:
|
||||
with self.get_connection() as conn:
|
||||
return [dict(row) for row in conn.execute("SELECT * FROM sessions ORDER BY is_starred DESC, sort_order ASC, updated_at DESC").fetchall()]
|
||||
|
||||
def create_session(self, title: str = "新对话") -> Dict:
|
||||
session_id = f"sess_{uuid.uuid4().hex[:12]}"
|
||||
now = int(time.time())
|
||||
sys_msg_id = f"msg_{uuid.uuid4().hex[:16]}"
|
||||
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
min_order = self.get_min_sort_order()
|
||||
|
||||
# 创建 session，直接将 system 消息设为初始叶子节点
|
||||
cursor.execute(
|
||||
"INSERT INTO sessions (id, title, created_at, updated_at, has_messages, sort_order, current_leaf_msg_id) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
(session_id, title, now, now, 0, min_order, sys_msg_id)
|
||||
)
|
||||
|
||||
cursor.execute("""
|
||||
INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, parent_id)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (sys_msg_id, session_id, "system", "你是一个优秀的助手！", "", 0, now, None))
|
||||
conn.commit()
|
||||
|
||||
return dict(cursor.execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone())
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 🆕 会话复制：深度克隆（全部分支 + 压缩标记 + 附件文件）
|
||||
# ------------------------------------------------------------------
|
||||
def _make_copy_title(self, base: str) -> str:
|
||||
"""生成不重名的副本标题：X → X (副本) → X (副本 2) → …"""
|
||||
root = re.sub(r"\s*\(副本(?:\s*\d+)?\)\s*$", "", base or "").strip() or "新对话"
|
||||
with self.get_connection() as conn:
|
||||
existing = {r[0] for r in conn.execute("SELECT title FROM sessions").fetchall()}
|
||||
cand = f"{root} (副本)"
|
||||
n = 2
|
||||
while cand in existing and n < 1000:
|
||||
cand = f"{root} (副本 {n})"
|
||||
n += 1
|
||||
return cand
|
||||
|
||||
def _dup_media_file(self, old_path: str, prefix: str):
|
||||
"""把一个项目内媒体文件复制成新名字。
|
||||
返回 (新绝对路径, 新项目相对路径)；文件不存在或异常返回 None。"""
|
||||
if not old_path:
|
||||
return None
|
||||
old_abs = old_path if os.path.isabs(old_path) else os.path.join(self.files_root, old_path)
|
||||
if not os.path.isfile(old_abs):
|
||||
return None
|
||||
ext = os.path.splitext(old_abs)[1] or ".bin"
|
||||
new_dir = os.path.join(self.files_root, "data", "attachments")
|
||||
os.makedirs(new_dir, exist_ok=True)
|
||||
new_abs = ""
|
||||
for _ in range(5):
|
||||
new_abs = os.path.join(new_dir, f"{prefix}_{uuid.uuid4().hex[:8]}{ext}")
|
||||
if not os.path.exists(new_abs):
|
||||
break
|
||||
shutil.copy2(old_abs, new_abs)
|
||||
new_rel = os.path.relpath(new_abs, self.files_root).replace("\\", "/")
|
||||
return new_abs, new_rel
|
||||
|
||||
def _copy_attachment_files(self, meta_json: Optional[str]) -> Optional[str]:
|
||||
"""附件深拷贝：image/pdf 的磁盘文件复制改名并改写 local_path；
|
||||
text 类型正文内联在 JSON 里，无需处理。
|
||||
任何异常/文件缺失都原样返回，绝不让复制整体失败。"""
|
||||
if not meta_json:
|
||||
return meta_json
|
||||
try:
|
||||
meta = json.loads(meta_json)
|
||||
except Exception:
|
||||
return meta_json
|
||||
if not isinstance(meta, dict) or not meta.get("attachments"):
|
||||
return meta_json
|
||||
changed = False
|
||||
for att in meta.get("attachments") or []:
|
||||
if not isinstance(att, dict):
|
||||
continue
|
||||
try:
|
||||
if att.get("type") in ("image", "pdf") and att.get("local_path"):
|
||||
got = self._dup_media_file(
|
||||
att["local_path"], "img" if att["type"] == "image" else "pdf")
|
||||
if got:
|
||||
att["local_path"] = got[1]
|
||||
if "abs_path" in att:
|
||||
att["abs_path"] = got[0]
|
||||
changed = True
|
||||
if att.get("type") == "pdf":
|
||||
for im in att.get("images") or []:
|
||||
if not isinstance(im, dict):
|
||||
continue
|
||||
got = self._dup_media_file(
|
||||
im.get("abs_path") or im.get("local_path"), "pdfimg")
|
||||
if got:
|
||||
im["abs_path"] = got[0]
|
||||
im["local_path"] = got[1]
|
||||
changed = True
|
||||
except Exception as e:
|
||||
print(f"[DB] ⚠️ 附件深拷贝失败（保留原路径）: {e}", flush=True)
|
||||
return json.dumps(meta, ensure_ascii=False) if changed else meta_json
|
||||
|
||||
def copy_session(self, session_id: str, new_title: Optional[str] = None,
|
||||
copy_attachments: bool = True) -> Optional[Dict]:
|
||||
"""📋 深度复制一个会话。
|
||||
|
||||
- messages 全部重新生成 ID，parent_id / current_leaf_msg_id 全量重映射
|
||||
→ 分支、压缩标记（role='compaction'）都原样保留
|
||||
- image/pdf 附件文件物理复制成新文件 → 副本自包含，删任意一方不影响另一方
|
||||
- 单事务写入；源会话零改动
|
||||
返回新会话 dict；源不存在返回 None。
|
||||
"""
|
||||
now = int(time.time())
|
||||
n_att = 0
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
src_row = cursor.execute(
|
||||
"SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
|
||||
if not src_row:
|
||||
return None
|
||||
src = dict(src_row)
|
||||
rows = [dict(r) for r in cursor.execute(
|
||||
"SELECT * FROM messages WHERE session_id = ? ORDER BY created_at ASC, rowid ASC",
|
||||
(session_id,)).fetchall()]
|
||||
|
||||
# ① 新 ID 映射（保持 comp_ / msg_ 前缀约定）
|
||||
idmap = {}
|
||||
for m in rows:
|
||||
pre = "comp_" if m.get("role") == "compaction" else "msg_"
|
||||
idmap[m["id"]] = f"{pre}{uuid.uuid4().hex[:16]}"
|
||||
|
||||
# ② 标题（重名自动递增）
|
||||
title = new_title or self._make_copy_title(src.get("title") or "新对话")
|
||||
|
||||
# ③ 消息 created_at 保相对间隔（同层兄弟排序不变）
|
||||
t0 = min((m.get("created_at") or 0) for m in rows) if rows else now
|
||||
|
||||
new_sid = f"sess_{uuid.uuid4().hex[:12]}"
|
||||
new_leaf = idmap.get(src.get("current_leaf_msg_id"))
|
||||
if rows and not new_leaf:
|
||||
# 兜底：源叶子不在链上（数据异常）→ 取副本里时间最新的一条
|
||||
last = max(rows, key=lambda m: (m.get("created_at") or 0))
|
||||
new_leaf = idmap.get(last["id"])
|
||||
print(f"[DB] ⚠️ copy_session 源叶子异常，回退 leaf={new_leaf}", flush=True)
|
||||
|
||||
# ④ 新会话（置列表顶部、不带星标、模式跟随源）
|
||||
cursor.execute(
|
||||
"INSERT INTO sessions (id, title, created_at, updated_at, has_messages, "
|
||||
"sort_order, is_starred, current_leaf_msg_id, mode) VALUES (?,?,?,?,?,?,?,?,?)",
|
||||
(new_sid, title, now, now, src.get("has_messages") or 0,
|
||||
self.get_min_sort_order(), 0, new_leaf, src.get("mode")))
|
||||
|
||||
# ⑤ 逐条复制消息（parent 重映射 + 附件深拷贝）
|
||||
for m in rows:
|
||||
meta = m.get("attachment_metadata")
|
||||
if copy_attachments and meta:
|
||||
new_meta = self._copy_attachment_files(meta)
|
||||
if new_meta != meta:
|
||||
n_att += 1
|
||||
meta = new_meta
|
||||
cursor.execute(
|
||||
"INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, "
|
||||
"created_at, attachment_metadata, parent_id, timeline, usage, stop_reason, error_message) "
|
||||
"VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)",
|
||||
(idmap[m["id"]], new_sid, m.get("role"), m.get("content"),
|
||||
m.get("reasoning"), m.get("is_ignored"),
|
||||
now + ((m.get("created_at") or 0) - t0),
|
||||
meta, idmap.get(m.get("parent_id")),
|
||||
m.get("timeline"), m.get("usage"),
|
||||
m.get("stop_reason"), m.get("error_message")))
|
||||
|
||||
conn.commit()
|
||||
out = dict(cursor.execute(
|
||||
"SELECT * FROM sessions WHERE id = ?", (new_sid,)).fetchone())
|
||||
|
||||
try:
|
||||
print(f"[DB] copy_session {session_id[:8]} → {new_sid[:8]} "
|
||||
f"消息={len(rows)} 附件深拷贝={n_att} 标题={title}", flush=True)
|
||||
except Exception:
|
||||
pass
|
||||
return out
|
||||
|
||||
def update_session_title(self, session_id: str, new_title: str):
|
||||
with self.get_connection() as conn:
|
||||
conn.execute("UPDATE sessions SET title = ?, updated_at = ? WHERE id = ?", (new_title, int(time.time()), session_id))
|
||||
conn.commit()
|
||||
|
||||
def delete_session(self, session_id: str):
|
||||
with self.get_connection() as conn:
|
||||
conn.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
|
||||
conn.commit()
|
||||
|
||||
# 下方其它基本Session功能保持不变...
|
||||
def mark_session_has_messages(self, session_id: str):
|
||||
with self.get_connection() as conn:
|
||||
conn.execute("UPDATE sessions SET has_messages = 1 WHERE id = ?", (session_id,))
|
||||
conn.commit()
|
||||
|
||||
def check_session_needs_title(self, session_id: str) -> bool:
|
||||
with self.get_connection() as conn:
|
||||
row = conn.execute("SELECT title, has_messages FROM sessions WHERE id = ?", (session_id,)).fetchone()
|
||||
if row: return row[0] == "新对话" and row[1] == 1
|
||||
return False
|
||||
|
||||
def get_min_sort_order(self) -> int:
|
||||
with self.get_connection() as conn:
|
||||
row = conn.execute("SELECT MIN(sort_order) FROM sessions").fetchone()
|
||||
return (row[0] or 0) - 1
|
||||
|
||||
def update_session_order(self, ordered_ids: list):
|
||||
with self.get_connection() as conn:
|
||||
for idx, sid in enumerate(ordered_ids):
|
||||
conn.execute("UPDATE sessions SET sort_order = ? WHERE id = ?", (idx, sid))
|
||||
conn.commit()
|
||||
|
||||
def update_session_star(self, session_id: str, is_starred: bool):
|
||||
with self.get_connection() as conn:
|
||||
conn.execute("UPDATE sessions SET is_starred = ?, updated_at = ? WHERE id = ?", (1 if is_starred else 0, int(time.time()), session_id))
|
||||
conn.commit()
|
||||
|
||||
def is_session_starred(self, session_id: str) -> bool:
|
||||
with self.get_connection() as conn:
|
||||
row = conn.execute("SELECT is_starred FROM sessions WHERE id = ?", (session_id,)).fetchone()
|
||||
return bool(row and row[0])
|
||||
|
||||
|
||||
# ==================== 🌟 核心：消息链表树操作 ====================
|
||||
|
||||
def get_session_leaf(self, session_id: str) -> Optional[str]:
|
||||
"""获取会话当前的叶子节点ID"""
|
||||
with self.get_connection() as conn:
|
||||
row = conn.execute("SELECT current_leaf_msg_id FROM sessions WHERE id = ?", (session_id,)).fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
def update_session_leaf(self, session_id: str, leaf_msg_id: str):
|
||||
"""切换时间线：手动更新当前会话的叶子节点"""
|
||||
with self.get_connection() as conn:
|
||||
conn.execute("UPDATE sessions SET current_leaf_msg_id = ?, updated_at = ? WHERE id = ?",
|
||||
(leaf_msg_id, int(time.time()), session_id))
|
||||
conn.commit()
|
||||
|
||||
def get_message_chain(self, session_id: str) -> List[Dict]:
|
||||
"""🚀 极客级递归拉取：顺藤摸瓜，只返回当前激活时间线上的消息！彻底断绝下游污染！"""
|
||||
leaf_id = self.get_session_leaf(session_id)
|
||||
if not leaf_id:
|
||||
return []
|
||||
|
||||
chain = []
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
current_id = leaf_id
|
||||
seen = set() # 🐛 环检测：数据异常（如自环标记）时快速退出，防主线程死循环卡死
|
||||
|
||||
# 使用 Python 循环向上回溯（比 SQL CTE 更好调试，性能对于本地几千条聊天来说在 1ms 内）
|
||||
while current_id and current_id not in seen:
|
||||
seen.add(current_id)
|
||||
cursor.execute("SELECT * FROM messages WHERE id = ?", (current_id,))
|
||||
msg = cursor.fetchone()
|
||||
if not msg:
|
||||
break
|
||||
chain.append(dict(msg))
|
||||
current_id = msg['parent_id']
|
||||
if current_id in seen:
|
||||
try:
|
||||
print(f"[DB] ⚠️ get_message_chain 检测到环（session={session_id}），已截断", flush=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 因为是向上回溯，拉出来的链条是反的，最后翻转一下恢复正序
|
||||
chain.reverse()
|
||||
return chain
|
||||
|
||||
def get_branch_info(self, parent_id: str) -> List[Dict]:
|
||||
"""获取某一父节点下的所有子分支消息 (第二阶段用于UI渲染 '1/3')"""
|
||||
if not parent_id: return []
|
||||
with self.get_connection() as conn:
|
||||
return [dict(row) for row in conn.execute(
|
||||
"SELECT * FROM messages WHERE parent_id = ? ORDER BY created_at ASC", (parent_id,)
|
||||
).fetchall()]
|
||||
|
||||
def add_message(self, session_id: str, role: str, content: str, parent_id: str,
|
||||
reasoning: str = "", is_ignored: bool = False,
|
||||
msg_id: Optional[str] = None, attachment_metadata: Optional[str] = None,
|
||||
timeline: Optional[str] = None,
|
||||
usage: Optional[str] = None,
|
||||
stop_reason: Optional[str] = None,
|
||||
error_message: Optional[str] = None) -> Dict:
|
||||
"""添加新消息，并自动将该消息设为当前会话的最新叶子节点
|
||||
🆕 P1: usage —— assistant 回复的精确 usage JSON（如 '{"input":..,"output":..}'）
|
||||
🆕 失败轮次: stop_reason/error_message —— 'error' 行入库但不回退叶子
|
||||
（对照 pi：错误也持久化，回放时由 build_api_context 决定取舍）"""
|
||||
if not msg_id: msg_id = f"msg_{uuid.uuid4().hex[:16]}"
|
||||
now = int(time.time())
|
||||
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
# 1. 插入消息
|
||||
cursor.execute("""
|
||||
INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, attachment_metadata, parent_id, timeline, usage, stop_reason, error_message)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (msg_id, session_id, role, content, reasoning, 1 if is_ignored else 0, now, attachment_metadata, parent_id, timeline, usage, stop_reason, error_message))
|
||||
|
||||
# 2. 自动更新 session 的叶子节点（时间线前推）
|
||||
cursor.execute("UPDATE sessions SET current_leaf_msg_id = ?, updated_at = ? WHERE id = ?",
|
||||
(msg_id, now, session_id))
|
||||
conn.commit()
|
||||
try:
|
||||
print(f"[DB] add_message role={role} 内容={len(content or '')}c "
|
||||
f"思考={len(reasoning or '')}c 时间线={'有' if timeline else '无'} "
|
||||
f"id={msg_id} session={session_id[:8]}", flush=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return dict(cursor.execute("SELECT * FROM messages WHERE id = ?", (msg_id,)).fetchone())
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 🆕 压缩持久化：链上标记点（子叶子回溯到标记即停）
|
||||
# 插入后链形：…→ cut_before → [MARK role="compaction"] → first_retained → …→叶子
|
||||
# 历史行全部保留（前端切会话渲染照常）；API 上下文构建在标记处截断。
|
||||
# 关键：不动 sessions.current_leaf_msg_id（绝不能用 add_message）；
|
||||
# INSERT + UPDATE 同一事务，不留断链窗口。
|
||||
# ------------------------------------------------------------------
|
||||
def insert_compaction_mark(self, session_id: str, summary: str,
|
||||
cut_before_id: str, first_retained_id: str,
|
||||
meta_json: Optional[str] = None) -> Optional[str]:
|
||||
"""在链上插入压缩标记行并把保留首条的 parent_id 改指到标记。返回 mark_id。"""
|
||||
if not cut_before_id or not first_retained_id:
|
||||
return None
|
||||
if cut_before_id == first_retained_id:
|
||||
# 🐛 防自环：同一行不能既做切点前又做保留首条（timeline 回放同 id 场景）
|
||||
try:
|
||||
print(f"[DB] ⚠️ insert_compaction_mark 拒绝自环 cut==retained={cut_before_id}", flush=True)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
mark_id = f"comp_{uuid.uuid4().hex[:16]}"
|
||||
now = int(time.time())
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""INSERT INTO messages
|
||||
(id, session_id, role, content, reasoning, is_ignored,
|
||||
created_at, attachment_metadata, parent_id, timeline, usage)
|
||||
VALUES (?, ?, 'compaction', ?, '', 1, ?, ?, ?, NULL, NULL)""",
|
||||
(mark_id, session_id, summary, now, meta_json, cut_before_id))
|
||||
cursor.execute("UPDATE messages SET parent_id = ? WHERE id = ?",
|
||||
(mark_id, first_retained_id))
|
||||
conn.commit()
|
||||
try:
|
||||
print(f"[DB] insert_compaction_mark session={session_id[:8]} "
|
||||
f"cut_before={cut_before_id} first_retained={first_retained_id} "
|
||||
f"mark={mark_id} summary={len(summary or '')}c", flush=True)
|
||||
except Exception:
|
||||
pass
|
||||
return mark_id
|
||||
# === 在 db_manager.py 中添加这个方法 ===
|
||||
def get_branch_leaf(self, msg_id: str) -> str:
|
||||
"""寻找一条时间线的最末端叶子节点"""
|
||||
current_id = msg_id
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
while True:
|
||||
# 寻找把当前节点作为父节点的子节点，按时间倒序取最新的一条（意味着它顺着最近被聊过的那条线往下走）
|
||||
cursor.execute("SELECT id FROM messages WHERE parent_id = ? ORDER BY created_at DESC LIMIT 1", (current_id,))
|
||||
child = cursor.fetchone()
|
||||
if child:
|
||||
current_id = child[0]
|
||||
else:
|
||||
break # 没有子节点了，它自己就是叶子！
|
||||
return current_id
|
||||
def delete_message_branch(self, session_id: str, msg_id: str):
|
||||
"""🚀 精准剪枝：删AI只删当前分支，删User连根拔起，并自动平滑回退时间线"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 1. 查出要删除的节点的父亲
|
||||
cursor.execute("SELECT parent_id FROM messages WHERE id = ?", (msg_id,))
|
||||
row = cursor.fetchone()
|
||||
if not row: return
|
||||
safe_parent_id = row['parent_id']
|
||||
|
||||
# 2. 目标就是传进来的 msg_id 本身 (不再强制上移到 parent)
|
||||
target_id = msg_id
|
||||
|
||||
# 3. 拉取全会话，构建亲属关系树
|
||||
cursor.execute("SELECT id, parent_id FROM messages WHERE session_id = ?", (session_id,))
|
||||
all_msgs = cursor.fetchall()
|
||||
|
||||
children_map = {}
|
||||
for m in all_msgs:
|
||||
pid = m['parent_id']
|
||||
if pid not in children_map:
|
||||
children_map[pid] = []
|
||||
children_map[pid].append(m['id'])
|
||||
|
||||
# 4. 广度优先搜索 (BFS)：找出目标消息及其所有子孙
|
||||
to_delete = set([target_id])
|
||||
queue = [target_id]
|
||||
while queue:
|
||||
curr = queue.pop(0)
|
||||
if curr in children_map:
|
||||
for child in children_map[curr]:
|
||||
if child not in to_delete:
|
||||
to_delete.add(child)
|
||||
queue.append(child)
|
||||
|
||||
# 5. 判断当前时间线的“叶子节点”是否在被波及的名单里
|
||||
cursor.execute("SELECT current_leaf_msg_id FROM sessions WHERE id = ?", (session_id,))
|
||||
leaf_row = cursor.fetchone()
|
||||
leaf_needs_update = leaf_row and leaf_row['current_leaf_msg_id'] in to_delete
|
||||
|
||||
# 6. 🌟 执行物理删除前，收集将被删除的附件元数据
|
||||
deleted_metadata = []
|
||||
for d_id in to_delete:
|
||||
# 先查出它的 metadata
|
||||
cursor.execute("SELECT attachment_metadata FROM messages WHERE id = ?", (d_id,))
|
||||
row = cursor.fetchone()
|
||||
if row and row['attachment_metadata']:
|
||||
deleted_metadata.append(row['attachment_metadata'])
|
||||
|
||||
# 然后再执行物理删除
|
||||
cursor.execute("DELETE FROM messages WHERE id = ?", (d_id,))
|
||||
|
||||
|
||||
# 7. 🌟 核心：如果时间线断了，自动寻找平滑降落点
|
||||
sibling_row = None
|
||||
if leaf_needs_update:
|
||||
# 尝试寻找被删节点的最新“兄弟姐妹” (例如删了分支2，寻找分支1)
|
||||
cursor.execute("SELECT id FROM messages WHERE parent_id = ? ORDER BY created_at DESC LIMIT 1", (safe_parent_id,))
|
||||
sibling_row = cursor.fetchone()
|
||||
|
||||
# 如果有兄弟，降落到兄弟；如果没兄弟(只有1次回答)，退回原点(提问)
|
||||
new_leaf = sibling_row['id'] if sibling_row else safe_parent_id
|
||||
cursor.execute("UPDATE sessions SET current_leaf_msg_id = ? WHERE id = ?", (new_leaf, session_id))
|
||||
|
||||
conn.commit()
|
||||
|
||||
# 8. 如果降落到了兄弟分支，兄弟可能还有下文，需再次对齐真实叶子节点
|
||||
if leaf_needs_update and sibling_row:
|
||||
real_leaf = self.get_branch_leaf(new_leaf)
|
||||
self.update_session_leaf(session_id, real_leaf)
|
||||
|
||||
return deleted_metadata # 🌟 返回被删除的元数据，交给 MainWindow 去粉碎文件
|
||||
@@ -0,0 +1,78 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""统一调试日志 + 调试器窗口控制协议（纯文件通信，与主窗口 UI 零耦合）
|
||||
|
||||
日志文件: data/debug_session.log （环境变量 HAOCODE_DEBUG_LOG 可覆盖）
|
||||
控制文件: data/debug_window.cmd （环境变量 HAOCODE_DEBUG_CMD 可覆盖）
|
||||
|
||||
三方写入协议:
|
||||
[USER] 用户在调试窗口输入框手输的观察情况
|
||||
[AGENT] 代理（外部脚本/命令行）注入的指令与备注
|
||||
[APP] 应用自身事件（上下文标签变化/usage/发送/完成/报错/压缩）
|
||||
[SYS] 调试窗口自身的开关事件
|
||||
|
||||
行格式: [YYYY-MM-DD HH:MM:SS.mmm] [TAG] 内容
|
||||
|
||||
控制协议: 代理往 debug_window.cmd 写入 "show" 或 "hide"（一行），
|
||||
app 侧 2s QTimer 轮询并消费（读完即删）。
|
||||
"""
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
|
||||
_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
_DATA_DIR = os.path.join(_ROOT, "data")
|
||||
|
||||
DEBUG_LOG_PATH = (os.environ.get("HAOCODE_DEBUG_LOG")
|
||||
or os.path.join(_DATA_DIR, "debug_session.log"))
|
||||
DEBUG_CMD_PATH = (os.environ.get("HAOCODE_DEBUG_CMD")
|
||||
or os.path.join(_DATA_DIR, "debug_window.cmd"))
|
||||
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
def _stamp() -> str:
|
||||
t = time.time()
|
||||
return (f"{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(t))}"
|
||||
f".{int(t * 1000) % 1000:03d}")
|
||||
|
||||
|
||||
def debug_log(msg: str, tag: str = "APP") -> None:
|
||||
"""线程安全追加一条日志（worker 线程亦可调用）。
|
||||
静默吞掉一切异常——本模块绝不影响主流程。"""
|
||||
try:
|
||||
with _lock:
|
||||
with open(DEBUG_LOG_PATH, "a", encoding="utf-8") as f:
|
||||
f.write(f"[{_stamp()}] [{tag}] {msg}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def poll_debug_cmd():
|
||||
"""消费控制文件，返回 'show' / 'hide' / None。由主线程 QTimer 周期调用。"""
|
||||
try:
|
||||
if os.path.exists(DEBUG_CMD_PATH):
|
||||
with open(DEBUG_CMD_PATH, "r", encoding="utf-8") as f:
|
||||
action = (f.read() or "").strip().lower()
|
||||
try:
|
||||
os.remove(DEBUG_CMD_PATH)
|
||||
except Exception:
|
||||
pass
|
||||
if action in ("show", "hide"):
|
||||
return action
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def autostart_debug_window(cfg: dict) -> bool:
|
||||
"""调试窗口随程序启动：cfg["debug_window_autostart"] 为 true（缺省也是 true）
|
||||
时向控制文件写 "show"，主窗口事件循环启动后 2s 轮询即开窗。
|
||||
纯文件操作（不依赖 Qt），返回是否写入。"""
|
||||
try:
|
||||
if not bool(cfg.get("debug_window_autostart", True)):
|
||||
return False
|
||||
with open(DEBUG_CMD_PATH, "w", encoding="utf-8") as f:
|
||||
f.write("show\n")
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
@@ -0,0 +1,446 @@
|
||||
"""
|
||||
core/llm_engine.py
|
||||
==================
|
||||
Qt 胶水层 —— 把 core.agent（pi 1:1 核心）接到 PyQt6 信号系统。
|
||||
|
||||
🌟 对照 pi 架构：
|
||||
AgentWorker(QThread)
|
||||
内部持有 Agent + AgentRunner（对照 pi 的 Agent + AgentSession 会话层）
|
||||
run() = pi 的 _runAgentPrompt：prompt → 循环 post-agent-run 恢复 → settle
|
||||
事件桥：AgentEvent → pyqtSignal（跨线程 queued 投递到 UI 线程）
|
||||
|
||||
ChatWorker(QThread)
|
||||
chat 模式：普通聊天单次流式（无工具循环/无重试/无压缩）。
|
||||
|
||||
TitleWorker(QThread)
|
||||
轻量单次流式补全（标题生成），保留原 LLMWorker 的信号面。
|
||||
|
||||
信号面（与旧 LLMWorker 兼容 + 新增工具事件）：
|
||||
chunk_received(str) # text_delta
|
||||
reasoning_received(str) # thinking_delta
|
||||
error_occurred(str) # 终局错误（已重试/压缩恢复仍失败）
|
||||
tool_execution_started(str, str, str) # (call_id, tool_name, args_json)
|
||||
tool_execution_updated(str, str) # (call_id, 执行中增量输出)
|
||||
tool_execution_finished(str, str, bool, str) # (call_id, tool_name, ok, 结果摘要)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openai import OpenAI
|
||||
from PyQt6.QtCore import QThread, pyqtSignal
|
||||
|
||||
from core.agent import (Agent, AgentConfig, AgentEvent, AgentMessage, AgentRunner,
|
||||
ModelConfig, RetryConfig, calculate_context_tokens,
|
||||
default_tools, from_openai_messages,
|
||||
openai_stream)
|
||||
from core.agent.stream_fn import _pick_reasoning, _pick_usage
|
||||
|
||||
CONFIG_PATH = os.path.join(os.path.dirname(__file__), "..", "data", "config.json")
|
||||
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
# 🌟 Agent 系统提示词文件（项目根目录，core 前面那个 .md）
|
||||
SYSTEM_PROMPT_FILE = os.path.join(PROJECT_ROOT, "SYSTEM_PROMPT.md")
|
||||
_FALLBACK_SYSTEM_PROMPT = (
|
||||
"你是 haocode 的本地智能体，可以读取/写入文件并执行 bash 命令"
|
||||
"（工作目录为项目根目录）。请用简体中文简洁地回答。"
|
||||
)
|
||||
|
||||
|
||||
def load_system_prompt() -> str:
|
||||
"""读取 SYSTEM_PROMPT.md；缺失时用兜底短提示词。"""
|
||||
try:
|
||||
with open(SYSTEM_PROMPT_FILE, "r", encoding="utf-8") as f:
|
||||
text = f.read().strip()
|
||||
if text:
|
||||
return text
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
return _FALLBACK_SYSTEM_PROMPT
|
||||
|
||||
|
||||
def _load_config() -> dict:
|
||||
if os.path.exists(CONFIG_PATH):
|
||||
try:
|
||||
with open(CONFIG_PATH, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
print(f"[llm_engine] 读取配置失败: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def _provider_info(config: dict, provider_name: str) -> dict:
|
||||
return config.get("providers", {}).get(provider_name, {}) or {}
|
||||
|
||||
|
||||
class AgentWorker(QThread):
|
||||
"""
|
||||
一次「发送」= 一个 AgentWorker（对照 pi：一个 AgentSession 实例）。
|
||||
多轮工具循环在一次 run() 内完成；UI 只管收信号 + 最终入库。
|
||||
"""
|
||||
|
||||
# ---- 与旧 LLMWorker 兼容的信号 ----
|
||||
chunk_received = pyqtSignal(str)
|
||||
reasoning_received = pyqtSignal(str)
|
||||
error_occurred = pyqtSignal(str)
|
||||
# ---- 新增：工具执行事件（pi tool_execution_* 事件）----
|
||||
tool_execution_started = pyqtSignal(str, str, str) # (call_id, name, args)
|
||||
tool_execution_updated = pyqtSignal(str, str) # (call_id, 增量输出)
|
||||
tool_execution_timed = pyqtSignal(str, int, int) # (call_id, 已耗秒, 超时秒) 🆕 bash 读秒
|
||||
tool_execution_finished = pyqtSignal(str, str, bool, str) # (call_id, name, ok, 结果)
|
||||
# 上下文压缩发生（UI 可提示「已自动压缩上下文」）
|
||||
context_compacted = pyqtSignal(dict)
|
||||
# 🆕 压缩开始（摘要 LLM 阻塞调用前）→ 前端显示「执行中」动态气泡
|
||||
# payload: {summary, before, after, duration_ms, path}
|
||||
compaction_started = pyqtSignal(str) # path: pre_prompt/overflow_compact/length_compact
|
||||
# 🆕 P1: 收到本轮精确 usage（UI 上下文标签做 usage 锚定，含 system+tools）
|
||||
usage_updated = pyqtSignal(dict)
|
||||
# 🆕 M3: 重试调度 / 重试结果（对照 pi onRetryScheduled/onRetryFinished）
|
||||
retry_scheduled = pyqtSignal(int, int, float, str) # (attempt, max_attempts, delay_ms, reason)
|
||||
retry_finished = pyqtSignal(bool)
|
||||
|
||||
def __init__(self, provider_name: str, model_name: str,
|
||||
openai_messages: List[Dict[str, Any]],
|
||||
tool_context: Optional[Dict[str, Any]] = None,
|
||||
enable_tools: bool = True,
|
||||
temperature: Optional[float] = None):
|
||||
"""
|
||||
openai_messages: build_api_context 的产物（OpenAI 格式）。
|
||||
最后一条 user 消息 = 本次提问；其余 = 历史上下文。
|
||||
"""
|
||||
super().__init__()
|
||||
self.provider_name = provider_name
|
||||
self.model_name = model_name
|
||||
self.openai_messages = openai_messages
|
||||
self.tool_context = tool_context or {"cwd": PROJECT_ROOT}
|
||||
self.enable_tools = enable_tools
|
||||
self._temperature_override = temperature
|
||||
|
||||
self.config = _load_config()
|
||||
self._aborted = False
|
||||
self._agent: Optional[Agent] = None
|
||||
self._runner: Optional[AgentRunner] = None
|
||||
self._compactions_before = 0
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 中断（对照 pi abort：当前流立刻停止，已产生的内容保留）
|
||||
# ------------------------------------------------------------------
|
||||
def abort(self):
|
||||
self._aborted = True
|
||||
if self._agent is not None:
|
||||
self._agent.abort()
|
||||
|
||||
# 兼容旧调用名
|
||||
def cancel(self):
|
||||
self.abort()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 模型配置（对照 pi models.json 条目 → ModelConfig）
|
||||
# ------------------------------------------------------------------
|
||||
def _build_model_config(self) -> ModelConfig:
|
||||
p = _provider_info(self.config, self.provider_name)
|
||||
context_window = int(p.get("model_contexts", {}).get(self.model_name, 128000))
|
||||
max_tokens = int(p.get("model_max_tokens", {}).get(self.model_name, 8192))
|
||||
temperature = (self._temperature_override
|
||||
if self._temperature_override is not None
|
||||
else float(self.config.get("temperature", 0.7)))
|
||||
return ModelConfig(
|
||||
provider=self.provider_name,
|
||||
name=self.model_name,
|
||||
context_window=context_window,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
api_key=p.get("api_key", ""),
|
||||
base_url=p.get("base_url", "https://api.openai.com/v1"),
|
||||
# 🆕 thinking reasoning_content 回传默认开启（ModelConfig.pass_reasoning=True）
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 压缩用的非流式 LLM 调用（对照 pi completeSimpleWithRetries）
|
||||
# 新签名 1:1 对照 pi：(prompt_text, system_prompt, max_tokens) -> str
|
||||
# ------------------------------------------------------------------
|
||||
def _make_summarize_fn(self, model: ModelConfig):
|
||||
def summarize(prompt_text: str, system_prompt: str, max_tokens: int) -> str:
|
||||
client = OpenAI(api_key=model.api_key, base_url=model.base_url,
|
||||
timeout=120.0, max_retries=0)
|
||||
resp = client.chat.completions.create(
|
||||
model=model.name,
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": prompt_text},
|
||||
],
|
||||
stream=False,
|
||||
max_tokens=max_tokens,
|
||||
temperature=0.3,
|
||||
)
|
||||
return (resp.choices[0].message.content or "").strip()
|
||||
|
||||
return summarize
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 事件桥：AgentEvent → Qt 信号（在工作线程内 emit，queued 投递）
|
||||
# ------------------------------------------------------------------
|
||||
def _bridge(self, e: AgentEvent):
|
||||
if e.type == "message_end" and e.message is not None:
|
||||
# 🆕 P1: 本轮刚结束 → 它的 usage 是「上一张照片」（精确覆盖
|
||||
# system+tools+历史+本轮输出），供 UI 显示锚定
|
||||
usage = e.message.usage or {}
|
||||
if calculate_context_tokens(usage) > 0:
|
||||
self.usage_updated.emit(dict(usage))
|
||||
if e.type == "message_update" and e.assistant_message_event is not None:
|
||||
ev = e.assistant_message_event
|
||||
# 🌟 字段是 ev.text（不是 ev.data）；_emit 会吞掉订阅者异常，
|
||||
# 写错字段 = 静默无流式输出，必须靠 test_agent_worker_stream_fn_wiring 守住
|
||||
if ev.type == "text_delta" and ev.text:
|
||||
self.chunk_received.emit(ev.text)
|
||||
elif ev.type == "thinking_delta" and ev.text:
|
||||
self.reasoning_received.emit(ev.text)
|
||||
return
|
||||
if e.type == "tool_execution_start" and e.tool_call is not None:
|
||||
try:
|
||||
args = json.dumps(e.tool_call.arguments, ensure_ascii=False) \
|
||||
if e.tool_call.arguments else (e.tool_call.raw_arguments or "")
|
||||
except Exception:
|
||||
args = str(e.tool_call.arguments)
|
||||
self.tool_execution_started.emit(e.tool_call.id, e.tool_call.name,
|
||||
str(args)[:2000])
|
||||
elif e.type == "tool_execution_update" and e.arg:
|
||||
self.tool_execution_updated.emit(e.tool_call.id, str(e.arg))
|
||||
elif e.type == "tool_execution_timer" and isinstance(
|
||||
getattr(e, "arg", None), tuple) and len(e.arg) == 2:
|
||||
# 🆕 bash 运行中每秒滴一次 → 前端气泡读秒
|
||||
self.tool_execution_timed.emit(e.tool_call.id, int(e.arg[0]),
|
||||
int(e.arg[1]))
|
||||
elif e.type == "tool_execution_end" and e.tool_call is not None:
|
||||
text = ""
|
||||
if e.result is not None:
|
||||
c = e.result.content
|
||||
text = c if isinstance(c, str) else \
|
||||
"".join(x.get("text", "") for x in c if isinstance(x, dict))
|
||||
# 🌟 修复: 原来 ok = is_error（反了）——成功的工具被显示成失败
|
||||
ok = not (e.result.is_error if e.result is not None else True)
|
||||
self.tool_execution_finished.emit(
|
||||
e.tool_call.id, e.tool_call.name, bool(ok), text[:20000]) # 长结果供 UI 展开
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 主入口 —— 对照 pi _runAgentPrompt
|
||||
# ------------------------------------------------------------------
|
||||
def run(self):
|
||||
provider = _provider_info(self.config, self.provider_name)
|
||||
if not provider:
|
||||
self.error_occurred.emit(f"未找到 provider 配置: {self.provider_name}")
|
||||
return
|
||||
|
||||
history = from_openai_messages(self.openai_messages)
|
||||
if not history or history[-1].role != "user":
|
||||
self.error_occurred.emit("上下文必须以 user 消息结尾")
|
||||
return
|
||||
last_user = history.pop() # 本次提问（其余为历史）
|
||||
|
||||
model = self._build_model_config()
|
||||
cfg = AgentConfig(
|
||||
model=model,
|
||||
# 🌟 系统提示词：SYSTEM_PROMPT.md（每次请求头部注入，不入历史）
|
||||
system_prompt=load_system_prompt() if self.enable_tools else "",
|
||||
tools=default_tools() if self.enable_tools else [],
|
||||
tool_context=self.tool_context,
|
||||
retry=RetryConfig(max_attempts=3, base_delay_ms=2000, factor=2.0),
|
||||
)
|
||||
agent = Agent(cfg)
|
||||
# 🌟 关键：注入流函数（对照 pi 的 agentLoopConfig.streamFn 注入点）。
|
||||
# 漏掉这一步 Agent 会抛 AgentError("未配置 stream_fn")。
|
||||
agent.set_stream_fn(openai_stream)
|
||||
agent.state.messages = list(history)
|
||||
agent.subscribe(self._bridge)
|
||||
|
||||
summarize = self._make_summarize_fn(model)
|
||||
runner = AgentRunner(agent, summarize_fn=summarize,
|
||||
# 🆕 M3: 重试回调 → Qt 信号（跨线程 queued 投递）
|
||||
on_retry_scheduled=lambda a, m, d, r:
|
||||
self.retry_scheduled.emit(int(a), int(m), float(d), str(r)),
|
||||
on_retry_finished=lambda ok:
|
||||
self.retry_finished.emit(bool(ok)),
|
||||
# 🆕 压缩开始回调 → 前端动态气泡
|
||||
on_compaction_started=lambda p:
|
||||
self.compaction_started.emit(str(p)),
|
||||
# 🆕 压缩完成回调 → 气泡即时定格（不等整轮结束）
|
||||
on_compaction_finished=lambda p:
|
||||
self.context_compacted.emit(dict(p)))
|
||||
# 🆕 接线轮中主动压缩检查（haocode 增强，偏离 pi 1:1）：
|
||||
# 内层循环每次 LLM 请求前主动检查，单条巨型工具输出不再依赖
|
||||
# 「一次失败往返 + 响应式兜底」
|
||||
agent.config.compact_fn = runner.compact_if_needed
|
||||
self._agent = agent
|
||||
self._runner = runner
|
||||
self._compactions_before = 0
|
||||
|
||||
try:
|
||||
result = runner.run(last_user)
|
||||
except Exception as e:
|
||||
if not self._aborted:
|
||||
self.error_occurred.emit(f"\n[Agent 异常]: {e}")
|
||||
return
|
||||
|
||||
if self._aborted:
|
||||
return # 取消路径：UI 已自行清理，不回退时间线
|
||||
|
||||
# 注：压缩事件（context_compacted）已由 on_compaction_finished 在压缩结束
|
||||
# 瞬间实时发出（气泡即时定格），这里不再重复 flush runner.compaction_events
|
||||
|
||||
if result.error is not None and result.stop_reason == "error":
|
||||
# 恢复逻辑（重试 + 压缩）全部用尽后的终局错误
|
||||
self.error_occurred.emit(
|
||||
f"\n[API 请求异常]: {result.error.message or '未知错误'}")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 供 UI 读取最终状态（入库用）
|
||||
# ------------------------------------------------------------------
|
||||
def final_messages(self) -> List[AgentMessage]:
|
||||
return list(self._agent.state.messages) if self._agent else []
|
||||
|
||||
|
||||
class ChatWorker(QThread):
|
||||
"""
|
||||
chat 模式 —— 普通聊天：单次流式补全（对照旧 LLMWorker 的行为）。
|
||||
不进入 agent 循环：无工具调用、无重试、无压缩。
|
||||
worker 模式请用 AgentWorker。
|
||||
"""
|
||||
|
||||
chunk_received = pyqtSignal(str)
|
||||
reasoning_received = pyqtSignal(str)
|
||||
error_occurred = pyqtSignal(str)
|
||||
# 🆕 P1: 收到精确 usage（UI 上下文标签 usage 锚定）
|
||||
usage_updated = pyqtSignal(dict)
|
||||
|
||||
def __init__(self, provider_name: str, model_name: str,
|
||||
openai_messages: List[Dict[str, Any]],
|
||||
temperature: Optional[float] = None):
|
||||
super().__init__()
|
||||
self.provider_name = provider_name
|
||||
self.model_name = model_name
|
||||
self.openai_messages = openai_messages
|
||||
self._temperature_override = temperature
|
||||
self.config = _load_config()
|
||||
self._cancelled = False
|
||||
|
||||
def cancel(self):
|
||||
self._cancelled = True
|
||||
|
||||
def abort(self):
|
||||
self._cancelled = True
|
||||
|
||||
def run(self):
|
||||
p = _provider_info(self.config, self.provider_name)
|
||||
if not p:
|
||||
self.error_occurred.emit(f"未找到 provider 配置: {self.provider_name}")
|
||||
return
|
||||
model = self._build_model_config(p)
|
||||
max_tokens = int(p.get("model_max_tokens", {}).get(self.model_name, 8192))
|
||||
try:
|
||||
client = OpenAI(api_key=p.get("api_key", ""),
|
||||
base_url=p.get("base_url", "https://api.openai.com/v1"),
|
||||
timeout=model.timeout_seconds, max_retries=0)
|
||||
response = client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=self.openai_messages,
|
||||
stream=True,
|
||||
temperature=model.temperature,
|
||||
max_tokens=max_tokens,
|
||||
# 🆕 P1: 拿精确 usage 供上下文标签锚定（对照 pi stream_options）
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
for chunk in response:
|
||||
if self._cancelled:
|
||||
try:
|
||||
response.close()
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
# 🆕 P1/M2: usage 读取（含 choice 兜底）
|
||||
_u = _pick_usage(chunk)
|
||||
if _u is not None and (
|
||||
getattr(_u, "prompt_tokens", 0) or getattr(_u, "completion_tokens", 0)):
|
||||
self.usage_updated.emit({
|
||||
"input": getattr(_u, "prompt_tokens", 0) or 0,
|
||||
"output": getattr(_u, "completion_tokens", 0) or 0,
|
||||
"cacheRead": 0, "cacheWrite": 0,
|
||||
})
|
||||
if not chunk.choices:
|
||||
continue
|
||||
delta = chunk.choices[0].delta
|
||||
# 🆕 M1: 思考字段优先级对照 pi（reasoning_content/reasoning/reasoning_text）
|
||||
_r = _pick_reasoning(delta)
|
||||
if _r:
|
||||
self.reasoning_received.emit(_r)
|
||||
if delta.content:
|
||||
self.chunk_received.emit(delta.content)
|
||||
except Exception as e:
|
||||
if not self._cancelled:
|
||||
self.error_occurred.emit(str(e))
|
||||
|
||||
def _build_model_config(self, p: dict) -> ModelConfig:
|
||||
context_window = int(p.get("model_contexts", {}).get(self.model_name, 128000))
|
||||
temperature = (self._temperature_override
|
||||
if self._temperature_override is not None
|
||||
else float(self.config.get("temperature", 0.7)))
|
||||
return ModelConfig(
|
||||
provider=self.provider_name,
|
||||
name=self.model_name,
|
||||
context_window=context_window,
|
||||
max_tokens=int(p.get("model_max_tokens", {}).get(self.model_name, 8192)),
|
||||
temperature=temperature,
|
||||
api_key=p.get("api_key", ""),
|
||||
base_url=p.get("base_url", "https://api.openai.com/v1"),
|
||||
# 🆕 thinking reasoning_content 回传默认开启
|
||||
)
|
||||
|
||||
|
||||
class TitleWorker(QThread):
|
||||
"""标题生成：单次流式补全（保留原信号面）"""
|
||||
|
||||
chunk_received = pyqtSignal(str)
|
||||
error_occurred = pyqtSignal(str)
|
||||
|
||||
def __init__(self, provider_name: str, model_name: str,
|
||||
openai_messages: List[Dict[str, Any]]):
|
||||
super().__init__()
|
||||
self.provider_name = provider_name
|
||||
self.model_name = model_name
|
||||
self.openai_messages = openai_messages
|
||||
self.config = _load_config()
|
||||
self._cancelled = False
|
||||
|
||||
def cancel(self):
|
||||
self._cancelled = True
|
||||
|
||||
def run(self):
|
||||
p = _provider_info(self.config, self.provider_name)
|
||||
if not p:
|
||||
return
|
||||
try:
|
||||
client = OpenAI(api_key=p.get("api_key", ""),
|
||||
base_url=p.get("base_url", "https://api.openai.com/v1"),
|
||||
timeout=60.0, max_retries=0)
|
||||
temperature = float(self.config.get("temperature", 0.7))
|
||||
response = client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=self.openai_messages,
|
||||
stream=True,
|
||||
temperature=temperature,
|
||||
max_tokens=60,
|
||||
)
|
||||
for chunk in response:
|
||||
if self._cancelled:
|
||||
try:
|
||||
response.close()
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
if chunk.choices:
|
||||
delta = chunk.choices[0].delta
|
||||
if delta.content:
|
||||
self.chunk_received.emit(delta.content)
|
||||
except Exception as e:
|
||||
if not self._cancelled:
|
||||
self.error_occurred.emit(str(e))
|
||||
@@ -0,0 +1,406 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
WebView2 后端（Windows 首选浏览器内核，失败自动回落 QtWebEngine）。
|
||||
|
||||
设计要点（来自 P1 实验 tests/_tmp_wv2_demo.py 的实测结论）：
|
||||
1. pythonnet + WebView2 .NET SDK（vendor/webview2/ 内 net462 Core.dll + webview2loader_x64.dll）
|
||||
2. 必须 OleInitialize（STA）后才能 CreateAsync
|
||||
3. 本机的 SDK 怪癖：传任何非空 user_data_folder 都报 RuntimeNotFound → 一律用默认 profile（ud=None）
|
||||
4. 残留 msedgewebview2.exe 会锁默认 profile（0x800700AA）→ 初始化前 taskkill
|
||||
5. .NET 版 CoreWebView2 不暴露子窗口 HWND → EnumChildWindows 找 Chrome_WidgetWin_* 类
|
||||
6. 子窗口天然是 Qt 顶层窗口的子 HWND，不需要 QWindow.fromWinId 包装（对子窗口会失败），
|
||||
由 Qt 布局算 slot 矩形后用 SetBoundsAndZoomFactor 驱动（父窗客户区物理像素）
|
||||
7. JS→Python 用 WebMessageReceived（JSON），Python→JS 用 ExecuteScriptAsync（与现有
|
||||
ChatBridge.run_js 生成的 JS 调用文本完全同构，前端零改动；仅 index.html 的
|
||||
window.bridge bootstrap 走双协议）
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import shutil
|
||||
import subprocess
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
WV2_DIR = os.path.join(ROOT, "vendor", "webview2")
|
||||
CORE_DLL = os.path.join(WV2_DIR, "net462_Microsoft.Web.WebView2.Core.dll")
|
||||
LOADER = os.path.join(WV2_DIR, "webview2loader_x64.dll")
|
||||
|
||||
_env = None
|
||||
_System = None # pythonnet 加载后缓存
|
||||
_INSTANCE_LOCK = {"fh": None}
|
||||
|
||||
|
||||
def _instance_lock_path():
|
||||
# 可覆盖（测试隔离用：HAOCODE_INSTANCE_LOCK_FILE 指向临时文件，
|
||||
# 否则测试会与正在运行的 app 争同一把锁 → 断言依赖环境）
|
||||
p = os.environ.get("HAOCODE_INSTANCE_LOCK_FILE")
|
||||
if p:
|
||||
return p
|
||||
d = os.path.join(ROOT, "data")
|
||||
try:
|
||||
os.makedirs(d, exist_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
return os.path.join(d, "app_instance.lock")
|
||||
|
||||
|
||||
def acquire_instance_lock():
|
||||
"""独占 data/app_instance.lock（非阻塞）。
|
||||
|
||||
返回 True = 本机唯一实例 → 可安全 taskkill 残留 msedgewebview2.exe
|
||||
False = 已有实例在跑 → 必须回落 QtWebEngine
|
||||
None = 平台不支持/异常（保守当“不确定”处理）
|
||||
|
||||
🐛 T0 根因：本函数存在之前，每个 MainWindow()（包括 offscreen 测试）都会走到
|
||||
_get_environment() 里的 `taskkill /F /IM msedgewebview2.exe`，
|
||||
把【当时正在运行的生产 app】的 WebView2 浏览器进程一并杀掉
|
||||
→ 它的 controller 变 disposed（set_bounds 报 0x8007139F）
|
||||
→ DOM 照渲染但视觉层永久空白（= “选中会话不渲染核心内容”）。
|
||||
"""
|
||||
if sys.platform != "win32":
|
||||
return None
|
||||
if _INSTANCE_LOCK["fh"] is not None:
|
||||
return True
|
||||
fh = None
|
||||
try:
|
||||
import msvcrt
|
||||
fh = open(_instance_lock_path(), "a+b")
|
||||
fh.seek(0) # 固定锁位置（锁定前 1 字节），不依赖 append 模式的当前位置
|
||||
msvcrt.locking(fh.fileno(), msvcrt.LK_NBLCK, 1)
|
||||
_INSTANCE_LOCK["fh"] = fh
|
||||
return True
|
||||
except Exception:
|
||||
try:
|
||||
if fh is not None:
|
||||
fh.close()
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def _wv2_allowed_here() -> bool:
|
||||
"""无头/测试环境一律不得碰 WebView2（共享默认 profile，一碰就可能误杀在生产实例）"""
|
||||
if os.environ.get("HAOCODE_FORCE_QTWEBENGINE", "") in ("1", "true", "True"):
|
||||
return False
|
||||
plat = (os.environ.get("QT_QPA_PLATFORM") or "").strip().lower()
|
||||
if plat and plat != "windows":
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _pump_wait(op, app, timeout=15.0):
|
||||
"""在 Qt 事件循环里等 .NET Task 完成（WebView2 初始化依赖消息泵）"""
|
||||
t0 = time.time()
|
||||
while not op.IsCompleted and time.time() - t0 < timeout:
|
||||
app.processEvents()
|
||||
time.sleep(0.005)
|
||||
if not op.IsCompleted:
|
||||
raise TimeoutError("WebView2 operation timed out")
|
||||
if op.IsFaulted:
|
||||
raise RuntimeError(str(op.Exception))
|
||||
return op.Result
|
||||
|
||||
|
||||
def _ts():
|
||||
"""启动链时间戳（WebView2 冷启动耗时诊断，保留）"""
|
||||
t = time.time()
|
||||
return time.strftime("%H:%M:%S", time.localtime(t)) + f".{int(t*1000) % 1000:03d}"
|
||||
|
||||
|
||||
def get_environment(app):
|
||||
"""初始化并返回 CoreWebView2Environment 单例；任何失败返回 None（调用方回落 QtWebEngine）"""
|
||||
global _env
|
||||
if _env is not None:
|
||||
return _env
|
||||
if sys.platform != "win32":
|
||||
return None
|
||||
if not (os.path.exists(CORE_DLL) and os.path.exists(LOADER)):
|
||||
return None
|
||||
# 🛡 守卫 1：无头/测试环境（QT_QPA_PLATFORM=offscreen 等）绝不启用 WebView2
|
||||
if not _wv2_allowed_here():
|
||||
print(f"[WV2] QT_QPA_PLATFORM={os.environ.get('QT_QPA_PLATFORM')!r} → 跳过 WebView2，"
|
||||
f"回落 QtWebEngine（无头环境不得触碰共享 profile）")
|
||||
return None
|
||||
# 🛡 守卫 2：本机已有实例在跑 → 不启用 WebView2、更不 taskkill（否则会把它的
|
||||
# 浏览器进程杀掉 → 对方 controller disposed → 聊天区永久空白）
|
||||
try:
|
||||
_lock_ok = acquire_instance_lock()
|
||||
except Exception:
|
||||
_lock_ok = None
|
||||
if _lock_ok is not True:
|
||||
print("[WV2] 检测到已有 haocode 实例在运行（instance lock 被占）→ "
|
||||
"本实例回落 QtWebEngine；已跳过 taskkill，不会影响对方渲染")
|
||||
return None
|
||||
try:
|
||||
os.environ["PATH"] = WV2_DIR + ";" + os.environ.get("PATH", "")
|
||||
try:
|
||||
os.add_dll_directory(WV2_DIR)
|
||||
except Exception:
|
||||
pass
|
||||
# .NET Core 的 P/Invoke 默认不查 CWD → 把 loader 复制到 CWD 一份
|
||||
try:
|
||||
cwd_loader = os.path.join(os.getcwd(), "WebView2Loader.dll")
|
||||
if not os.path.exists(cwd_loader):
|
||||
shutil.copyfile(LOADER, cwd_loader)
|
||||
except Exception:
|
||||
pass
|
||||
import ctypes
|
||||
ctypes.windll.ole32.OleInitialize(None) # STA（COM 初始化要求）
|
||||
# 清残留浏览器进程（锁默认 profile 会导致 0x800700AA）
|
||||
# 🛡 只有【本机唯一实例】才会走到这里（已在上面用 instance lock 保证），
|
||||
# 否则会把兄弟实例的浏览器进程杀掉 → 对方控制器 disposed → 聊天区空白
|
||||
print(f"[WV2] {_ts()} warmup: taskkill 残留进程（唯一实例，安全）...")
|
||||
try:
|
||||
subprocess.run(["taskkill", "/F", "/IM", "msedgewebview2.exe"],
|
||||
capture_output=True, timeout=10)
|
||||
time.sleep(1.0)
|
||||
except Exception:
|
||||
pass
|
||||
import clr
|
||||
global _System
|
||||
clr.AddReference(CORE_DLL)
|
||||
import System
|
||||
_System = System
|
||||
from Microsoft.Web.WebView2.Core import CoreWebView2Environment
|
||||
# ⚠️ ud=None：本机 SDK 对非空 user_data 路径报 RuntimeNotFound（实测怪癖）
|
||||
_env = _pump_wait(CoreWebView2Environment.CreateAsync(None, None, None), app)
|
||||
print(f"[WV2] {_ts()} Runtime ready: {_env.BrowserVersionString}")
|
||||
return _env
|
||||
except Exception as ex:
|
||||
_env = None
|
||||
print(f"[WV2] init failed → fallback to QtWebEngine: {ex}")
|
||||
return None
|
||||
|
||||
|
||||
class Wv2Session:
|
||||
"""一个 WebView2 实例：controller + 子窗口 + 消息泵 + bounds 驱动"""
|
||||
|
||||
def __init__(self, env, parent_hwnd, app):
|
||||
self.app = app
|
||||
print(f"[WV2] {_ts()} controller create begin (hwnd={parent_hwnd:#x})")
|
||||
self.controller = _pump_wait(
|
||||
env.CreateCoreWebView2ControllerAsync(_System.IntPtr(int(parent_hwnd))), app)
|
||||
print(f"[WV2] {_ts()} controller ready")
|
||||
self.core = self.controller.CoreWebView2
|
||||
try:
|
||||
self.core.Settings.AreDefaultContextMenusEnabled = False
|
||||
self.core.Settings.IsZoomControlEnabled = False
|
||||
self.core.Settings.IsStatusBarEnabled = False
|
||||
except Exception:
|
||||
pass
|
||||
self.child_hwnd = 0
|
||||
self.on_message = None # callable(dict) —— main_window 绑定到 ChatBridge
|
||||
self.on_load_finished = None # callable() —— WebView2View 绑定 loadFinished 信号
|
||||
self._msg_handler = self._make_msg_handler()
|
||||
self.core.add_WebMessageReceived(self._msg_handler)
|
||||
self._nw_handler = self._make_navigated_handler()
|
||||
self.core.add_NavigationCompleted(self._nw_handler)
|
||||
# 🆕 真异步 JS 执行队列：单一定时器轮询完成，回调从主线程定时器发出。
|
||||
# （旧版 execute_js_async 是阻塞忙等 + processEvents → 从定时器/事件回调里调用
|
||||
# 时产生重入嵌套事件循环 → COM 事件分发崩溃 → 启动卡死/拖动时 webview 不刷新）
|
||||
self._js_pending = [] # [task, cb, t0]
|
||||
from PyQt6.QtCore import QTimer
|
||||
self._js_pump = QTimer()
|
||||
self._js_pump.setInterval(25)
|
||||
self._js_pump.timeout.connect(self._js_pump_tick)
|
||||
self._js_pump.start()
|
||||
# 🆕 预热：立即导航 about:blank，让 msedgewebview2 进程/GPU 在 UI 构建期间冷启动
|
||||
# （实测本机首次真实页面导航需 12-15s，预热后降到 ~1s）
|
||||
try:
|
||||
print(f"[WV2] {_ts()} warmup: Navigate about:blank")
|
||||
self.core.Navigate("about:blank")
|
||||
except Exception as ex:
|
||||
print("[WV2] warmup navigate error:", ex)
|
||||
|
||||
# ---------- 事件 ----------
|
||||
def _make_msg_handler(self):
|
||||
import json
|
||||
from System import EventHandler
|
||||
from Microsoft.Web.WebView2.Core import CoreWebView2WebMessageReceivedEventArgs
|
||||
|
||||
def handler(sender, args):
|
||||
try:
|
||||
data = json.loads(args.WebMessageAsJson)
|
||||
if self.on_message:
|
||||
self.on_message(data)
|
||||
except Exception as ex:
|
||||
print("[WV2] message error:", ex)
|
||||
return EventHandler[CoreWebView2WebMessageReceivedEventArgs](handler)
|
||||
|
||||
def _make_navigated_handler(self):
|
||||
from System import EventHandler
|
||||
from Microsoft.Web.WebView2.Core import CoreWebView2NavigationCompletedEventArgs
|
||||
|
||||
def handler(sender, args):
|
||||
try:
|
||||
src = self.core.Source or ""
|
||||
except Exception:
|
||||
src = "?"
|
||||
print(f"[WV2] {_ts()} NavigationCompleted src={src}")
|
||||
# 预热页（about:blank）的加载完成不触发 loadFinished（避免 JS 探针空转）
|
||||
if src.startswith("about:blank"):
|
||||
return
|
||||
if self.on_load_finished:
|
||||
self.on_load_finished()
|
||||
return EventHandler[CoreWebView2NavigationCompletedEventArgs](handler)
|
||||
|
||||
# ---------- 子窗口发现（.NET 不暴露 HWND，轮询枚举） ----------
|
||||
def find_child_once(self):
|
||||
"""单次探测；找到返回 hwnd，否则 0（非阻塞，供 UI 线程定时器调）"""
|
||||
import ctypes
|
||||
import ctypes.wintypes as wt
|
||||
u32 = ctypes.windll.user32
|
||||
parent = wt.HWND(self.controller.ParentWindow.ToInt64())
|
||||
found = []
|
||||
|
||||
@ctypes.WINFUNCTYPE(wt.BOOL, wt.HWND, wt.LPVOID)
|
||||
def cb(h, _):
|
||||
buf = ctypes.create_unicode_buffer(256)
|
||||
u32.GetClassNameW(h, buf, 256)
|
||||
if buf.value.startswith("Chrome_WidgetWin"):
|
||||
found.append(h)
|
||||
return True
|
||||
|
||||
u32.EnumChildWindows(parent, cb, None)
|
||||
if found:
|
||||
self.child_hwnd = int(found[0])
|
||||
return self.child_hwnd
|
||||
return 0
|
||||
|
||||
def find_child(self):
|
||||
"""同步等待并返回 Chrome_WidgetWin_* 子 HWND（最多 10s）"""
|
||||
import ctypes
|
||||
import ctypes.wintypes as wt
|
||||
u32 = ctypes.windll.user32
|
||||
parent = wt.HWND(self.controller.ParentWindow.ToInt64())
|
||||
for _ in range(200):
|
||||
found = []
|
||||
|
||||
@ctypes.WINFUNCTYPE(wt.BOOL, wt.HWND, wt.LPVOID)
|
||||
def cb(h, _):
|
||||
buf = ctypes.create_unicode_buffer(256)
|
||||
u32.GetClassNameW(h, buf, 256)
|
||||
if buf.value.startswith("Chrome_WidgetWin"):
|
||||
found.append(h)
|
||||
return True
|
||||
|
||||
u32.EnumChildWindows(parent, cb, None)
|
||||
if found:
|
||||
self.child_hwnd = int(found[0])
|
||||
return self.child_hwnd
|
||||
self.app.processEvents()
|
||||
time.sleep(0.05)
|
||||
return 0
|
||||
|
||||
# ---------- 几何 ----------
|
||||
def set_visible(self, visible: bool):
|
||||
try:
|
||||
self.controller.IsVisible = bool(visible)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def set_bounds(self, left, top, width, height):
|
||||
"""父窗客户区物理像素"""
|
||||
try:
|
||||
from System.Drawing import Rectangle as _Rect
|
||||
self.controller.SetBoundsAndZoomFactor(_Rect(int(left), int(top),
|
||||
int(width), int(height)), 1.0)
|
||||
except Exception as ex:
|
||||
print("[WV2] set_bounds error:", ex)
|
||||
|
||||
def child_size(self):
|
||||
"""子窗口当前屏幕像素尺寸（观察用）"""
|
||||
if not self.child_hwnd:
|
||||
return (0, 0)
|
||||
import ctypes
|
||||
import ctypes.wintypes as wt
|
||||
r = wt.RECT()
|
||||
if ctypes.windll.user32.GetWindowRect(wt.HWND(self.child_hwnd), ctypes.byref(r)):
|
||||
return (r.right - r.left, r.bottom - r.top)
|
||||
return (0, 0)
|
||||
|
||||
# ---------- JS ----------
|
||||
def navigate(self, url: str):
|
||||
print(f"[WV2] {_ts()} navigate {url[:80]}")
|
||||
try:
|
||||
self.core.Navigate(url)
|
||||
except Exception as ex:
|
||||
print("[WV2] navigate error:", ex)
|
||||
|
||||
def execute_js(self, script: str):
|
||||
"""fire-and-forget（ChatBridge.run_js 的替换，JS 文本完全同构）"""
|
||||
self._js_run(script, None)
|
||||
|
||||
def execute_js_async(self, script: str, cb):
|
||||
"""带回调执行：真异步，回调在主线程定时器 tick 中发出（绝不阻塞）"""
|
||||
self._js_run(script, cb)
|
||||
|
||||
def _js_run(self, script: str, cb):
|
||||
try:
|
||||
task = self.core.ExecuteScriptWithResultAsync(script)
|
||||
except Exception as ex:
|
||||
print("[WV2] execute_js error:", ex)
|
||||
if cb:
|
||||
try:
|
||||
cb(None)
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
self._js_pending.append([task, cb, time.time()])
|
||||
|
||||
def _js_pump_tick(self):
|
||||
if not self._js_pending:
|
||||
return
|
||||
import json
|
||||
remaining = []
|
||||
for task, cb, t0 in self._js_pending:
|
||||
done = False
|
||||
try:
|
||||
done = task.IsCompleted
|
||||
except Exception:
|
||||
done = True
|
||||
if not done and time.time() - t0 > 10:
|
||||
done = True # 10s 安全超时（渲染器死亡时不永久卡队列）
|
||||
if not done:
|
||||
remaining.append([task, cb, t0])
|
||||
continue
|
||||
result = None
|
||||
try:
|
||||
if task.IsCompleted and not task.IsFaulted:
|
||||
r = task.Result # CoreWebView2ExecuteScriptResult 包装结构
|
||||
if getattr(r, "Succeeded", True):
|
||||
s = r.ResultAsJson # JSON 编码字符串（或 null）
|
||||
if s:
|
||||
result = json.loads(s)
|
||||
if isinstance(result, str):
|
||||
try:
|
||||
result = json.loads(result)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as ex:
|
||||
print("[WV2] js result decode error:", ex)
|
||||
if cb:
|
||||
try:
|
||||
cb(result)
|
||||
except Exception as ex:
|
||||
print("[WV2] js callback error:", ex)
|
||||
self._js_pending = remaining
|
||||
|
||||
def close(self):
|
||||
try:
|
||||
self._js_pump.stop()
|
||||
except Exception:
|
||||
pass
|
||||
self._js_pending = []
|
||||
try:
|
||||
self.core.remove_WebMessageReceived(self._msg_handler)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
self.core.remove_NavigationCompleted(self._nw_handler)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
self.controller.Close()
|
||||
except Exception:
|
||||
pass
|
||||
Reference in New Issue
Block a user