Files
omnigent-ai--omnigent/omnigent/kiro_native_session_forwarder.py
2026-07-13 13:12:00 +08:00

678 lines
26 KiB
Python

"""Structured session forwarder for the kiro-native harness.
Kiro CLI persists chat turns under ``~/.kiro/sessions/cli`` as session metadata
plus JSONL message records. The native Kiro terminal path injects web prompts
into the TUI; this forwarder mirrors Kiro's persisted assistant messages back
into the Omnigent conversation with ``external_conversation_item`` events.
"""
from __future__ import annotations
import asyncio
import contextlib
import json
import logging
import os
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
import httpx
from omnigent.kiro_native_bridge import write_forwarder_ready
_logger = logging.getLogger(__name__)
_DEFAULT_POLL_INTERVAL_S = 0.7
_POST_TIMEOUT_S = 30.0
_DISCOVERY_SKEW_MS = 10_000
_STATE_FILE = "kiro_session_forwarder.json"
_SUPERVISOR_INITIAL_BACKOFF_S = 1.0
_SUPERVISOR_MAX_BACKOFF_S = 30.0
_SUPERVISOR_HEALTHY_UPTIME_S = 60.0
@dataclass
class _ForwardState:
"""Durable cursor for one Kiro JSONL session file."""
session_id: str | None = None
byte_offset: int = 0
@dataclass(frozen=True)
class _KiroConversationMessage:
"""One conversation message parsed from Kiro's JSONL store."""
message_id: str
role: str
text: str
def _kiro_cli_sessions_dir(home: Path | None = None) -> Path:
"""Return Kiro CLI's session directory for this user."""
return (home or Path.home()) / ".kiro" / "sessions" / "cli"
def _read_state(bridge_dir: Path) -> _ForwardState:
"""Load the persisted forward cursor, or a cold default."""
try:
data = json.loads((bridge_dir / _STATE_FILE).read_text(encoding="utf-8"))
except (OSError, ValueError):
return _ForwardState()
session_id = data.get("session_id")
byte_offset = data.get("byte_offset")
return _ForwardState(
session_id=session_id if isinstance(session_id, str) and session_id else None,
byte_offset=byte_offset if isinstance(byte_offset, int) and byte_offset >= 0 else 0,
)
def _write_state(bridge_dir: Path, state: _ForwardState) -> None:
"""Persist the forward cursor atomically."""
bridge_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
with contextlib.suppress(OSError):
os.chmod(bridge_dir, 0o700)
tmp = bridge_dir / (_STATE_FILE + ".tmp")
tmp.write_text(
json.dumps({"session_id": state.session_id, "byte_offset": state.byte_offset}),
encoding="utf-8",
)
os.replace(tmp, bridge_dir / _STATE_FILE)
def _parse_iso_epoch_ms(value: object) -> int:
"""Parse Kiro's ISO timestamp string into epoch milliseconds."""
if not isinstance(value, str) or not value:
return 0
try:
parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
except ValueError:
return 0
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=UTC)
return int(parsed.timestamp() * 1000)
def _same_workspace(left: object, right: str) -> bool:
"""Return whether Kiro metadata cwd matches the runner workspace."""
if not isinstance(left, str) or not left:
return False
try:
return Path(left).expanduser().resolve() == Path(right).expanduser().resolve()
except OSError:
return left == right
# Competing-candidate sets already warned about, so the ambiguous branch logs
# once per distinct collision instead of every poll tick for the session's life.
# ponytail: grows by one entry per distinct concurrent collision — too rare to
# bound; add a cap only if a pathological churn ever shows up.
_ambiguous_discovery_warned: set[frozenset[str]] = set()
def _warn_ambiguous_discovery(workspace: str, session_ids: list[str]) -> None:
"""Warn once per distinct set of competing same-workspace Kiro sessions.
Discovery returns ``None`` for both "no candidate yet" (transient) and ">=2
candidates" (won't ever bind until one goes away). This makes the latter
diagnosable without flooding logs on every ~0.7s poll.
"""
key = frozenset(session_ids)
if key in _ambiguous_discovery_warned:
return
_ambiguous_discovery_warned.add(key)
_logger.warning(
"kiro session discovery ambiguous; %d same-workspace sessions above the "
"launch floor, binding none to avoid cross-talk; workspace=%s sessions=%s",
len(session_ids),
workspace,
sorted(session_ids),
)
def _discover_kiro_session_jsonl(
*,
workspace: str,
launch_epoch_ms: int,
sessions_dir: Path | None = None,
) -> tuple[str, Path] | None:
"""Find this Omnigent session's Kiro JSONL file — only when it's unambiguous.
Candidates are Kiro sessions in the same workspace (``cwd``) with a parseable
``created_at`` at/after the launch floor (``launch_epoch_ms`` minus a small
skew, which already drops pre-existing sessions). We bind **only when exactly
one** session qualifies.
Each Kiro session is its own JSONL keyed by Kiro's minted id, so two fresh
sessions launched in the same workspace within the skew window both qualify.
Picking newest-by-``updated_at`` (the old behaviour) would latch onto
whichever session most recently emitted a turn — i.e. it can bind the *other*
session's transcript and silently cross-talk it into this conversation. With
two or more candidates we can't tell which JSONL is ours, so we return
``None`` and retry rather than guess (logged once via
:func:`_warn_ambiguous_discovery` so the ambiguous case is distinct from "not
written yet"). A brief delay is safe; mirroring the wrong conversation is not.
Mirrors cursor-native's "bind only when exactly one chat qualifies"
(:func:`omnigent.cursor_native_forwarder._discover_store`).
The resume/fork path doesn't reach here: when the Kiro id is already known the
caller binds it directly via :func:`_kiro_session_jsonl_for_id`.
"""
root = sessions_dir or _kiro_cli_sessions_dir()
if not root.is_dir():
return None
floor_ms = max(0, launch_epoch_ms - _DISCOVERY_SKEW_MS)
candidates: list[tuple[str, Path]] = []
for metadata_path in root.glob("*.json"):
session_id = metadata_path.stem
jsonl_path = root / f"{session_id}.jsonl"
if not jsonl_path.is_file():
continue
try:
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
except (OSError, ValueError):
continue
if not isinstance(metadata, dict) or not _same_workspace(metadata.get("cwd"), workspace):
continue
created_ms = _parse_iso_epoch_ms(metadata.get("created_at"))
# Require a parseable created_at at/after the floor. A fresh session always
# stamps created_at, so a missing/garbled one can't be confirmed in-window;
# dropping it stops an undateable straggler from poisoning the "exactly one"
# count and silently blocking discovery forever.
if not created_ms or created_ms < floor_ms:
continue
candidates.append((session_id, jsonl_path))
if len(candidates) > 1:
_warn_ambiguous_discovery(workspace, [session_id for session_id, _ in candidates])
return None
return candidates[0] if candidates else None
def _kiro_session_jsonl_for_id(
session_id: str,
*,
workspace: str,
sessions_dir: Path | None = None,
) -> Path | None:
"""Return the JSONL path for a known Kiro session id, if it is usable."""
root = sessions_dir or _kiro_cli_sessions_dir()
metadata_path = root / f"{session_id}.json"
jsonl_path = root / f"{session_id}.jsonl"
if not jsonl_path.is_file():
return None
try:
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
except (OSError, ValueError):
return None
if not isinstance(metadata, dict) or not _same_workspace(metadata.get("cwd"), workspace):
return None
return jsonl_path
def _read_new_kiro_messages(
jsonl_path: Path,
byte_offset: int,
) -> tuple[list[_KiroConversationMessage], int]:
"""Read conversation messages after *byte_offset* from Kiro's JSONL file."""
messages: list[_KiroConversationMessage] = []
try:
with jsonl_path.open("rb") as handle:
handle.seek(byte_offset)
# Advance only past newline-terminated lines: Kiro appends to this
# JSONL live, so the final line may be a record mid-write (no
# trailing ``\n``). Persisting ``handle.tell()`` past such a partial
# line would skip the record once Kiro finishes writing it. Hold the
# offset at the last complete line and re-read the tail next poll.
offset = byte_offset
for raw_line in handle:
if not raw_line.endswith(b"\n"):
break
offset += len(raw_line)
try:
line = raw_line.decode("utf-8")
except UnicodeDecodeError:
continue
message = _parse_kiro_jsonl_line(line)
if message is not None:
messages.append(message)
return messages, offset
except OSError:
return [], byte_offset
def _parse_kiro_jsonl_line(line: str) -> _KiroConversationMessage | None:
"""Parse one Kiro JSONL line into a mirrorable conversation message."""
try:
record = json.loads(line)
except ValueError:
return None
if not isinstance(record, dict):
return None
kind = record.get("kind")
if kind == "Prompt":
role = "user"
elif kind == "AssistantMessage":
role = "assistant"
else:
return None
data = record.get("data")
if not isinstance(data, dict):
return None
message_id = data.get("message_id")
if not isinstance(message_id, str) or not message_id:
return None
text = _kiro_content_text(data.get("content")).strip()
if not text:
return None
return _KiroConversationMessage(message_id=message_id, role=role, text=text)
def _kiro_content_text(content: object) -> str:
"""Join text blocks from Kiro's persisted message content."""
if isinstance(content, str):
return content
if not isinstance(content, list):
return ""
parts: list[str] = []
for block in content:
if not isinstance(block, dict):
continue
if block.get("kind") == "text" and isinstance(block.get("data"), str):
parts.append(block["data"])
elif block.get("type") in {"text", "output_text"} and isinstance(block.get("text"), str):
parts.append(block["text"])
return "\n".join(parts)
async def _post_conversation_message(
client: httpx.AsyncClient,
*,
session_id: str,
agent_name: str,
message: _KiroConversationMessage,
) -> None:
"""POST one Kiro message as an external conversation item."""
if message.role == "assistant":
item_data = {
"role": "assistant",
"agent": agent_name,
"content": [{"type": "output_text", "text": message.text}],
}
else:
item_data = {
"role": "user",
"content": [{"type": "input_text", "text": message.text}],
}
resp = await client.post(
f"/v1/sessions/{session_id}/events",
json={
"type": "external_conversation_item",
"data": {
"item_type": "message",
"item_data": item_data,
"response_id": f"kiro:{message.message_id}",
},
},
)
resp.raise_for_status()
async def _patch_external_session_id(
client: httpx.AsyncClient,
*,
session_id: str,
external_session_id: str,
) -> None:
"""Persist Kiro's native CLI session id onto the Omnigent session."""
resp = await client.patch(
f"/v1/sessions/{session_id}",
json={"external_session_id": external_session_id},
)
# The server rejects overwrites with a different id. Forwarding must keep
# running in that case; losing chat mirroring would be worse than failing to
# improve cold resume for an already-conflicted session.
if resp.status_code >= 400:
_logger.warning(
"AP rejected Kiro external_session_id PATCH (%s); session=%s kiro_session=%s",
resp.status_code,
session_id,
external_session_id,
)
return
def _read_kiro_cumulative_credits(metadata_path: Path) -> tuple[float | None, str | None]:
"""Sum the per-turn credit metering from a Kiro session ``.json`` snapshot.
kiro-cli meters in credits, not tokens: each turn under
``session_state.conversation_metadata.user_turn_metadatas`` carries a
``metering_usage`` list of ``{"value": <float>, "unit": "credit"}`` entries
(token counts are 0), and the CLI shows a per-turn ``Credits:`` line. The
cumulative session cost is the sum of every turn's credit values. This data
lives only in the ``.json`` snapshot, not the ``.jsonl`` transcript the
forwarder tails.
:returns: ``(cumulative_credits, model_id)``, or ``(None, None)`` when the
file is missing/unparseable or carries no metering yet.
"""
try:
raw = metadata_path.read_text(encoding="utf-8")
except OSError:
return None, None
try:
data = json.loads(raw)
except ValueError:
return None, None
if not isinstance(data, dict):
return None, None
session_state = data.get("session_state")
if not isinstance(session_state, dict):
return None, None
conversation_metadata = session_state.get("conversation_metadata")
turns = (
conversation_metadata.get("user_turn_metadatas")
if isinstance(conversation_metadata, dict)
else None
)
if not isinstance(turns, list):
return None, None
total = 0.0
saw_credit = False
for turn in turns:
if not isinstance(turn, dict):
continue
metering = turn.get("metering_usage")
if not isinstance(metering, list):
continue
for entry in metering:
if not isinstance(entry, dict):
continue
value = entry.get("value")
if isinstance(value, int | float) and not isinstance(value, bool):
total += float(value)
saw_credit = True
if not saw_credit:
return None, None
model_id: str | None = None
rts_model_state = session_state.get("rts_model_state")
if isinstance(rts_model_state, dict):
model_info = rts_model_state.get("model_info")
if isinstance(model_info, dict):
candidate = model_info.get("model_id")
if isinstance(candidate, str) and candidate:
model_id = candidate
return total, model_id
def _read_kiro_current_model(metadata_path: Path) -> str | None:
"""Return kiro's current model id from the session ``.json`` snapshot.
Reads ``session_state.rts_model_state.model_info.model_id`` (e.g. ``"auto"``
or ``"claude-haiku-4.5"``), independent of the metering read above so it is
available at launch before any turn. kiro updates it in place on a ``/model``
switch, so polling it mirrors both the launched model and live TUI switches.
:returns: The model id, or ``None`` when the file is missing/unparseable or
carries no model state yet.
"""
try:
raw = metadata_path.read_text(encoding="utf-8")
except OSError:
return None
try:
data = json.loads(raw)
except ValueError:
return None
if not isinstance(data, dict):
return None
session_state = data.get("session_state")
if not isinstance(session_state, dict):
return None
rts_model_state = session_state.get("rts_model_state")
if not isinstance(rts_model_state, dict):
return None
model_info = rts_model_state.get("model_info")
if not isinstance(model_info, dict):
return None
model_id = model_info.get("model_id")
return model_id if isinstance(model_id, str) and model_id else None
async def _post_external_model_change(
client: httpx.AsyncClient,
*,
session_id: str,
model: str,
) -> None:
"""Mirror kiro's current model to the web as ``external_model_change``.
The server persists this as ``model_override`` (so the picker shows the real
model instead of falling back to the harness name) and deliberately does NOT
forward a ``/model`` back to the runner, so mirroring the model the TUI is
already on cannot loop. Mirrors cursor-native's terminal->web model mirror.
"""
resp = await client.post(
f"/v1/sessions/{session_id}/events",
json={"type": "external_model_change", "data": {"model": model}},
)
resp.raise_for_status()
async def _post_session_cost(
client: httpx.AsyncClient,
*,
session_id: str,
cumulative_cost_usd: float,
model: str | None,
) -> None:
"""POST Kiro's cumulative credit spend as authoritative session cost.
kiro-cli reports cost in credits and there is no credit->USD conversion
available, so credits are forwarded 1:1 into ``cumulative_cost_usd`` (the
same convention the Copilot relay uses for its AI-credit total). The server
treats this value as authoritative and monotonic, in preference to
token x catalog pricing, via the ``external_session_usage`` event used by
the claude-/codex-native forwarders.
"""
data: dict[str, object] = {"cumulative_cost_usd": cumulative_cost_usd}
if model:
data["model"] = model
resp = await client.post(
f"/v1/sessions/{session_id}/events",
json={"type": "external_session_usage", "data": data},
)
resp.raise_for_status()
async def forward_kiro_session_to_omnigent(
*,
base_url: str,
headers: dict[str, str],
session_id: str,
bridge_dir: Path,
agent_name: str,
workspace: str,
launch_epoch_ms: int,
expected_session_id: str | None = None,
poll_interval_s: float = _DEFAULT_POLL_INTERVAL_S,
auth: httpx.Auth | None = None,
) -> None:
"""Tail Kiro's session JSONL and mirror assistant messages into AP."""
state = _read_state(bridge_dir)
jsonl_path: Path | None = None
timeout = httpx.Timeout(_POST_TIMEOUT_S)
mirrored_external_session_id: str | None = None
last_posted_cost: float | None = None
last_posted_model: str | None = None
async with httpx.AsyncClient(
base_url=base_url, headers=headers, auth=auth, timeout=timeout
) as client:
while True:
try:
if state.session_id is None or jsonl_path is None or not jsonl_path.exists():
discovered: tuple[str, Path] | None = None
if expected_session_id:
expected_path = await asyncio.to_thread(
_kiro_session_jsonl_for_id,
expected_session_id,
workspace=workspace,
)
if expected_path is not None:
discovered = (expected_session_id, expected_path)
elif discovered is None:
discovered = await asyncio.to_thread(
_discover_kiro_session_jsonl,
workspace=workspace,
launch_epoch_ms=launch_epoch_ms,
)
if discovered is not None:
discovered_session_id, discovered_path = discovered
if state.session_id != discovered_session_id:
state = _ForwardState(session_id=discovered_session_id, byte_offset=0)
_write_state(bridge_dir, state)
jsonl_path = discovered_path
if jsonl_path is not None and state.session_id is not None:
if mirrored_external_session_id != state.session_id:
await _patch_external_session_id(
client,
session_id=session_id,
external_session_id=state.session_id,
)
mirrored_external_session_id = state.session_id
messages, byte_offset = await asyncio.to_thread(
_read_new_kiro_messages,
jsonl_path,
state.byte_offset,
)
for message in messages:
# Mirror the transcript only. Running/idle status for
# kiro-native is owned by the PTY watcher's ``emit_status``
# (``runner/resource_registry.py``), matching the other
# forwarder-backed native harnesses (goose/qwen/hermes),
# whose forwarders mirror the transcript, not the
# running/idle session status. Posting status here too
# double-sourced it (#1137).
await _post_conversation_message(
client,
session_id=session_id,
agent_name=agent_name,
message=message,
)
if byte_offset != state.byte_offset:
state.byte_offset = byte_offset
_write_state(bridge_dir, state)
write_forwarder_ready(bridge_dir)
# Forward kiro's credit metering as authoritative session
# cost. It lives in the ``.json`` snapshot (sibling of the
# tailed ``.jsonl``), not the transcript, and is cumulative;
# the server treats ``cumulative_cost_usd`` as monotonic, so
# only post when it advances.
cumulative_cost, cost_model = await asyncio.to_thread(
_read_kiro_cumulative_credits, jsonl_path.with_suffix(".json")
)
if cumulative_cost is not None and (
last_posted_cost is None or cumulative_cost > last_posted_cost
):
await _post_session_cost(
client,
session_id=session_id,
cumulative_cost_usd=cumulative_cost,
model=cost_model,
)
last_posted_cost = cumulative_cost
# Mirror kiro's current model so the web picker shows the real
# model (not the harness name) at launch and after a live
# ``/model`` switch. Read independently of metering so it is
# available before the first turn; posted only when it changes.
current_model = await asyncio.to_thread(
_read_kiro_current_model, jsonl_path.with_suffix(".json")
)
if current_model is not None and current_model != last_posted_model:
await _post_external_model_change(
client, session_id=session_id, model=current_model
)
last_posted_model = current_model
except asyncio.CancelledError:
raise
except Exception:
_logger.exception(
"kiro session forwarder poll failed; session=%s bridge_dir=%s",
session_id,
bridge_dir,
)
await asyncio.sleep(poll_interval_s)
def _supervisor_monotonic() -> float:
"""Indirection so tests can stub the supervisor clock."""
return time.monotonic()
async def _supervisor_sleep(seconds: float) -> None:
"""Indirection so tests can stub supervisor sleep."""
await asyncio.sleep(seconds)
async def supervise_kiro_session_forwarder(
*,
base_url: str,
headers: dict[str, str],
session_id: str,
bridge_dir: Path,
agent_name: str,
workspace: str,
launch_epoch_ms: int,
expected_session_id: str | None = None,
poll_interval_s: float = _DEFAULT_POLL_INTERVAL_S,
auth: httpx.Auth | None = None,
) -> None:
"""Run the Kiro session forwarder under a restart supervisor."""
backoff_s = _SUPERVISOR_INITIAL_BACKOFF_S
while True:
run_started_at = _supervisor_monotonic()
crash_exc: Exception | None = None
try:
await forward_kiro_session_to_omnigent(
base_url=base_url,
headers=headers,
session_id=session_id,
bridge_dir=bridge_dir,
agent_name=agent_name,
workspace=workspace,
launch_epoch_ms=launch_epoch_ms,
expected_session_id=expected_session_id,
poll_interval_s=poll_interval_s,
auth=auth,
)
_logger.warning(
"kiro session forwarder returned unexpectedly; restarting; "
"session=%s bridge_dir=%s",
session_id,
bridge_dir,
)
except asyncio.CancelledError:
raise
except Exception as exc: # noqa: BLE001 - supervisor restarts on any crash
crash_exc = exc
if _supervisor_monotonic() - run_started_at >= _SUPERVISOR_HEALTHY_UPTIME_S:
backoff_s = _SUPERVISOR_INITIAL_BACKOFF_S
if crash_exc is not None:
_logger.error(
"kiro session forwarder crashed; restarting in %.1fs; session=%s bridge_dir=%s",
backoff_s,
session_id,
bridge_dir,
exc_info=crash_exc,
)
await _supervisor_sleep(backoff_s)
backoff_s = min(backoff_s * 2, _SUPERVISOR_MAX_BACKOFF_S)
__all__ = [
"forward_kiro_session_to_omnigent",
"supervise_kiro_session_forwarder",
]