wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修改

Browse Source
This commit is contained in:
hupenglong1
2026-05-22 19:48:47 +08:00
parent 20690cdef3
commit 5311e6d97c
31 changed files with 4620 additions and 472 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent_prompting.py
+3
View File
@@ -249,6 +249,9 @@ def get_using_your_tools_section(enabled_tool_names: set[str]) -> str:
items.append(
'需要可靠保存文件时，可以在 python_exec.code 中使用 pathlib.Path(...).write_text(...)、json.dump/json.dumps、csv.writer 或流式逐行写入；这比把大段内容塞进 write_file 参数更稳定。'
)
items.append(
'python_exec 默认 timeout=30s。读 CSV/JSONL 大文件、pandas 处理、批量校验、训练评测脚本等可能 >30s 的任务，调用时显式传 timeout_seconds（常规 120-300、训练/长跑 600+），不要靠默认值。被 timeout(1) kill 后会以 timed_out=true 报错并提示同样的内容，看到就重跑并调大。'
)
if 'python_package' in enabled_tool_names:
items.append(
'当 python_exec 因缺少 pandas、pyarrow、openpyxl 等 Python 包失败时，使用 python_package 在当前用户独立 venv 中安装缺失包，然后重试；不要安装到系统 Python 或项目 .venv。'
src/agent_tool_core.py
+7
View File
@@ -10,6 +10,7 @@ from .agent_types import AgentPermissions, AgentRuntimeConfig, ToolExecutionResu
if TYPE_CHECKING:
from .account_runtime import AccountRuntime
from .ask_user_runtime import AskUserRuntime
from .bash_bg_store import BgTaskSpec, BgTaskStatus
from .config_runtime import ConfigRuntime
from .lsp_runtime import LSPRuntime
from .mcp_runtime import MCPRuntime
@@ -58,6 +59,12 @@ class ToolExecutionContext:
team_runtime: 'TeamRuntime | None' = None
workflow_runtime: 'WorkflowRuntime | None' = None
worktree_runtime: 'WorktreeRuntime | None' = None
bg_register: Callable[['BgTaskSpec'], None] | None = None
bg_status_query: Callable[[str], 'BgTaskStatus | None'] | None = None
bg_kill_request: Callable[[str], bool] | None = None
account_id: str | None = None
session_id: str | None = None
run_id: str | None = None
ToolHandler = Callable[
src/agent_tool_specs/execution.py
+69 -1
View File
@@ -44,7 +44,7 @@ def build_execution_tools(handlers: Mapping[str, ToolHandler]) -> list[AgentTool
'timeout_seconds': {
'type': 'number',
'minimum': 1,
'description': '可选超时时间，默认使用当前会话命令超时。',
'description': '可选超时时间，默认 30 秒。读大 CSV/JSONL、pandas 分析、批量校验、训练评测脚本等可能 >30s 的任务必须显式传值（常规 120-300、训练/长跑 600+）。',
},
'max_output_chars': {
'type': 'integer',
@@ -98,16 +98,84 @@ def build_execution_tools(handlers: Mapping[str, ToolHandler]) -> list[AgentTool
'不要用 bash 执行 python、python3、pip 或 .venv/bin/python；结构化文件分析、'
'JSON/JSONL 处理、批量校验、数据抽样和快速计算必须优先使用 python_exec。'
'Python 包安装必须优先使用 python_package。'
'需要等待远端长任务（>30s 的训练、文件落盘、外部 API 回调）时，'
'设置 run_in_background=true 让命令在后台 detach 跑、立刻返回 task_id；'
'默认 wait_for_completion=true 时进程结束后会自动起新一轮把结果送回。'
'后台任务输出落在 <remote_workspace>/.bg_tasks/<task_id>/output。'
),
parameters={
'type': 'object',
'properties': {
'command': {'type': 'string'},
'run_in_background': {
'type': 'boolean',
'description': (
'后台 detach 跑。立刻返回 task_id，不等命令结束。'
'适合 >30s 的远端任务、需要等异步事件的场景。'
),
},
'wait_for_completion': {
'type': 'boolean',
'description': (
'仅在 run_in_background=true 时生效。'
'true（默认）：进程结束后由后端自动起新一轮把结果作为 system 消息送回；'
'false：不自动续，agent 需自行调 bash_status 取结果。'
),
},
},
'required': ['command'],
},
handler=resolve_handler(handlers, 'bash', 'execution'),
),
AgentTool(
name='bash_status',
description=(
'查询通过 bash(run_in_background=true) 启动的后台任务状态。'
'block=true 会阻塞到任务终止或 timeout_seconds 到期再返回；'
'block=false（默认）立刻返回当前快照。'
'返回 status (running/completed/failed/cancelled)、exit_code 和最近输出预览。'
),
parameters={
'type': 'object',
'properties': {
'task_id': {
'type': 'string',
'description': 'bash run_in_background 启动时返回的 task_id。',
},
'block': {
'type': 'boolean',
'description': '是否等到任务终止再返回；默认 false。',
},
'timeout_seconds': {
'type': 'number',
'minimum': 0,
'maximum': 600,
'description': 'block=true 时最长等待秒数，默认 30，上限 600。',
},
},
'required': ['task_id'],
},
handler=resolve_handler(handlers, 'bash_status', 'execution'),
),
AgentTool(
name='bash_kill',
description=(
'主动终止通过 bash(run_in_background=true) 启动的后台任务。'
'注意：用户在前端按 stop 取消当前 run 不会自动杀后台进程（detach 的本意），'
'需要明确调用本工具才能终止远端进程。'
),
parameters={
'type': 'object',
'properties': {
'task_id': {
'type': 'string',
'description': 'bash run_in_background 启动时返回的 task_id。',
},
},
'required': ['task_id'],
},
handler=resolve_handler(handlers, 'bash_kill', 'execution'),
),
AgentTool(
name='sleep',
description='Pause execution briefly for bounded local wait flows.',
src/agent_tools.py
+321 -39
View File
@@ -7,7 +7,9 @@ import io
import json
import os
import re
import secrets
import selectors
import shlex
import shutil
import signal
import subprocess
@@ -32,6 +34,7 @@ from .agent_tool_specs.data_agent import build_data_agent_tools
from .agent_tool_specs.execution import build_execution_tools
from .agent_tool_specs.files import build_file_tools
from .agent_types import DEFAULT_MAX_TURNS, ToolExecutionResult
from .bash_bg_store import BgTaskSpec, BgTaskStatus
from .data_agent_records import (
DataRecordError,
confirm_generation_goal,
@@ -117,6 +120,8 @@ def default_tool_registry() -> dict[str, AgentTool]:
'python_exec': _run_python_exec,
'python_package': _run_python_package,
'bash': _run_bash,
'bash_status': _run_bash_status,
'bash_kill': _run_bash_kill,
'sleep': _sleep,
}),
AgentTool(
@@ -1734,46 +1739,29 @@ def _ensure_shell_allowed(command: str, context: ToolExecutionContext) -> None:
'Shell command appears to write platform code. '
'Platform code directories are read-only in data-agent sessions.'
)
if context.permissions.allow_destructive_shell_commands:
return
destructive_patterns = [
r'(^|[;&|])\s*rm\s',
r'(^|[;&|])\s*mv\s',
r'(^|[;&|])\s*dd\s',
r'(^|[;&|])\s*shutdown\s',
r'(^|[;&|])\s*reboot\s',
r'(^|[;&|])\s*mkfs',
r'(^|[;&|])\s*chmod\s+-r\s+777',
r'(^|[;&|])\s*chown\s+-r',
r'(^|[;&|])\s*git\s+reset\s+--hard',
r'(^|[;&|])\s*git\s+clean\s+-fd',
r'(^|[;&|])\s*:\s*>\s*',
]
lowered = command.lower()
if any(re.search(pattern, lowered) for pattern in destructive_patterns):
raise ToolPermissionError(
'Potentially destructive shell command blocked. Re-run with --unsafe to allow it.'
)
return
def _looks_like_platform_code_shell_write(
command: str,
context: ToolExecutionContext,
) -> bool:
if context.jupyter_runtime is not None:
return False
if not _is_platform_app_root(context.root):
return False
lowered = command.lower()
write_marker = re.search(
r'>|(^|[\s;&|])(tee|touch|python|python3|sed\s+-i|perl\s+-pi)\b',
r'(?<![0-9&])>(?!&)|(^|[\s;&|])(tee|touch|python|python3|sed\s+-i|perl\s+-pi)\b',
lowered,
)
if write_marker is None:
return False
root_prefix = context.root.as_posix().lower()
platform_fragments = [
f'{name}/' for name in _PLATFORM_READONLY_DIRS
f'{root_prefix}/{name}/' for name in _PLATFORM_READONLY_DIRS
] + [
f'{context.root.as_posix().lower()}/{name}/'
for name in _PLATFORM_READONLY_DIRS
f'./{name}/' for name in _PLATFORM_READONLY_DIRS
]
return any(fragment in lowered for fragment in platform_fragments)
@@ -2322,6 +2310,26 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
]
if result.timed_out:
payload.insert(0, f'timed_out=true\ntimeout_seconds={timeout_seconds:g}')
# 被 timeout(1) 砍掉时 ok=True 会让 agent 误以为脚本跑完了；翻成 ok=False，
# 错误消息里直接给出建议（调大 timeout、拆分操作、流式读 CSV 等）。
raise ToolExecutionError(
_truncate_output(
'\n'.join(
[
*payload,
(
f'[hint] python_exec 在 {timeout_seconds:g}s 被 timeout(1) kill，没有跑完。\n'
'首选做法：把 timeout_seconds 调大（数据分析类任务起步给 300，处理大 CSV/JSONL 给 600+）。\n'
'不要把脚本拆成多个更小的 python_exec 调用——每次调用都要重新 import、重读数据、'
'重跑一遍带完整对话上下文的远端 round-trip，拆得越碎总耗时越长。'
'把多步分析（解析、聚合、统计、写出）合并到一个脚本里才是对的。\n'
'只有当数据本身远超内存或单次跑要 1h+ 时，才考虑 pandas chunksize / 流式过滤先把数据降量。'
),
]
).strip(),
max_output_chars,
)
)
return (
_truncate_output('\n'.join(payload).strip(), max_output_chars),
{
@@ -2393,21 +2401,17 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
stdout.rstrip(),
'[stderr]',
stderr.rstrip(),
(
f'[hint] python_exec 在 {timeout_seconds:g}s 被 kill，没有跑完。\n'
'首选做法：把 timeout_seconds 调大（数据分析类任务起步给 300，处理大 CSV/JSONL 给 600+）。\n'
'不要把脚本拆成多个更小的 python_exec 调用——每次调用都要重新 import、重读数据、'
'重跑一遍带完整对话上下文的远端 round-trip，拆得越碎总耗时越长。'
'把多步分析（解析、聚合、统计、写出）合并到一个脚本里才是对的。\n'
'只有当数据本身远超内存或单次跑要 1h+ 时，才考虑 pandas chunksize / 流式过滤先把数据降量。'
),
]
return (
_truncate_output('\n'.join(payload).strip(), max_output_chars),
{
'action': 'python_exec',
'mode': mode,
'interpreter': interpreter,
'scratchpad_directory': str(context.scratchpad_directory) if context.scratchpad_directory else '',
'python_env_dir': str(context.python_env_dir) if context.python_env_dir else '',
'timed_out': True,
'timeout_seconds': timeout_seconds,
'stdout_preview': _snapshot_text(stdout),
'stderr_preview': _snapshot_text(stderr),
'output_preview': _snapshot_text('\n'.join(payload).strip()),
},
raise ToolExecutionError(
_truncate_output('\n'.join(payload).strip(), max_output_chars)
)
except ToolExecutionError:
raise
@@ -2690,7 +2694,15 @@ def _resolve_python_interpreter(context: ToolExecutionContext) -> str:
def _run_bash(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
command = _require_string(arguments, 'command')
if arguments.get('run_in_background'):
return _run_bash_background(arguments, context)
raw_command = arguments.get('command')
if not isinstance(raw_command, str) or not raw_command.strip():
return (
'(no-op: empty bash command — pass a non-empty `command` string '
'to actually execute something)'
)
command = raw_command
_ensure_shell_allowed(command, context)
if context.jupyter_runtime is not None:
result = context.jupyter_runtime.run_command(
@@ -2752,6 +2764,276 @@ def _run_bash(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
)
def _new_bg_task_id() -> str:
return f'bg_{secrets.token_hex(4)}'
def _bg_extract_pid(stdout: str) -> int | None:
match = re.search(r'BG_TASK_PID=(\d+)', stdout)
if match is None:
return None
try:
return int(match.group(1))
except ValueError:
return None
def _bg_remote_launcher(
*,
task_dir: str,
output_path: str,
pid_path: str,
exit_code_path: str,
user_command: str,
) -> str:
inner = (
f'( {user_command} ) > {shlex.quote(output_path)} 2>&1'
f' ; echo $? > {shlex.quote(exit_code_path)}'
)
return (
f'mkdir -p {shlex.quote(task_dir)}\n'
f'nohup bash -c {shlex.quote(inner)} </dev/null '
'>/dev/null 2>&1 &\n'
'PID=$!\n'
f'echo $PID > {shlex.quote(pid_path)}\n'
'echo "BG_TASK_PID=$PID"\n'
)
def _run_bash_background(
arguments: dict[str, Any],
context: ToolExecutionContext,
) -> tuple[str, dict[str, Any]]:
raw_command = arguments.get('command')
if not isinstance(raw_command, str) or not raw_command.strip():
return (
'(no-op: empty bash command — pass a non-empty `command` string '
'to actually execute something)',
{'action': 'bash_background_noop'},
)
command = raw_command
_ensure_shell_allowed(command, context)
wait_for_completion = bool(arguments.get('wait_for_completion', True))
task_id = _new_bg_task_id()
if context.jupyter_runtime is not None:
workspace = context.jupyter_runtime.binding.workspace_cwd
task_dir = f'{workspace}/.bg_tasks/{task_id}'
output_path = f'{task_dir}/output'
pid_path = f'{task_dir}/pid'
exit_code_path = f'{task_dir}/exit_code'
launcher = _bg_remote_launcher(
task_dir=task_dir,
output_path=output_path,
pid_path=pid_path,
exit_code_path=exit_code_path,
user_command=command,
)
result = context.jupyter_runtime.run_command(
launcher,
timeout_seconds=min(context.command_timeout_seconds, 15.0),
max_output_chars=4096,
cancel_event=context.cancel_event,
)
if result.exit_code != 0:
raise ToolExecutionError(
f'后台 bash 启动失败 exit_code={result.exit_code}\n'
f'stderr={result.stderr.strip()}'
)
pid = _bg_extract_pid(result.stdout)
if pid is None:
raise ToolExecutionError(
f'后台 bash 启动后未拿到 pid\nstdout={result.stdout.strip()}'
)
spec = BgTaskSpec(
task_id=task_id,
account_key=context.account_id or '',
session_id=context.session_id or '',
run_id=context.run_id or '',
pid=pid,
task_dir=task_dir,
output_path=output_path,
pid_path=pid_path,
exit_code_path=exit_code_path,
command=command,
started_at=time.time(),
wait_for_completion=wait_for_completion,
)
if context.bg_register is not None:
context.bg_register(spec)
content = (
f'已在远端后台启动\n'
f'task_id={task_id}\n'
f'pid={pid}\n'
f'task_dir={task_dir}\n'
f'output={output_path}\n'
f'wait_for_completion={wait_for_completion}\n'
'完成后将自动起新一轮把结果送回；'
'中途可用 bash_status(task_id) 查看，bash_kill(task_id) 终止。'
)
return (
content,
{
'action': 'remote_bash_bg',
'task_id': task_id,
'pid': pid,
'task_dir': task_dir,
'output_path': output_path,
'command': command,
'wait_for_completion': wait_for_completion,
},
)
# Local fallback: spawn detached subprocess with file redirects.
base_dir = (context.scratchpad_directory or context.root) / '.bg_tasks' / task_id
base_dir.mkdir(parents=True, exist_ok=True)
output_path = base_dir / 'output'
pid_path = base_dir / 'pid'
exit_code_path = base_dir / 'exit_code'
inner = (
f'( {command} ) > {shlex.quote(str(output_path))} 2>&1'
f' ; echo $? > {shlex.quote(str(exit_code_path))}'
)
process = subprocess.Popen(
['bash', '-c', inner],
cwd=_execution_cwd(context),
stdin=subprocess.DEVNULL,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
env=_build_subprocess_env(context),
start_new_session=True,
)
pid_path.write_text(f'{process.pid}\n', encoding='utf-8')
spec = BgTaskSpec(
task_id=task_id,
account_key=context.account_id or '',
session_id=context.session_id or '',
run_id=context.run_id or '',
pid=process.pid,
task_dir=str(base_dir),
output_path=str(output_path),
pid_path=str(pid_path),
exit_code_path=str(exit_code_path),
command=command,
started_at=time.time(),
wait_for_completion=wait_for_completion,
)
if context.bg_register is not None:
context.bg_register(spec)
content = (
f'已在本地后台启动\n'
f'task_id={task_id}\n'
f'pid={process.pid}\n'
f'output={output_path}\n'
)
return (
content,
{
'action': 'bash_bg',
'task_id': task_id,
'pid': process.pid,
'task_dir': str(base_dir),
'output_path': str(output_path),
'command': command,
'wait_for_completion': wait_for_completion,
},
)
def _format_bash_status(status: BgTaskStatus, max_output_chars: int) -> str:
finished = (
time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(status.finished_at))
if status.finished_at is not None
else '-'
)
started = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(status.started_at))
body = [
f'task_id={status.task_id}',
f'status={status.status}',
f'pid={status.pid}',
f'started_at={started}',
f'finished_at={finished}',
f'exit_code={status.exit_code if status.exit_code is not None else "-"}',
f'auto_resumed={status.auto_resumed}',
'[output_preview]',
status.output_preview or '(empty)',
]
return _truncate_output('\n'.join(body), max_output_chars)
def _run_bash_status(
arguments: dict[str, Any],
context: ToolExecutionContext,
) -> tuple[str, dict[str, Any]]:
task_id = _require_string(arguments, 'task_id')
if context.bg_status_query is None:
raise ToolExecutionError('bg_status_query 未注入；当前运行环境不支持 bash 后台任务')
block = bool(arguments.get('block', False))
timeout_seconds = float(arguments.get('timeout_seconds', 30))
timeout_seconds = max(0.0, min(timeout_seconds, 600.0))
deadline = time.monotonic() + timeout_seconds if block else None
while True:
status = context.bg_status_query(task_id)
if status is None:
raise ToolExecutionError(f'未找到 task_id={task_id}')
if not block or status.status != 'running':
return (
_format_bash_status(status, context.max_output_chars),
{
'action': 'bash_status',
'task_id': task_id,
'status': status.status,
'exit_code': status.exit_code,
'auto_resumed': status.auto_resumed,
},
)
if deadline is None or time.monotonic() >= deadline:
return (
_format_bash_status(status, context.max_output_chars),
{
'action': 'bash_status',
'task_id': task_id,
'status': status.status,
'timed_out_block': True,
},
)
if context.cancel_event is not None and context.cancel_event.is_set():
return (
_format_bash_status(status, context.max_output_chars),
{
'action': 'bash_status',
'task_id': task_id,
'status': status.status,
'cancelled': True,
},
)
time.sleep(min(2.0, max(0.5, deadline - time.monotonic())))
def _run_bash_kill(
arguments: dict[str, Any],
context: ToolExecutionContext,
) -> tuple[str, dict[str, Any]]:
task_id = _require_string(arguments, 'task_id')
if context.bg_kill_request is None:
raise ToolExecutionError('bg_kill_request 未注入；当前运行环境不支持 bash 后台任务')
killed = context.bg_kill_request(task_id)
content = (
f'已终止 task_id={task_id}'
if killed
else f'未能终止 task_id={task_id}（可能已结束或不存在）'
)
return (
content,
{
'action': 'bash_kill',
'task_id': task_id,
'killed': killed,
},
)
def _web_fetch(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
raw_url = _require_string(arguments, 'url')
max_chars = _coerce_int(arguments, 'max_chars', context.max_output_chars)
src/agent_types.py
+1 -1
View File
@@ -157,7 +157,7 @@ DEFAULT_MAX_TURNS = 50
class AgentRuntimeConfig:
cwd: Path
max_turns: int = DEFAULT_MAX_TURNS
command_timeout_seconds: float = 30.0
command_timeout_seconds: float = 300.0
max_output_chars: int = 50000 # ≈ 20k token (Chinese/code mix ≈ 2.5 chars/token)
stream_model_responses: bool = False
auto_snip_threshold_tokens: int | None = None
src/bash_bg_store.py
+211
View File
@@ -0,0 +1,211 @@
"""Persist background bash task records.
Sibling to run_state_store.py: a SQLite table tracking bg shell tasks the
agent spawns via bash(run_in_background=true). The store survives backend
restarts so polling can be resumed; the actual remote process is owned by
nohup on the Jupyter terminal and is independent of this store.
"""
from __future__ import annotations
import sqlite3
import time
from dataclasses import dataclass
from pathlib import Path
from threading import RLock
from typing import Any
ACTIVE_BG_STATUSES = {'running'}
TERMINAL_BG_STATUSES = {'completed', 'failed', 'cancelled'}
@dataclass(frozen=True)
class BgTaskSpec:
"""Payload an agent-side handler hands to the backend on registration."""
task_id: str
account_key: str
session_id: str
run_id: str
pid: int
task_dir: str
output_path: str
pid_path: str
exit_code_path: str
command: str
started_at: float
wait_for_completion: bool
@dataclass(frozen=True)
class BgTaskStatus:
"""Snapshot the agent reads back via bash_status."""
task_id: str
status: str
pid: int
started_at: float
finished_at: float | None
exit_code: int | None
output_path: str
output_preview: str
auto_resumed: bool
class BashBgStore:
def __init__(self, db_path: Path) -> None:
self.db_path = db_path
self._lock = RLock()
def record_start(self, spec: BgTaskSpec) -> None:
now = time.time()
with self._connect() as conn:
conn.execute(
"""
insert into bash_bg_tasks (
task_id, account_key, session_id, run_id, pid,
task_dir, output_path, pid_path, exit_code_path,
command, status, started_at, updated_at,
wait_for_completion, auto_resumed
)
values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'running', ?, ?, ?, 0)
on conflict(task_id) do nothing
""",
(
spec.task_id,
spec.account_key,
spec.session_id,
spec.run_id,
spec.pid,
spec.task_dir,
spec.output_path,
spec.pid_path,
spec.exit_code_path,
spec.command,
spec.started_at,
now,
1 if spec.wait_for_completion else 0,
),
)
def mark_completed(
self,
task_id: str,
*,
exit_code: int,
finished_at: float | None = None,
) -> None:
finished = finished_at if finished_at is not None else time.time()
status = 'completed' if exit_code == 0 else 'failed'
with self._connect() as conn:
conn.execute(
"""
update bash_bg_tasks
set status = ?, exit_code = ?, finished_at = ?, updated_at = ?
where task_id = ? and status = 'running'
""",
(status, exit_code, finished, finished, task_id),
)
def mark_killed(self, task_id: str) -> None:
now = time.time()
with self._connect() as conn:
conn.execute(
"""
update bash_bg_tasks
set status = 'cancelled', finished_at = ?, updated_at = ?
where task_id = ? and status = 'running'
""",
(now, now, task_id),
)
def mark_auto_resumed(self, task_id: str) -> bool:
"""Set auto_resumed=1 atomically; return True if we won the race."""
with self._connect() as conn:
cursor = conn.execute(
"""
update bash_bg_tasks
set auto_resumed = 1, updated_at = ?
where task_id = ? and auto_resumed = 0
""",
(time.time(), task_id),
)
return cursor.rowcount > 0
def list_active(self) -> list[dict[str, Any]]:
with self._connect() as conn:
rows = conn.execute(
"select * from bash_bg_tasks where status = 'running'"
).fetchall()
return [dict(row) for row in rows]
def list_for_session(
self,
account_key: str,
session_id: str,
) -> list[dict[str, Any]]:
with self._connect() as conn:
rows = conn.execute(
"""
select * from bash_bg_tasks
where account_key = ? and session_id = ?
order by started_at desc
""",
(account_key, session_id),
).fetchall()
return [dict(row) for row in rows]
def get(self, task_id: str) -> dict[str, Any] | None:
with self._connect() as conn:
row = conn.execute(
"select * from bash_bg_tasks where task_id = ?",
(task_id,),
).fetchone()
return dict(row) if row is not None else None
def _connect(self) -> sqlite3.Connection:
with self._lock:
self.db_path.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(self.db_path, timeout=10)
conn.row_factory = sqlite3.Row
self._init_schema(conn)
return conn
def _init_schema(self, conn: sqlite3.Connection) -> None:
conn.execute('pragma journal_mode=wal')
conn.execute(
"""
create table if not exists bash_bg_tasks (
task_id text primary key,
account_key text not null,
session_id text not null,
run_id text not null default '',
pid integer not null,
task_dir text not null,
output_path text not null,
pid_path text not null,
exit_code_path text not null,
command text not null,
status text not null,
started_at real not null,
finished_at real,
exit_code integer,
updated_at real not null,
wait_for_completion integer not null default 1,
auto_resumed integer not null default 0
)
"""
)
conn.execute(
"""
create index if not exists idx_bash_bg_session
on bash_bg_tasks(account_key, session_id, started_at)
"""
)
conn.execute(
"""
create index if not exists idx_bash_bg_active
on bash_bg_tasks(status)
"""
)
src/bash_security.py
-7
View File
@@ -1251,11 +1251,4 @@ def check_shell_security(
if result.is_misparsing:
return (False, f'Security check: {result.message}')
# Check destructive patterns if not allowed
if not allow_destructive:
warning = get_destructive_command_warning(command)
if warning:
return (False, f'Potentially destructive command blocked: {warning}. '
'Re-run with --unsafe to allow it.')
return (True, '')
src/jupyter_runtime.py
+55 -1
View File
@@ -233,12 +233,41 @@ class JupyterRuntimeSession:
'persisted_at': time.time(),
}
@property
def chat_workspace_root(self) -> str:
"""Per-user-per-chat root on shared NFS so the SFT training pod
(which doesn't mount the jupyter pod's private workspace) can read
and write the same artifacts. Layout:
/mnt/wangsenhao/autoresearch-zk-users/<account_id>/<session_id>/"""
account_part = sanitize_remote_path_part(self.binding.account_id)
session_part = sanitize_remote_path_part(self.binding.session_id)
return (
f'/mnt/wangsenhao/autoresearch-zk-users/'
f'{account_part}/{session_part}'
)
def bootstrap_workspace(
self,
*,
timeout_seconds: float = 20.0,
project_root: Path | None = None,
) -> None:
chat_root = self.chat_workspace_root
nfs_users_root = '/mnt/wangsenhao/autoresearch-zk-users'
probe = self.run_command(
(
f'mkdir -p {shlex.quote(nfs_users_root)} && '
f'touch {shlex.quote(nfs_users_root)}/.write_probe && '
f'rm -f {shlex.quote(nfs_users_root)}/.write_probe'
),
timeout_seconds=timeout_seconds,
max_output_chars=2000,
)
if probe.exit_code != 0:
raise JupyterRuntimeError(
'远端 jupyter pod 未挂载 /mnt/wangsenhao 或不可写：'
+ (probe.stdout.strip() or probe.stderr.strip() or 'unknown error')
)
result = self.run_command(
(
'mkdir -p '
@@ -247,7 +276,10 @@ class JupyterRuntimeSession:
f'{shlex.quote(self.binding.workspace_cwd)}/scratchpad '
f'{shlex.quote(self.binding.workspace_root)}/.runtime/uploads '
f'{shlex.quote(self.binding.workspace_root)}/runtime_uploads '
f'{shlex.quote(self.account_runtime_root)}/python'
f'{shlex.quote(self.account_runtime_root)}/python '
f'{shlex.quote(chat_root)}/results '
f'{shlex.quote(chat_root)}/sft_output '
f'{shlex.quote(chat_root)}/scripts'
),
timeout_seconds=timeout_seconds,
max_output_chars=4000,
@@ -259,6 +291,27 @@ class JupyterRuntimeSession:
self.ensure_python_environment(timeout_seconds=max(timeout_seconds, 300.0))
if project_root is not None:
self.sync_runtime_bundle(project_root, timeout_seconds=max(timeout_seconds, 180.0))
self._sync_chat_workspace_scripts(timeout_seconds=timeout_seconds)
def _sync_chat_workspace_scripts(self, *, timeout_seconds: float = 20.0) -> None:
if not self.skills_root:
return
chat_root = self.chat_workspace_root
src = f'{self.skills_root}/model-iteration/scripts/prepare_and_train_sft.py'
dst = f'{chat_root}/scripts/prepare_and_train_sft.py'
result = self.run_command(
(
f'if [ -f {shlex.quote(src)} ]; then '
f'cp {shlex.quote(src)} {shlex.quote(dst)}; '
f'fi'
),
timeout_seconds=timeout_seconds,
max_output_chars=2000,
)
if result.exit_code != 0:
raise JupyterRuntimeError(
result.stdout.strip() or 'Unable to sync chat workspace SFT script.'
)
def to_dict(self) -> dict[str, Any]:
payload = self.binding.to_dict()
@@ -835,6 +888,7 @@ print(f"{'appended' if MODE == 'a' else 'wrote'} {target} ({len(content)} chars)
'ZK_AGENT_OUTPUT': f'{self.binding.workspace_cwd}/output',
'ZK_AGENT_SCRATCHPAD': f'{self.binding.workspace_cwd}/scratchpad',
'ZK_AGENT_PYTHON_ENV': self.python_env_path,
'AUTORESEARCH_CHAT_ROOT': self.chat_workspace_root,
}
if self.platform_root:
exports['ZK_AGENT_PLATFORM_ROOT'] = self.platform_root
src/openai_compat.py
+104 -14
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import json
import os
import sys
from typing import Any, Iterator
from urllib import error, request
@@ -31,7 +32,7 @@ ANTHROPIC_MESSAGES_MODEL_PREFIXES = (
)
ANTHROPIC_VERSION = '2023-06-01'
ANTHROPIC_MAX_TOKENS = 4096
ANTHROPIC_MAX_TOKENS = 16384
DEFAULT_MODEL_IDLE_TIMEOUT_SECONDS = 60.0
@@ -218,6 +219,57 @@ def _uses_anthropic_messages_api(model: str, base_url: str) -> bool:
)
_UNKNOWN_EVENT_LOG_LIMIT = 8
_unknown_event_seen: dict[str, int] = {}
def _log_unknown_anthropic_event(kind: str, payload: dict[str, Any]) -> None:
"""Emit a one-line stderr warning the first few times we see an unknown event.
The streaming handler used to silently drop any block/delta type it didn't
recognise. When the upstream proxy started emitting frames we didn't decode,
the model's output budget was burned with no observable effect. This log
surfaces the situation without flooding logs on every turn.
"""
seen = _unknown_event_seen.get(kind, 0)
if seen >= _UNKNOWN_EVENT_LOG_LIMIT:
return
_unknown_event_seen[kind] = seen + 1
try:
snippet = json.dumps(payload, ensure_ascii=False)[:400]
except Exception: # pragma: no cover - defensive
snippet = repr(payload)[:400]
print(
f'[openai_compat] unknown anthropic stream event {kind!r}: {snippet}',
file=sys.stderr,
flush=True,
)
def _mark_last_message_cache_breakpoint(messages: list[dict[str, Any]]) -> None:
"""Attach cache_control to the last content block of the last message.
Anthropic prompt cache requires the breakpoint at the *end* of the cacheable
prefix. The next turn appends a new message → the previous last becomes the
second-to-last with its breakpoint still in the prefix, so everything up to
that point is reused on the next call.
"""
if not messages:
return
last = messages[-1]
content = last.get('content')
if not isinstance(content, list) or not content:
return
last_block = content[-1]
if not isinstance(last_block, dict):
return
new_block = dict(last_block)
new_block['cache_control'] = {'type': 'ephemeral'}
new_content = list(content)
new_content[-1] = new_block
last['content'] = new_content
def _anthropic_base_url(base_url: str) -> str:
base = base_url.rstrip('/')
if base.lower().endswith('/anthropic'):
@@ -486,25 +538,46 @@ class OpenAICompatClient:
if blocks:
self._append_anthropic_message(anthropic_messages, 'user', blocks)
system_text = '\n\n'.join(system_parts) if system_parts else ''
if output_schema is not None:
schema_hint = (
f'请严格输出符合 JSON Schema `{output_schema.name}` 的 JSON，'
'不要输出额外解释。'
)
system_text = f'{system_text}\n\n{schema_hint}'.strip() if system_text else schema_hint
cache_enabled = (
os.environ.get('CLAW_DISABLE_PROMPT_CACHE', '').strip().lower()
not in {'1', 'true', 'yes', 'on'}
)
payload: dict[str, Any] = {
'model': self.config.model,
'messages': anthropic_messages,
'max_tokens': ANTHROPIC_MAX_TOKENS,
'stream': stream,
}
if system_parts:
payload['system'] = '\n\n'.join(system_parts)
if system_text:
if cache_enabled:
payload['system'] = [
{
'type': 'text',
'text': system_text,
'cache_control': {'type': 'ephemeral'},
}
]
else:
payload['system'] = system_text
converted_tools = self._anthropic_tools(tools)
if converted_tools:
if cache_enabled:
converted_tools = list(converted_tools)
last_tool = dict(converted_tools[-1])
last_tool['cache_control'] = {'type': 'ephemeral'}
converted_tools[-1] = last_tool
payload['tools'] = converted_tools
if output_schema is not None:
schema_hint = (
f'请严格输出符合 JSON Schema `{output_schema.name}` 的 JSON，'
'不要输出额外解释。'
)
payload['system'] = (
f'{payload.get("system", "")}\n\n{schema_hint}'.strip()
)
if cache_enabled and anthropic_messages:
_mark_last_message_cache_breakpoint(anthropic_messages)
return payload
def _anthropic_headers(self) -> dict[str, str]:
@@ -604,7 +677,12 @@ class OpenAICompatClient:
content_block = event_payload.get('content_block')
if not isinstance(block_index, int) or not isinstance(content_block, dict):
continue
if content_block.get('type') != 'tool_use':
block_type = content_block.get('type')
if block_type != 'tool_use':
if block_type not in ('text', 'thinking', 'redacted_thinking'):
_log_unknown_anthropic_event(
'content_block_start', event_payload
)
continue
tool_index = next_tool_index
next_tool_index += 1
@@ -636,7 +714,8 @@ class OpenAICompatClient:
delta = event_payload.get('delta')
if not isinstance(delta, dict):
continue
if delta.get('type') == 'text_delta':
delta_type = delta.get('type')
if delta_type == 'text_delta':
text = delta.get('text')
if isinstance(text, str) and text:
yield StreamEvent(
@@ -645,7 +724,7 @@ class OpenAICompatClient:
raw_event=event_payload,
)
continue
if delta.get('type') == 'input_json_delta':
if delta_type == 'input_json_delta':
block_index = event_payload.get('index')
partial_json = delta.get('partial_json')
if (
@@ -661,6 +740,17 @@ class OpenAICompatClient:
raw_event=event_payload,
)
continue
if delta_type in ('thinking_delta', 'signature_delta'):
# Extended-thinking blocks: we don't surface them as
# content (the runtime has nowhere to put them), but
# they previously vanished silently and burned the
# whole 4k output budget before tool_use could start.
# Logging once is enough to confirm they're flowing.
continue
_log_unknown_anthropic_event(
f'content_block_delta:{delta_type}', event_payload
)
continue
if event_type == 'message_delta':
delta = event_payload.get('delta')
if isinstance(delta, dict) and isinstance(delta.get('stop_reason'), str):
src/session_store.py
+1 -1
View File
@@ -223,7 +223,7 @@ def deserialize_runtime_config(payload: JSONDict) -> AgentRuntimeConfig:
return AgentRuntimeConfig(
cwd=Path(str(payload['cwd'])).resolve(),
max_turns=int(payload.get('max_turns', DEFAULT_MAX_TURNS)),
command_timeout_seconds=float(payload.get('command_timeout_seconds', 30.0)),
command_timeout_seconds=float(payload.get('command_timeout_seconds', 300.0)),
max_output_chars=int(payload.get('max_output_chars', 50000)),
stream_model_responses=bool(payload.get('stream_model_responses', False)),
auto_snip_threshold_tokens=_optional_int(payload.get('auto_snip_threshold_tokens')),