439 lines
18 KiB
Python
439 lines
18 KiB
Python
"""Filesystem bridge + tmux injection for the kimi-native terminal harness.
|
|
|
|
The runner launches the ``kimi`` TUI in a private tmux pane and records
|
|
that pane's socket + target here via :func:`write_tmux_target`. The harness
|
|
executor then delivers Omnigent web-UI messages into the *same* pane via
|
|
:func:`inject_user_message` (tmux bracketed paste + Enter) — the kimi analog
|
|
of claude-native's tmux send-keys bridge. This is what wires the web-UI chat box
|
|
to the running Kimi TUI (and, since the web UI embeds that pane, the message
|
|
shows in both surfaces).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
#: Env var carrying the bridge dir into the harness executor process.
|
|
BRIDGE_DIR_ENV_VAR = "HARNESS_KIMI_NATIVE_BRIDGE_DIR"
|
|
|
|
_BRIDGE_ROOT = Path(os.environ.get("TMPDIR", "/tmp")) / f"omnigent-{os.getuid()}" / "kimi-native"
|
|
_TMUX_FILE = "tmux.json"
|
|
# Omnigent routing details the kimi hook subprocess reads to reach the server.
|
|
# Mirrors claude-native's ``permission_hook.json`` (server URL + auth headers +
|
|
# the active Omnigent session). Written by the runner at terminal-create time;
|
|
# read by :mod:`omnigent.kimi_native_hook` (PreToolUse deny-gate + the
|
|
# PermissionRequest read-only surface).
|
|
_HOOK_CONFIG_FILE = "hook_config.json"
|
|
_TMUX_READY_TIMEOUT_S = 30.0
|
|
_TMUX_SEND_TIMEOUT_S = 5.0
|
|
_POLL_INTERVAL_S = 0.15
|
|
_PASTE_SETTLE_S = 0.1 # let the TUI commit a paste before the separate submit Enter
|
|
_PASTE_BUFFER = "omnigent-kimi-paste"
|
|
# How long to wait for the pasted text to become visible in the pane before
|
|
# sending Enter — submitting before the TUI commits the paste folds the Enter
|
|
# into the paste as a newline and the message sits unsent.
|
|
_PASTE_COMMIT_TIMEOUT_S = 5.0
|
|
# kimi TUI readiness markers — substrings that appear once the TUI has mounted
|
|
# its input box, gating the pre-paste wait in ``_settle_pane``. The footer
|
|
# context-window indicator (``context: 6.5% (17.0k/262.1k)``) is verified
|
|
# present in EVERY mounted state of a live K2.7 session — idle, thinking, and
|
|
# while an approval menu is up — so the gate returns on the first capture
|
|
# instead of blocking the full readiness timeout before each paste. (The old
|
|
# ``"Plan, search, build"`` / ``"Add a follow-up"`` strings were carried over
|
|
# from cursor-native unverified and never matched, so every injection ate the
|
|
# whole 30s timeout — the web→TUI latency the markers were meant to avoid.)
|
|
_INPUT_READY_MARKERS = ("context:",)
|
|
_TRUST_MARKER = "Trust this workspace"
|
|
|
|
|
|
def bridge_dir_for_session_id(session_id: str) -> Path:
|
|
"""Return the per-session bridge dir, e.g. ``/tmp/omnigent-<uid>/kimi-native/<hash>``."""
|
|
digest = hashlib.sha256(session_id.encode("utf-8")).hexdigest()[:32]
|
|
return _BRIDGE_ROOT / digest
|
|
|
|
|
|
def bridge_root() -> Path:
|
|
"""Return the configured Kimi-native bridge root."""
|
|
return _BRIDGE_ROOT
|
|
|
|
|
|
def _ensure_dir(path: Path) -> None:
|
|
"""Create *path* (and parents) with owner-only permissions."""
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
with contextlib.suppress(OSError):
|
|
os.chmod(path, 0o700)
|
|
|
|
|
|
def build_kimi_native_spawn_env(session_id: str) -> dict[str, str]:
|
|
"""Build the ``HARNESS_KIMI_NATIVE_*`` env the harness executor reads."""
|
|
bridge_dir = bridge_dir_for_session_id(session_id)
|
|
_ensure_dir(bridge_dir)
|
|
return {
|
|
BRIDGE_DIR_ENV_VAR: str(bridge_dir),
|
|
}
|
|
|
|
|
|
def write_hook_config(
|
|
bridge_dir: Path,
|
|
*,
|
|
server_url: str,
|
|
headers: dict[str, str],
|
|
session_id: str,
|
|
) -> None:
|
|
"""Record the Omnigent routing details the kimi hook subprocess reads.
|
|
|
|
The PreToolUse / PermissionRequest hook commands receive only
|
|
``--bridge-dir`` on their command line (no secrets); they read the
|
|
server URL, auth headers, and active session id from this file. Mirrors
|
|
:func:`omnigent.claude_native_bridge` ``permission_hook.json`` plumbing.
|
|
|
|
:param bridge_dir: The kimi-native bridge dir.
|
|
:param server_url: Omnigent server base URL, e.g. ``"http://127.0.0.1:8787"``.
|
|
:param headers: Auth headers to replay on the hook's POSTs (may be empty).
|
|
:param session_id: The Omnigent session the hook events belong to.
|
|
"""
|
|
_ensure_dir(bridge_dir)
|
|
payload = {
|
|
"ap_server_url": server_url,
|
|
"ap_auth_headers": dict(headers),
|
|
"session_id": session_id,
|
|
"updated_at": time.time(),
|
|
}
|
|
tmp = bridge_dir / (_HOOK_CONFIG_FILE + ".tmp")
|
|
tmp.write_text(json.dumps(payload), encoding="utf-8")
|
|
with contextlib.suppress(OSError):
|
|
os.chmod(tmp, 0o600)
|
|
os.replace(tmp, bridge_dir / _HOOK_CONFIG_FILE)
|
|
|
|
|
|
def read_hook_config(bridge_dir: Path) -> dict[str, Any]:
|
|
"""Read Omnigent routing details for the kimi hook subprocess.
|
|
|
|
:param bridge_dir: The kimi-native bridge dir.
|
|
:returns: ``{"ap_server_url", "ap_auth_headers", "session_id"}`` (or an
|
|
empty dict when the file is absent or malformed).
|
|
"""
|
|
try:
|
|
raw = (bridge_dir / _HOOK_CONFIG_FILE).read_text(encoding="utf-8")
|
|
except OSError:
|
|
return {}
|
|
try:
|
|
data = json.loads(raw)
|
|
except ValueError:
|
|
return {}
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
|
|
def read_active_session_id(bridge_dir: Path) -> str | None:
|
|
"""Return the Omnigent session id recorded for the hook subprocess.
|
|
|
|
:param bridge_dir: The kimi-native bridge dir.
|
|
:returns: The session id, or ``None`` when unset / malformed.
|
|
"""
|
|
session_id = read_hook_config(bridge_dir).get("session_id")
|
|
return session_id if isinstance(session_id, str) and session_id else None
|
|
|
|
|
|
def write_tmux_target(
|
|
bridge_dir: Path,
|
|
*,
|
|
socket_path: Path,
|
|
tmux_target: str,
|
|
pid: int | None = None,
|
|
) -> None:
|
|
"""Advertise the tmux socket + target for the running Kimi terminal."""
|
|
_ensure_dir(bridge_dir)
|
|
payload: dict[str, Any] = {
|
|
"socket_path": str(socket_path),
|
|
"tmux_target": tmux_target,
|
|
"updated_at": time.time(),
|
|
}
|
|
if pid is not None:
|
|
payload["pid"] = pid
|
|
tmp = bridge_dir / (_TMUX_FILE + ".tmp")
|
|
tmp.write_text(json.dumps(payload), encoding="utf-8")
|
|
os.replace(tmp, bridge_dir / _TMUX_FILE)
|
|
|
|
|
|
def read_tmux_info(bridge_dir: Path) -> dict[str, str] | None:
|
|
"""Return ``{socket_path, tmux_target}`` from ``tmux.json``, or ``None``."""
|
|
try:
|
|
raw = (bridge_dir / _TMUX_FILE).read_text(encoding="utf-8")
|
|
except OSError:
|
|
return None
|
|
try:
|
|
data = json.loads(raw)
|
|
except ValueError:
|
|
return None
|
|
socket_path = data.get("socket_path")
|
|
tmux_target = data.get("tmux_target")
|
|
if (
|
|
isinstance(socket_path, str)
|
|
and socket_path
|
|
and isinstance(tmux_target, str)
|
|
and tmux_target
|
|
):
|
|
return {"socket_path": socket_path, "tmux_target": tmux_target}
|
|
return None
|
|
|
|
|
|
def _wait_for_tmux_info(bridge_dir: Path, *, timeout_s: float) -> dict[str, str]:
|
|
"""Block until ``tmux.json`` is advertised, or raise on timeout."""
|
|
deadline = time.monotonic() + timeout_s
|
|
while time.monotonic() < deadline:
|
|
info = read_tmux_info(bridge_dir)
|
|
if info is not None:
|
|
return info
|
|
time.sleep(_POLL_INTERVAL_S)
|
|
raise RuntimeError(f"kimi-native tmux target was not advertised within {timeout_s:.0f}s")
|
|
|
|
|
|
def _run_tmux(socket_path: str, *args: str) -> None:
|
|
"""Invoke ``tmux -S <socket> <args...>`` and raise on failure."""
|
|
try:
|
|
proc = subprocess.run(
|
|
["tmux", "-S", socket_path, *args],
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=_TMUX_SEND_TIMEOUT_S,
|
|
)
|
|
except subprocess.TimeoutExpired as exc:
|
|
raise RuntimeError(f"tmux command timed out after {_TMUX_SEND_TIMEOUT_S}s") from exc
|
|
if proc.returncode != 0:
|
|
detail = proc.stderr.strip() or proc.stdout.strip() or "<no output>"
|
|
raise RuntimeError(f"tmux command failed (rc={proc.returncode}): {detail}")
|
|
|
|
|
|
def _capture_pane(socket_path: str, tmux_target: str) -> str:
|
|
"""Capture the visible pane contents; ``""`` on any failure (treat as not-ready)."""
|
|
try:
|
|
proc = subprocess.run(
|
|
["tmux", "-S", socket_path, "capture-pane", "-p", "-t", tmux_target],
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=_TMUX_SEND_TIMEOUT_S,
|
|
)
|
|
except (subprocess.TimeoutExpired, OSError):
|
|
return ""
|
|
return proc.stdout if proc.returncode == 0 else ""
|
|
|
|
|
|
def _paste_payload_bytes(text: str) -> bytes:
|
|
r"""Encode text for ``tmux load-buffer``: line breaks → CR, tabs kept, other
|
|
control bytes dropped (a stray ESC would close the bracketed-paste early)."""
|
|
normalized = text.replace("\r\n", "\n").replace("\r", "\n")
|
|
body = bytearray()
|
|
for ch in normalized:
|
|
if ch == "\n":
|
|
body.append(0x0D)
|
|
continue
|
|
if ch == "\t":
|
|
body.append(0x09)
|
|
continue
|
|
if ord(ch) < 0x20:
|
|
continue
|
|
body.extend(ch.encode("utf-8"))
|
|
return bytes(body)
|
|
|
|
|
|
def _session_alive(socket_path: str, tmux_target: str) -> bool:
|
|
"""Return whether the tmux session/pane still exists (the TUI is running)."""
|
|
try:
|
|
proc = subprocess.run(
|
|
["tmux", "-S", socket_path, "has-session", "-t", tmux_target],
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=_TMUX_SEND_TIMEOUT_S,
|
|
)
|
|
except (subprocess.TimeoutExpired, OSError):
|
|
return False
|
|
return proc.returncode == 0
|
|
|
|
|
|
def _submit_needle(content: str) -> str:
|
|
"""A stable single-line substring used to confirm the paste rendered in the pane."""
|
|
for line in content.splitlines():
|
|
stripped = line.strip()
|
|
if len(stripped) >= 4:
|
|
return stripped[:24]
|
|
stripped = content.strip()
|
|
return stripped[:24] if len(stripped) >= 4 else ""
|
|
|
|
|
|
def _settle_pane(socket_path: str, tmux_target: str, *, timeout_s: float) -> None:
|
|
"""Best-effort wait until the Kimi input box is ready to receive a paste.
|
|
|
|
Accepts the first-run "Trust this workspace" modal (sends ``a`` at most once)
|
|
so the input box can mount, then returns as soon as the TUI chrome is present
|
|
(see :data:`_INPUT_READY_MARKERS`). Falls through after the timeout (e.g. a
|
|
boot that never renders) rather than raising — the paste still lands.
|
|
"""
|
|
deadline = time.monotonic() + timeout_s
|
|
trust_accepted = False
|
|
while time.monotonic() < deadline:
|
|
pane = _capture_pane(socket_path, tmux_target)
|
|
if any(marker in pane for marker in _INPUT_READY_MARKERS):
|
|
return
|
|
# One-shot, only when no input marker is up (so a later transcript that
|
|
# merely echoes the phrase can't spray repeated keystrokes into the TUI).
|
|
if not trust_accepted and _TRUST_MARKER in pane:
|
|
trust_accepted = True
|
|
with contextlib.suppress(RuntimeError):
|
|
_run_tmux(socket_path, "send-keys", "-t", tmux_target, "a")
|
|
time.sleep(_POLL_INTERVAL_S)
|
|
|
|
|
|
def inject_user_message(
|
|
bridge_dir: Path,
|
|
*,
|
|
content: str,
|
|
timeout_s: float = _TMUX_READY_TIMEOUT_S,
|
|
) -> None:
|
|
"""Deliver a web-UI user message into the Kimi TUI via a tmux bracketed paste.
|
|
|
|
Clears any leftover draft, pastes *content* (multi-line safe via
|
|
``load-buffer``/``paste-buffer -p`` so interior newlines stay data, not
|
|
submits), settles, then submits with Enter.
|
|
|
|
:param bridge_dir: The kimi-native bridge dir holding ``tmux.json``.
|
|
:param content: User text (non-empty).
|
|
:param timeout_s: Per-readiness-gate timeout.
|
|
:raises RuntimeError: If the tmux target is never advertised or a tmux
|
|
command fails.
|
|
"""
|
|
if not content:
|
|
raise RuntimeError("kimi-native injection requires non-empty content")
|
|
info = _wait_for_tmux_info(bridge_dir, timeout_s=timeout_s)
|
|
socket_path = info["socket_path"]
|
|
tmux_target = info["tmux_target"]
|
|
# Fast-fail if the TUI already exited: otherwise _settle_pane polls a dead
|
|
# pane for the full timeout and the web message is silently lost. A clear
|
|
# error lets run_turn surface ExecutorError so the UI can say "restart".
|
|
if not _session_alive(socket_path, tmux_target):
|
|
raise RuntimeError(
|
|
"kimi terminal is no longer running (the TUI exited); restart the session"
|
|
)
|
|
_settle_pane(socket_path, tmux_target, timeout_s=timeout_s)
|
|
# Clear any leftover draft: Home (C-a) + kill-to-end (C-k).
|
|
_run_tmux(socket_path, "send-keys", "-t", tmux_target, "C-a")
|
|
_run_tmux(socket_path, "send-keys", "-t", tmux_target, "C-k")
|
|
with tempfile.NamedTemporaryFile(
|
|
dir=bridge_dir, prefix="paste_", suffix=".bin", delete=False
|
|
) as paste_file:
|
|
# Trailing newline absorbs any trailing backslash so it can't escape Enter.
|
|
paste_file.write(_paste_payload_bytes(content + "\n"))
|
|
paste_path = paste_file.name
|
|
try:
|
|
_run_tmux(socket_path, "load-buffer", "-b", _PASTE_BUFFER, paste_path)
|
|
_run_tmux(
|
|
socket_path,
|
|
"paste-buffer",
|
|
"-p", # bracketed-paste markers — the TUI keeps newlines as data
|
|
"-d", # drop the buffer after pasting
|
|
"-b",
|
|
_PASTE_BUFFER,
|
|
"-t",
|
|
tmux_target,
|
|
)
|
|
finally:
|
|
with contextlib.suppress(OSError):
|
|
os.unlink(paste_path)
|
|
# Wait until the paste is visibly committed to the input box before Enter.
|
|
# Submitting mid-paste folds the Enter in as a newline (the kimi TUI
|
|
# coalesces rapid stdin bursts), leaving the message unsent. Poll for the
|
|
# text, then submit; fall through to a blind submit if no needle is usable.
|
|
needle = _submit_needle(content)
|
|
if needle:
|
|
deadline = time.monotonic() + _PASTE_COMMIT_TIMEOUT_S
|
|
while time.monotonic() < deadline:
|
|
if needle in _capture_pane(socket_path, tmux_target):
|
|
break
|
|
time.sleep(_POLL_INTERVAL_S)
|
|
time.sleep(_PASTE_SETTLE_S)
|
|
_run_tmux(socket_path, "send-keys", "-t", tmux_target, "Enter")
|
|
|
|
|
|
def inject_interrupt(bridge_dir: Path, *, timeout_s: float = _TMUX_READY_TIMEOUT_S) -> None:
|
|
"""Cancel the in-flight Kimi turn by sending ``Escape`` to the pane.
|
|
|
|
kimi stops a running turn on a single ``Escape`` (verified live).
|
|
The harness ``run_turn`` returns right after the paste, so the runner's
|
|
in-process cancel floor can't reach the turn — this is the analog of
|
|
:func:`inject_user_message` for the web UI's Stop button.
|
|
|
|
:raises RuntimeError: If the tmux target is not advertised or send-keys fails.
|
|
"""
|
|
info = _wait_for_tmux_info(bridge_dir, timeout_s=timeout_s)
|
|
# No ``-l``: tmux must interpret ``Escape`` as a key name.
|
|
_run_tmux(info["socket_path"], "send-keys", "-t", info["tmux_target"], "Escape")
|
|
|
|
|
|
#: Tool-independent label proving kimi's permission menu is on screen — guards
|
|
#: against injecting a stray digit once the prompt was already answered.
|
|
_PERMISSION_PROMPT_MARKER = "Approve once"
|
|
|
|
#: Web-UI approve/deny → option digit in kimi's fixed numbered menu
|
|
#: (1=Approve once, 2=Approve for session, 3=Reject, 4=Reject with feedback).
|
|
#: "Approve once" re-prompts each call so Omnigent governs every one.
|
|
APPROVE_KEY = "1"
|
|
DENY_KEY = "3"
|
|
|
|
|
|
def inject_approval_keystroke(
|
|
bridge_dir: Path, *, key: str, timeout_s: float = _TMUX_READY_TIMEOUT_S
|
|
) -> bool:
|
|
"""Answer kimi's tool-permission menu by typing an option digit + Enter.
|
|
|
|
kimi's permission prompt is a numbered select whose footer documents
|
|
``1/2/3/4 choose · ↵ confirm``; the web-UI Approve/Deny buttons map to
|
|
:data:`APPROVE_KEY` / :data:`DENY_KEY`. This types *key* then ``Enter``.
|
|
|
|
Captures the pane first and injects ONLY when the permission menu is
|
|
actually showing (:data:`_PERMISSION_PROMPT_MARKER`), so a web verdict that
|
|
lands after the user already answered in the terminal (or after the prompt
|
|
closed) is a no-op rather than a stray keystroke leaking into whatever is on
|
|
screen next.
|
|
|
|
:param bridge_dir: The kimi-native bridge dir holding ``tmux.json``.
|
|
:param key: The option digit to select (e.g. :data:`APPROVE_KEY`).
|
|
:returns: ``True`` if the keystroke was injected; ``False`` if the
|
|
permission menu was not present (already answered / closed / TUI gone).
|
|
:raises RuntimeError: If the tmux target is not advertised or send-keys fails.
|
|
"""
|
|
info = _wait_for_tmux_info(bridge_dir, timeout_s=timeout_s)
|
|
socket_path = info["socket_path"]
|
|
tmux_target = info["tmux_target"]
|
|
if not _session_alive(socket_path, tmux_target):
|
|
return False
|
|
if _PERMISSION_PROMPT_MARKER not in _capture_pane(socket_path, tmux_target):
|
|
return False
|
|
# ``key`` is a single documented option digit; Enter confirms (the footer
|
|
# lists "choose" and "confirm" separately, so a digit selects and ↵ commits).
|
|
_run_tmux(socket_path, "send-keys", "-t", tmux_target, key)
|
|
_run_tmux(socket_path, "send-keys", "-t", tmux_target, "Enter")
|
|
return True
|
|
|
|
|
|
def kill_session(bridge_dir: Path, *, timeout_s: float = _TMUX_READY_TIMEOUT_S) -> None:
|
|
"""Hard-stop the Kimi session by killing its tmux session.
|
|
|
|
Terminates ``kimi`` and the pane outright — the analog of the
|
|
user manually exiting the attached TUI, for the web UI's "Stop session"
|
|
affordance. Mirrors :func:`omnigent.claude_native_bridge.kill_session`.
|
|
|
|
:raises RuntimeError: If the tmux target is not advertised or kill-session fails.
|
|
"""
|
|
info = _wait_for_tmux_info(bridge_dir, timeout_s=timeout_s)
|
|
_run_tmux(info["socket_path"], "kill-session", "-t", info["tmux_target"])
|