coding-cli-runtime 0.1.0__py3-none-any.whl

coding_cli_runtime/session_execution.py

@@ -0,0 +1,604 @@
+"""Shared async session execution and transcript mirroring helpers."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import shutil
+import signal
+import time
+from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+def _default_source_exists(source: Any) -> bool:
+    return isinstance(source, Path) and source.exists()
+
+
+def _default_describe_source(source: Any) -> str:
+    if isinstance(source, Path):
+        return str(source)
+    return str(source)
+
+
+def _default_copy_source(source: Any, destination: Path) -> None:
+    if not isinstance(source, Path):
+        raise TypeError(f"Unsupported transcript source type: {type(source)!r}")
+    shutil.copy2(source, destination)
+
+
+def _default_source_identity(source: Any) -> Any:
+    return source
+
+
+@dataclass(frozen=True)
+class TranscriptMirrorStrategy:
+    provider_label: str
+    destination: Path
+    locate_source: Callable[[], object | None]
+    source_identity: Callable[[object], object] = _default_source_identity
+    source_exists: Callable[[object], bool] = _default_source_exists
+    describe_source: Callable[[object], str] = _default_describe_source
+    copy_source: Callable[[object, Path], None] = _default_copy_source
+    final_copy_source: Callable[[object, Path], None] | None = None
+    poll_interval: float = 1.0
+    refresh_source_each_poll: bool = False
+    progress_callback: Callable[[SessionProgressEvent], None] | None = None
+    activity_byte_threshold: int = 1024
+    idle_heartbeat_seconds: float = 15.0
+    final_retry_attempts: int = 1
+    final_retry_initial_delay_seconds: float = 0.0
+    final_retry_backoff_multiplier: float = 1.0
+
+
+@dataclass(frozen=True)
+class SessionRetryDecision:
+    retry: bool
+    delay_seconds: float = 0.0
+    category: str | None = None
+
+
+@dataclass(frozen=True)
+class InteractiveCliRunResult:
+    stdout: bytes
+    stderr: bytes
+    returncode: int | None
+    conversation_path: Path | None
+    cli_attempt_used: int
+    max_cli_attempts: int
+    started_at: float
+
+
+@dataclass(frozen=True)
+class SessionProgressEvent:
+    event_type: str
+    conversation_path: Path
+    source_description: str | None = None
+    activity_kind: str | None = None
+    mirrored_bytes: int | None = None
+    last_activity_at: float | None = None
+
+
+class SessionExecutionTimeoutError(RuntimeError):
+    def __init__(
+        self,
+        message: str,
+        *,
+        conversation_path: Path | None = None,
+        started_at: float | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.conversation_path = conversation_path
+        self.started_at = started_at
+
+
+def _mark_isolated_process_group(proc: asyncio.subprocess.Process) -> None:
+    try:
+        proc._llm_eval_isolated_process_group = True  # type: ignore[attr-defined]
+    except Exception:
+        pass
+
+
+def _has_isolated_process_group(proc: asyncio.subprocess.Process) -> bool:
+    return bool(getattr(proc, "_llm_eval_isolated_process_group", False))
+
+
+def _close_process_transport(
+    proc: asyncio.subprocess.Process | None,
+    *,
+    logger: logging.Logger | None = None,
+    label: str = "process",
+) -> None:
+    if proc is None:
+        return
+    transport = getattr(proc, "_transport", None)
+    if transport is None:
+        return
+    try:
+        transport.close()
+    except Exception as exc:
+        if logger:
+            logger.debug("Failed to close %s transport: %s", label, exc)
+
+
+def _safe_file_size(path: Path) -> int | None:
+    try:
+        return path.stat().st_size
+    except OSError:
+        return None
+
+
+def _emit_progress_event(
+    *,
+    callback: Callable[[SessionProgressEvent], None] | None,
+    event: SessionProgressEvent,
+    logger: logging.Logger,
+    job_name: str,
+    phase_tag: str,
+) -> None:
+    if callback is None:
+        return
+    try:
+        callback(event)
+    except Exception as exc:
+        logger.debug(
+            "[%s] %sprogress callback failed for %s: %s",
+            job_name,
+            phase_tag,
+            event.event_type,
+            exc,
+        )
+
+
+async def mirror_session_transcript(
+    *,
+    job_name: str,
+    phase_tag: str,
+    logger: logging.Logger,
+    proc: asyncio.subprocess.Process,
+    strategy: TranscriptMirrorStrategy,
+) -> Path | None:
+    destination = strategy.destination
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        destination.unlink(missing_ok=True)
+    except FileNotFoundError:
+        pass
+    destination.touch()
+
+    source: object | None = None
+    source_key: object | None = None
+    announced = False
+    last_mirrored_bytes: int | None = None
+    last_activity_at = time.time()
+    last_heartbeat_emit_at = 0.0
+    activity_byte_threshold = max(1, int(strategy.activity_byte_threshold))
+    idle_heartbeat_seconds = max(0.0, float(strategy.idle_heartbeat_seconds))
+
+    def emit_progress(
+        event_type: str,
+        *,
+        activity_kind: str,
+        source_description: str | None = None,
+        mirrored_bytes: int | None = None,
+        last_activity_ts: float | None = None,
+    ) -> None:
+        _emit_progress_event(
+            callback=strategy.progress_callback,
+            event=SessionProgressEvent(
+                event_type=event_type,
+                conversation_path=destination,
+                source_description=source_description,
+                activity_kind=activity_kind,
+                mirrored_bytes=mirrored_bytes,
+                last_activity_at=last_activity_ts,
+            ),
+            logger=logger,
+            job_name=job_name,
+            phase_tag=phase_tag,
+        )
+
+    while True:
+        latest_source = None
+        latest_source_key = None
+        if source is None or strategy.refresh_source_each_poll:
+            latest_source = strategy.locate_source()
+        if latest_source is not None:
+            latest_source_key = strategy.source_identity(latest_source)
+        source_changed = latest_source is not None and latest_source_key != source_key
+        if latest_source is not None:
+            if source is not None:
+                if source_changed:
+                    logger.info(
+                        "[%s] %sswitching %s conversation source to %s",
+                        job_name,
+                        phase_tag,
+                        strategy.provider_label,
+                        strategy.describe_source(latest_source),
+                    )
+            source = latest_source
+            source_key = latest_source_key
+        if source is not None and strategy.source_exists(source):
+            source_description = strategy.describe_source(source)
+            first_attachment = not announced
+            if first_attachment:
+                logger.info(
+                    "[%s] %sstreaming %s conversation from %s",
+                    job_name,
+                    phase_tag,
+                    strategy.provider_label,
+                    source_description,
+                )
+            try:
+                strategy.copy_source(source, destination)
+                mirrored_bytes = _safe_file_size(destination)
+                now = time.time()
+                if first_attachment or source_changed:
+                    attach_kind = "attached" if first_attachment else "source_switched"
+                    emit_progress(
+                        "transcript_attached",
+                        activity_kind=attach_kind,
+                        source_description=source_description,
+                        mirrored_bytes=mirrored_bytes,
+                        last_activity_ts=now,
+                    )
+                    last_heartbeat_emit_at = now
+                if mirrored_bytes is not None and mirrored_bytes != last_mirrored_bytes:
+                    bytes_delta = abs(mirrored_bytes - (last_mirrored_bytes or 0))
+                    last_mirrored_bytes = mirrored_bytes
+                    last_activity_at = now
+                    if bytes_delta >= activity_byte_threshold:
+                        emit_progress(
+                            "provider_call_heartbeat",
+                            activity_kind="transcript_growth",
+                            source_description=source_description,
+                            mirrored_bytes=mirrored_bytes,
+                            last_activity_ts=last_activity_at,
+                        )
+                        last_heartbeat_emit_at = now
+                if first_attachment:
+                    announced = True
+            except OSError as exc:
+                logger.debug(
+                    "[%s] %sfailed mirroring %s transcript: %s",
+                    job_name,
+                    phase_tag,
+                    strategy.provider_label,
+                    exc,
+                )
+        now = time.time()
+        if (
+            proc.returncode is None
+            and idle_heartbeat_seconds > 0
+            and strategy.progress_callback is not None
+            and now - last_activity_at >= idle_heartbeat_seconds
+            and now - last_heartbeat_emit_at >= idle_heartbeat_seconds
+        ):
+            emit_progress(
+                "provider_call_heartbeat",
+                activity_kind="idle",
+                source_description=(
+                    strategy.describe_source(source)
+                    if source is not None and strategy.source_exists(source)
+                    else None
+                ),
+                mirrored_bytes=(
+                    last_mirrored_bytes
+                    if last_mirrored_bytes is not None
+                    else _safe_file_size(destination)
+                ),
+                last_activity_ts=last_activity_at,
+            )
+            last_heartbeat_emit_at = now
+        if proc.returncode is not None:
+            break
+        await asyncio.sleep(strategy.poll_interval)
+
+    retry_attempts = max(1, int(strategy.final_retry_attempts))
+    retry_delay = max(0.0, float(strategy.final_retry_initial_delay_seconds))
+    retry_multiplier = max(1.0, float(strategy.final_retry_backoff_multiplier))
+    for attempt_index in range(retry_attempts):
+        if attempt_index > 0 and retry_delay > 0:
+            await asyncio.sleep(retry_delay)
+            retry_delay *= retry_multiplier
+        latest_source = strategy.locate_source()
+        latest_source_key = (
+            strategy.source_identity(latest_source) if latest_source is not None else None
+        )
+        final_source_changed = latest_source is not None and latest_source_key != source_key
+        if latest_source is not None:
+            if source is not None:
+                if final_source_changed:
+                    logger.info(
+                        "[%s] %sswitching %s conversation source to %s",
+                        job_name,
+                        phase_tag,
+                        strategy.provider_label,
+                        strategy.describe_source(latest_source),
+                    )
+            source = latest_source
+            source_key = latest_source_key
+        else:
+            source = source or latest_source
+        if source is not None and strategy.source_exists(source):
+            try:
+                final_copy = strategy.final_copy_source or strategy.copy_source
+                final_copy(source, destination)
+                final_size = _safe_file_size(destination)
+                if not announced or final_source_changed:
+                    emit_progress(
+                        "transcript_attached",
+                        activity_kind="source_switched" if final_source_changed else "attached",
+                        source_description=strategy.describe_source(source),
+                        mirrored_bytes=final_size,
+                        last_activity_ts=time.time(),
+                    )
+                return destination
+            except OSError as exc:
+                logger.warning(
+                    "[%s] %sfailed final copy for %s transcript: %s",
+                    job_name,
+                    phase_tag,
+                    strategy.provider_label,
+                    exc,
+                )
+        else:
+            logger.debug(
+                "[%s] %s%s transcript not found",
+                job_name,
+                phase_tag,
+                strategy.provider_label,
+            )
+    return destination if destination.exists() else None
+
+
+async def _terminate_process(
+    proc: asyncio.subprocess.Process,
+    *,
+    logger: logging.Logger | None = None,
+    label: str = "process",
+    timeout: float = 5.0,
+) -> None:
+    if proc.returncode is not None:
+        _close_process_transport(proc, logger=logger, label=label)
+        return
+    isolated_group = (
+        os.name != "nt"
+        and _has_isolated_process_group(proc)
+        and isinstance(getattr(proc, "pid", None), int)
+        and int(proc.pid) > 0
+    )
+    try:
+        if isolated_group:
+            os.killpg(int(proc.pid), signal.SIGTERM)
+        else:
+            proc.terminate()
+    except ProcessLookupError:
+        _close_process_transport(proc, logger=logger, label=label)
+        return
+    try:
+        await asyncio.wait_for(proc.wait(), timeout=timeout)
+        _close_process_transport(proc, logger=logger, label=label)
+        return
+    except TimeoutError:
+        pass
+    try:
+        if isolated_group:
+            os.killpg(int(proc.pid), signal.SIGKILL)
+        else:
+            proc.kill()
+    except ProcessLookupError:
+        _close_process_transport(proc, logger=logger, label=label)
+        return
+    try:
+        await asyncio.wait_for(proc.wait(), timeout=timeout)
+    except TimeoutError:
+        if logger:
+            logger.warning("Timed out while killing %s", label)
+    finally:
+        _close_process_transport(proc, logger=logger, label=label)
+
+
+async def _finalize_transcript_task(
+    *,
+    transcript_task: asyncio.Task[Path | None] | None,
+    logger: logging.Logger,
+    job_name: str,
+    phase_tag: str,
+) -> Path | None:
+    if transcript_task is None:
+        return None
+    try:
+        return await asyncio.wait_for(transcript_task, timeout=5)
+    except TimeoutError:
+        transcript_task.cancel()
+        try:
+            await transcript_task
+        except asyncio.CancelledError:
+            logger.debug(
+                "[%s] %sconversation log capture canceled to avoid blocking",
+                job_name,
+                phase_tag,
+            )
+        except Exception as exc:
+            logger.debug(
+                "[%s] %sconversation log capture canceled with error: %s",
+                job_name,
+                phase_tag,
+                exc,
+            )
+    except Exception as exc:
+        logger.debug(
+            "[%s] %sconversation log capture failed: %s",
+            job_name,
+            phase_tag,
+            exc,
+        )
+    return None
+
+
+async def _cancel_transcript_task(
+    *,
+    transcript_task: asyncio.Task[Path | None] | None,
+    logger: logging.Logger,
+    job_name: str,
+    phase_tag: str,
+    reason: str,
+) -> None:
+    if transcript_task is None or transcript_task.done():
+        return
+    transcript_task.cancel()
+    try:
+        await transcript_task
+    except asyncio.CancelledError:
+        logger.debug(
+            "[%s] %sconversation log capture canceled after %s",
+            job_name,
+            phase_tag,
+            reason,
+        )
+    except Exception as exc:
+        logger.debug(
+            "[%s] %sconversation log cancel failed: %s",
+            job_name,
+            phase_tag,
+            exc,
+        )
+
+
+async def run_interactive_session(
+    *,
+    cmd_parts: Sequence[str],
+    cwd: Path,
+    stdin_text: str | None,
+    logger: logging.Logger,
+    provider_label: str = "cli",
+    job_name: str = "session",
+    phase_tag: str = "",
+    process_label: str = "process",
+    timeout_seconds: float | None = None,
+    env: dict[str, str] | None = None,
+    create_process: Callable[..., Awaitable[asyncio.subprocess.Process]] | None = None,
+    transcript_factory: Callable[[asyncio.subprocess.Process], Awaitable[Path | None]]
+    | None = None,
+    max_cli_attempts: int = 1,
+    retry_decider: Callable[[InteractiveCliRunResult, int, int], SessionRetryDecision | bool]
+    | None = None,
+) -> InteractiveCliRunResult:
+    process_factory = create_process or asyncio.create_subprocess_exec
+    stdin_payload = stdin_text.encode("utf-8") if stdin_text is not None else None
+    last_result: InteractiveCliRunResult | None = None
+
+    for cli_attempt in range(1, max(1, max_cli_attempts) + 1):
+        started_at = time.time()
+        creation_kwargs = {
+            "stdout": asyncio.subprocess.PIPE,
+            "stderr": asyncio.subprocess.PIPE,
+            "stdin": (
+                asyncio.subprocess.PIPE if stdin_text is not None else asyncio.subprocess.DEVNULL
+            ),
+            "cwd": str(cwd),
+            "env": env,
+        }
+        if os.name != "nt":
+            creation_kwargs["start_new_session"] = True
+        proc = await process_factory(
+            *cmd_parts,
+            **creation_kwargs,  # type: ignore[arg-type]
+        )
+        if os.name != "nt":
+            _mark_isolated_process_group(proc)
+        transcript_task: asyncio.Task[Path | None] | None = (
+            asyncio.create_task(transcript_factory(proc))  # type: ignore[arg-type]
+            if transcript_factory is not None
+            else None
+        )
+        try:
+            if timeout_seconds and timeout_seconds > 0:
+                stdout, stderr = await asyncio.wait_for(
+                    proc.communicate(input=stdin_payload),
+                    timeout=timeout_seconds,
+                )
+            else:
+                stdout, stderr = await proc.communicate(input=stdin_payload)
+        except asyncio.CancelledError as exc:
+            await _terminate_process(proc, logger=logger, label=process_label)
+            conversation_path = await _finalize_transcript_task(
+                transcript_task=transcript_task,
+                logger=logger,
+                job_name=job_name,
+                phase_tag=phase_tag,
+            )
+            try:
+                exc.conversation_path = conversation_path  # type: ignore[attr-defined]
+                exc.started_at = started_at  # type: ignore[attr-defined]
+            except Exception:
+                pass
+            raise
+        except TimeoutError:
+            logger.error(
+                "[%s] %s%s run timed out after %ss",
+                job_name,
+                phase_tag,
+                provider_label,
+                timeout_seconds,
+            )
+            await _terminate_process(proc, logger=logger, label=process_label)
+            conversation_path = await _finalize_transcript_task(
+                transcript_task=transcript_task,
+                logger=logger,
+                job_name=job_name,
+                phase_tag=phase_tag,
+            )
+            raise SessionExecutionTimeoutError(
+                f"{provider_label.lower()} run timed out after {timeout_seconds}s",
+                conversation_path=conversation_path,
+                started_at=started_at,
+            ) from None
+
+        conversation_path = await _finalize_transcript_task(
+            transcript_task=transcript_task,
+            logger=logger,
+            job_name=job_name,
+            phase_tag=phase_tag,
+        )
+        result = InteractiveCliRunResult(
+            stdout=stdout,
+            stderr=stderr,
+            returncode=proc.returncode,
+            conversation_path=conversation_path,
+            cli_attempt_used=cli_attempt,
+            max_cli_attempts=max(1, max_cli_attempts),
+            started_at=started_at,
+        )
+        last_result = result
+        if proc.returncode == 0:
+            return result
+
+        if retry_decider is None or cli_attempt >= max(1, max_cli_attempts):
+            return result
+
+        decision = retry_decider(result, cli_attempt, max(1, max_cli_attempts))
+        if isinstance(decision, bool):
+            decision = SessionRetryDecision(retry=decision)
+        if not decision.retry:
+            return result
+        logger.warning(
+            "[%s] %s%s retryable failure (%s, %s/%s); retrying in %.1fs",
+            job_name,
+            phase_tag,
+            provider_label,
+            decision.category or "retryable",
+            cli_attempt,
+            max(1, max_cli_attempts),
+            decision.delay_seconds,
+        )
+        if decision.delay_seconds > 0:
+            await asyncio.sleep(decision.delay_seconds)
+
+    if last_result is None:
+        raise RuntimeError(f"{provider_label.lower()} run failed before launch")
+    return last_result