@agentunion/kite 1.5.0 → 1.6.0

package/extensions/services/dataclaw/core/config.py

@@ -0,0 +1,50 @@
+"""
+配置加载器：通过 module.md 的 business_config 读取配置
+"""
+import os
+import re
+import json
+from typing import Any, Dict
+
+_config: Dict[str, Any] = {}
+
+
+def load_config(module_config: Dict[str, Any] = None) -> Dict[str, Any]:
+    """从 module_config 中读取 business_config 指定的文件"""
+    if module_config is None:
+        # 兜底：直接读取 config.json5
+        config_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json5")
+    else:
+        config_file = module_config.get("business_config")
+        if not config_file:
+            return {}
+        config_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), config_file)
+
+    try:
+        import pyjson5
+        with open(config_path, encoding="utf-8") as f:
+            cfg = pyjson5.load(f)
+        print(f"[dataclaw] config loaded from {config_path}")
+        return cfg
+    except ImportError:
+        with open(config_path, encoding="utf-8") as f:
+            text = f.read()
+        text = re.sub(r'//.*?$', '', text, flags=re.MULTILINE)
+        text = re.sub(r',\s*([}\]])', r'\1', text)
+        return json.loads(text)
+    except FileNotFoundError:
+        print(f"[dataclaw] WARNING: config not found at {config_path}, using defaults")
+        return {}
+    except Exception as e:
+        print(f"[dataclaw] ERROR: failed to load config: {e}")
+        return {}
+
+
+def get_config() -> Dict[str, Any]:
+    return _config
+
+
+def init_config(module_config: Dict[str, Any] = None) -> Dict[str, Any]:
+    global _config
+    _config = load_config(module_config)
+    return _config

package/extensions/services/dataclaw/core/database.py

@@ -0,0 +1,70 @@
+"""
+MySQL 数据库连接池（aiomysql + SQLAlchemy 2.0 异步）
+"""
+from typing import AsyncGenerator
+from urllib.parse import quote_plus
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
+from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy import text
+
+_engine = None
+_async_session_factory = None
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+def init_db(cfg: dict):
+    global _engine, _async_session_factory
+    host = cfg.get("host", "127.0.0.1")
+    port = cfg.get("port", 3306)
+    user = cfg.get("user", "root")
+    password = quote_plus(cfg.get("password", ""))
+    database = cfg.get("database", "dataclaw")
+    pool_size = cfg.get("pool_size", 10)
+    max_overflow = cfg.get("max_overflow", 5)
+    pool_recycle = cfg.get("pool_recycle", 3600)
+    pool_pre_ping = cfg.get("pool_pre_ping", True)
+
+    url = f"mysql+aiomysql://{user}:{password}@{host}:{port}/{database}?charset=utf8mb4"
+    _engine = create_async_engine(
+        url,
+        pool_size=pool_size,
+        max_overflow=max_overflow,
+        pool_recycle=pool_recycle,
+        pool_pre_ping=pool_pre_ping,
+        echo=False,
+    )
+    _async_session_factory = async_sessionmaker(
+        _engine, class_=AsyncSession, expire_on_commit=False
+    )
+    print(f"[dataclaw] database pool initialized -> {host}:{port}/{database}")
+
+
+async def get_session() -> AsyncGenerator[AsyncSession, None]:
+    async with _async_session_factory() as session:
+        try:
+            yield session
+        except Exception:
+            await session.rollback()
+            raise
+        finally:
+            await session.close()
+
+
+async def check_connection() -> tuple[bool, str]:
+    """检查数据库连通性，返回 (ok, error_message)"""
+    try:
+        async with _async_session_factory() as session:
+            await session.execute(text("SELECT 1"))
+        return True, ""
+    except Exception as e:
+        return False, str(e)
+
+
+async def close_db():
+    global _engine
+    if _engine:
+        await _engine.dispose()
+        print("[dataclaw] database pool closed")

package/extensions/services/dataclaw/entry.py

@@ -0,0 +1,416 @@
+"""
+DataClaw V2.0 — Agent 数据共享平台
+一次采集，万次复用。
+"""
+
+import builtins
+import json
+import os
+import re
+import socket
+import sys
+import threading
+import time
+import traceback
+from datetime import datetime, timezone
+
+import asyncio
+import uvicorn
+
+# 预导入 server 模块（在等待 stdin 之前）
+from server import DataClawServer
+
+# ── Module configuration ──
+
+def _load_module_config() -> dict:
+    _this_dir = os.path.dirname(os.path.abspath(__file__))
+    module_md = os.path.join(_this_dir, "module.md")
+    result = {"name": "", "preferred_port": 0, "advertise_ip": "0.0.0.0"}
+
+    if not os.path.exists(module_md):
+        print(f"[dataclaw] ERROR: module.md not found")
+        sys.exit(1)
+
+    try:
+        with open(module_md, encoding="utf-8") as f:
+            text = f.read()
+        m = re.match(r'^---\s*\n(.*?)\n---', text, re.DOTALL)
+        if not m:
+            print(f"[dataclaw] ERROR: Missing YAML frontmatter in module.md")
+            sys.exit(1)
+        try:
+            import yaml
+            fm = yaml.safe_load(m.group(1)) or {}
+        except ImportError:
+            # fallback: simple key:value parser
+            fm = {}
+            for line in m.group(1).splitlines():
+                if ':' in line and not line.strip().startswith('-'):
+                    k, _, v = line.partition(':')
+                    fm[k.strip()] = v.strip()
+        raw_name = str(fm.get("name", "")).strip()
+        if not raw_name:
+            print("[dataclaw] ERROR: Missing name in module.md")
+            sys.exit(1)
+        result["name"] = re.sub(r'[^a-zA-Z0-9_\-]', '', raw_name)
+        if "preferred_port" in fm:
+            try:
+                result["preferred_port"] = int(fm["preferred_port"])
+            except (ValueError, TypeError):
+                pass
+        if "advertise_ip" in fm:
+            result["advertise_ip"] = str(fm["advertise_ip"])
+        # max_connections（弹性连接上限）
+        try:
+            result["max_connections"] = max(1, min(10, int(fm.get("max_connections", 1))))
+        except (ValueError, TypeError):
+            result["max_connections"] = 1
+        # business_config（从 businesses 数组取第一个 config_file）
+        businesses = fm.get("businesses")
+        if isinstance(businesses, list) and businesses:
+            cfg_file = businesses[0].get("config_file") if isinstance(businesses[0], dict) else None
+            if cfg_file:
+                result["business_config"] = str(cfg_file).strip()
+    except SystemExit:
+        raise
+    except Exception as e:
+        print(f"[dataclaw] ERROR: Failed to read module.md: {e}")
+        sys.exit(1)
+
+    return result
+
+_module_config = _load_module_config()
+MODULE_NAME = _module_config["name"]
+
+# ── Safe stdout/stderr ──
+
+class _SafeWriter:
+    def __init__(self, stream):
+        self._stream = stream
+    def write(self, s):
+        try: self._stream.write(s)
+        except (BrokenPipeError, OSError): pass
+    def flush(self):
+        try: self._stream.flush()
+        except (BrokenPipeError, OSError): pass
+    def __getattr__(self, name):
+        return getattr(self._stream, name)
+
+sys.stdout = _SafeWriter(sys.stdout)
+sys.stderr = _SafeWriter(sys.stderr)
+
+# ── Log system ──
+
+_builtin_print = builtins.print
+_start_ts = time.monotonic()
+_last_ts = time.monotonic()
+_ANSI_RE = re.compile(r"\033\[[0-9;]*m")
+_log_lock = threading.Lock()
+_log_latest_path = None
+_log_daily_path = None
+_log_daily_date = ""
+_log_dir = None
+_crash_log_path = None
+
+def _strip_ansi(s: str) -> str:
+    return _ANSI_RE.sub("", s)
+
+def _resolve_daily_log_path():
+    global _log_daily_path, _log_daily_date
+    if not _log_dir:
+        return
+    today = datetime.now().strftime("%Y-%m-%d")
+    if today == _log_daily_date and _log_daily_path:
+        return
+    month_dir = os.path.join(_log_dir, today[:7])
+    os.makedirs(month_dir, exist_ok=True)
+    _log_daily_path = os.path.join(month_dir, f"{today}.log")
+    _log_daily_date = today
+
+def _write_log(plain_line: str):
+    with _log_lock:
+        if _log_latest_path:
+            try:
+                with open(_log_latest_path, "a", encoding="utf-8") as f:
+                    f.write(plain_line)
+            except Exception:
+                pass
+        _resolve_daily_log_path()
+        if _log_daily_path:
+            try:
+                with open(_log_daily_path, "a", encoding="utf-8") as f:
+                    f.write(plain_line)
+            except Exception:
+                pass
+
+def _write_crash(exc_type, exc_value, exc_tb, thread_name=None, severity="critical", handled=False):
+    record = {
+        "timestamp": datetime.now(timezone.utc).isoformat(),
+        "module": MODULE_NAME,
+        "thread": thread_name or threading.current_thread().name,
+        "exception_type": exc_type.__name__ if exc_type else "Unknown",
+        "exception_message": str(exc_value),
+        "traceback": "".join(traceback.format_exception(exc_type, exc_value, exc_tb)),
+        "severity": severity,
+        "handled": handled,
+    }
+    line = json.dumps(record, ensure_ascii=False) + "\n"
+    if _crash_log_path:
+        try:
+            with open(_crash_log_path, "a", encoding="utf-8") as f:
+                f.write(line)
+        except Exception:
+            pass
+
+def _setup_exception_hooks():
+    _orig = sys.excepthook
+    def _hook(t, v, tb):
+        _write_crash(t, v, tb, severity="critical", handled=False)
+        _orig(t, v, tb)
+    sys.excepthook = _hook
+
+def _tprint(*args, **kwargs):
+    global _last_ts
+    now = time.monotonic()
+    elapsed = now - _start_ts
+    delta = now - _last_ts
+    _last_ts = now
+    elapsed_str = f"{elapsed*1000:.0f}ms" if elapsed < 1 else (f"{elapsed:.1f}s" if elapsed < 100 else f"{elapsed:.0f}s")
+    delta_str = "" if delta < 0.001 else (f"+{delta*1000:.0f}ms" if delta < 1 else f"+{delta:.1f}s")
+    ts = datetime.now().strftime("%H:%M:%S.%f")[:-3]
+    _builtin_print(*args, **kwargs)
+    if _log_latest_path or _log_daily_path:
+        sep = kwargs.get("sep", " ")
+        end = kwargs.get("end", "\n")
+        raw = sep.join(str(a) for a in args) + end
+        plain = _strip_ansi(raw)
+        timing = f"[{ts}|{elapsed_str}{('|' + delta_str) if delta_str else ''}] "
+        _write_log(timing + plain)
+
+builtins.print = _tprint
+
+# ── Port binding ──
+
+def _bind_port(preferred: int) -> int:
+    for port in range(preferred, preferred + 100):
+        try:
+            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            s.bind(('0.0.0.0', port))
+            s.close()
+            return port
+        except OSError:
+            continue
+    raise RuntimeError(f"No free port found in range {preferred}\u2013{preferred+99}")
+
+# ── Entry point ──
+
+def _fmt_elapsed(t0):
+    ms = int((time.monotonic() - t0) * 1000)
+    return f"{ms}ms" if ms < 1000 else f"{ms/1000:.1f}s"
+
+
+def _run_local_mode(kernel_port: int):
+    global _log_dir, _log_latest_path, _crash_log_path
+
+    # 1. Init log directory
+    data_dir = os.environ.get("KITE_MODULE_DATA", "")
+    if data_dir:
+        _log_dir = os.path.join(data_dir, "log")
+        os.makedirs(_log_dir, exist_ok=True)
+        _log_latest_path = os.path.join(_log_dir, "latest.log")
+        _crash_log_path = os.path.join(_log_dir, "crashes.jsonl")
+        try:
+            open(_log_latest_path, "w", encoding="utf-8").close()
+        except Exception:
+            pass
+
+    _setup_exception_hooks()
+
+    _t0 = time.monotonic()
+
+    if not kernel_port:
+        print(f"[{MODULE_NAME}] WARNING: Missing KITE_KERNEL_PORT, running without Kernel")
+    else:
+        print(f"[{MODULE_NAME}] Kernel port: {kernel_port} ({_fmt_elapsed(_t0)})")
+
+    # 3. Bind port
+    host = _module_config.get("advertise_ip", "0.0.0.0")
+    preferred = _module_config.get("preferred_port") or 20010
+    port = _bind_port(preferred)
+    if port is None:
+        print(f"[{MODULE_NAME}] ERROR: Cannot bind to any port, exiting")
+        sys.exit(1)
+    print(f"[{MODULE_NAME}] bound to port {port} ({_fmt_elapsed(_t0)})")
+
+    # 4. Read token from stdin (JSON protocol: {"token": "xxx"})
+    line = sys.stdin.readline().strip()
+    if not line:
+        print(f"\033[31m[{MODULE_NAME}] ERROR: stdin closed\033[0m")
+        sys.exit(1)
+
+    try:
+        msg = json.loads(line)
+    except json.JSONDecodeError as e:
+        print(f"\033[31m[{MODULE_NAME}] ERROR: Invalid JSON from stdin: {e}\033[0m")
+        sys.exit(1)
+
+    if "error" in msg:
+        print(f"\033[31m[{MODULE_NAME}] 启动失败: {msg.get('message', 'unknown')}\033[0m")
+        sys.exit(1)
+
+    token = msg.get("token", "")
+    if not token:
+        print(f"\033[31m[{MODULE_NAME}] ERROR: No token in stdin message\033[0m")
+        sys.exit(1)
+
+    # 5. Create server instance
+
+    server = DataClawServer(
+        token=token,
+        kernel_port=kernel_port,
+        advertise_ip=host,
+        port=port,
+        data_dir=data_dir,
+        max_connections=_module_config.get("max_connections", 1),
+        module_config=_module_config,
+    )
+
+    app = server.build_app()
+
+    # 6. Start with uvicorn
+    display_host = "localhost" if host == "0.0.0.0" else host
+    url = f"http://{display_host}:{port}"
+    print(f"[{MODULE_NAME}] Starting on {host}:{port} ({_fmt_elapsed(_t0)})")
+    print(f"[{MODULE_NAME}] \033[32m✓ DataClaw ready: {url}\033[0m")
+
+    try:
+        config = uvicorn.Config(
+            app,
+            host=host,
+            port=port,
+            log_level="warning",
+            access_log=False,
+        )
+        uv_server = uvicorn.Server(config)
+        server._uv_server = uv_server
+        uv_server.run()
+    except Exception as e:
+        _write_crash(type(e), e, e.__traceback__, severity="critical", handled=True)
+        print(f"\033[31m[{MODULE_NAME}] ERROR: {e}\033[0m")
+        sys.exit(1)
+
+    if server._exit_code != 0:
+        sys.exit(server._exit_code)
+
+
+def _run_remote_mode(gateway_url: str):
+    global _log_dir, _log_latest_path, _crash_log_path
+    home = os.environ.get("HOME") or os.environ.get("USERPROFILE") or os.path.expanduser("~")
+    data_dir = os.path.join(home, ".kite", "remote", MODULE_NAME)
+    os.makedirs(data_dir, exist_ok=True)
+    _log_dir = os.path.join(data_dir, "log")
+    os.makedirs(_log_dir, exist_ok=True)
+    _log_latest_path = os.path.join(_log_dir, "latest.log")
+    try:
+        with open(_log_latest_path, "w", encoding="utf-8") as f: pass
+    except Exception: _log_latest_path = None
+    _crash_log_path = os.path.join(_log_dir, "crashes.jsonl")
+    try:
+        with open(_crash_log_path, "w", encoding="utf-8") as f: pass
+    except Exception: _crash_log_path = None
+    _resolve_daily_log_path()
+    _setup_exception_hooks()
+
+    _t0 = time.monotonic()
+    kite_token = _get_kite_token(MODULE_NAME, gateway_url)
+    print(f"[{MODULE_NAME}] Remote mode: gateway={gateway_url}")
+
+    host = _module_config.get("advertise_ip", "0.0.0.0")
+    preferred = _module_config.get("preferred_port") or 20010
+    port = _bind_port(preferred)
+    if port is None:
+        print(f"[{MODULE_NAME}] ERROR: Cannot bind to any port")
+        sys.exit(1)
+
+    server = DataClawServer(
+        token="",
+        kernel_port=0,
+        advertise_ip=host,
+        port=port,
+        data_dir=data_dir,
+        max_connections=_module_config.get("max_connections", 1),
+        gateway_url=gateway_url,
+        kite_token=kite_token,
+    )
+
+    app = server.build_app()
+
+    try:
+        config = uvicorn.Config(
+            app,
+            host=host,
+            port=port,
+            log_level="warning",
+            access_log=False,
+        )
+        uv_server = uvicorn.Server(config)
+        server._uv_server = uv_server
+        uv_server.run()
+    except Exception as e:
+        _write_crash(type(e), e, e.__traceback__, severity="critical", handled=True)
+        print(f"\033[31m[{MODULE_NAME}] ERROR: {e}\033[0m")
+        sys.exit(1)
+    if hasattr(server, '_auth_failed') and server._auth_failed:
+        _clear_token_cache(MODULE_NAME, gateway_url)
+    if server._exit_code != 0:
+        sys.exit(server._exit_code)
+
+
+# ── Token management ──
+
+def _gateway_to_filename(gateway_url: str) -> str:
+    try:
+        from urllib.parse import urlparse
+        parsed = urlparse(gateway_url)
+        host = parsed.hostname or "unknown"
+        port = parsed.port or (443 if parsed.scheme == "wss" else 80)
+        return f"{host}-{port}.json".replace(":", "-").replace("/", "-")
+    except Exception: return "default.json"
+
+def _load_token_cache(module_name: str, gateway_url: str) -> dict | None:
+    home = os.environ.get("HOME") or os.environ.get("USERPROFILE") or os.path.expanduser("~")
+    token_file = os.path.join(home, ".kite", "remote", module_name, "tokens", _gateway_to_filename(gateway_url))
+    if not os.path.exists(token_file): return None
+    try:
+        with open(token_file, "r") as f: return json.load(f)
+    except Exception: return None
+
+def _clear_token_cache(module_name: str, gateway_url: str):
+    home = os.environ.get("HOME") or os.environ.get("USERPROFILE") or os.path.expanduser("~")
+    token_file = os.path.join(home, ".kite", "remote", module_name, "tokens", _gateway_to_filename(gateway_url))
+    try: os.remove(token_file)
+    except Exception: pass
+
+def _get_kite_token(module_name: str, gateway_url: str) -> str:
+    token = os.environ.get("KITE_TOKEN")
+    if token: return token
+    cache = _load_token_cache(module_name, gateway_url)
+    if cache: return cache["token"]
+    print(f"[{module_name}] No token for {gateway_url}")
+    print(f"  export KITE_TOKEN=<token>")
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+    kernel_port = os.environ.get("KITE_KERNEL_PORT")
+    if kernel_port:
+        _run_local_mode(int(kernel_port))
+    else:
+        config = _load_module_config()
+        gateway_url = config.get("gateway_url") or os.environ.get("KITE_GATEWAY_URL")
+        if gateway_url:
+            _run_remote_mode(gateway_url)
+        else:
+            print(f"[{MODULE_NAME}] ERROR: No KITE_KERNEL_PORT and no gateway_url")
+            sys.exit(1)