dataface 0.1.2__py3-none-any.whl

dataface/cli/commands/chat.py

@@ -0,0 +1,645 @@
+"""Terminal chat CLI (`dft chat`)."""
+
+from __future__ import annotations
+
+import atexit
+import json
+import re
+import shutil
+import signal
+import tempfile
+import time
+from collections.abc import Callable
+from contextlib import AbstractContextManager, nullcontext
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated, Any, Literal
+
+if TYPE_CHECKING:
+    from dataface.ai.external_mcp import ExternalMCPManager
+
+import typer
+from rich.console import Console
+from rich.panel import Panel
+from rich.syntax import Syntax
+from rich.table import Table
+
+from dataface.agent_api import chat as _api
+from dataface.agent_api._state import dft_home
+from dataface.agent_api.file_refs import expand_file_refs
+from dataface.ai.context import DatafaceAIContext
+from dataface.ai.events import (
+    AgentDone,
+    AgentError,
+    ContentDelta,
+    ThinkingStatus,
+    ToolCallEvent,
+    ToolResultEvent,
+)
+from dataface.ai.llm import create_client, infer_provider
+from dataface.cli.commands._agent_input import (
+    PromptToolkitInput,
+    select_input_layer,
+)
+
+# Max JSON lines before we truncate and write the full result to a tmpfile.
+_JSON_TRUNCATE_LINES = 80
+
+# Completed stream status — used by _run_agent_session callers.
+_SessionStatus = Literal["done", "interrupted", "errored"]
+
+
+# ---------------------------------------------------------------------------
+# Interrupt sentinel — BaseException so it bypasses 'except Exception' in tools
+# ---------------------------------------------------------------------------
+
+
+class _AgentInterrupt(BaseException):
+    """Raised by the SIGINT handler to cancel an in-flight LLM/tool call."""
+
+
+# ---------------------------------------------------------------------------
+# Session dataclass (CLI REPL state — holds console, not persisted)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class Session:
+    """Mutable REPL state for one dft chat session."""
+
+    client: Any
+    context: DatafaceAIContext
+    console: Console
+    messages: list[dict[str, Any]] = field(default_factory=list)
+
+
+# ---------------------------------------------------------------------------
+# Slash-command helpers
+# ---------------------------------------------------------------------------
+
+
+def _cmd_exit(session: Session, args: list[str]) -> None:
+    session.console.print("[dim]Session ended.[/dim]")
+    raise SystemExit(0)
+
+
+def _cmd_clear(session: Session, args: list[str]) -> None:
+    session.messages.clear()
+    # Rebuild the client to reset any provider-side conversation state
+    # (e.g. OpenAIClient._previous_response_id / _sent_message_count).
+    # Without this, OpenAI's incremental mode would compute an empty input
+    # slice on the next prompt and silently return nothing.
+    session.client = create_client(model=session.client.model)
+    session.console.print("[dim]Conversation history cleared.[/dim]")
+
+
+def _cmd_model(session: Session, args: list[str]) -> None:
+    if not args:
+        session.console.print(
+            f"[dim]Current model: {session.client.model} "
+            f"(provider: {session.client.provider})[/dim]"
+        )
+        return
+    requested_model = args[0]
+    requested_provider = infer_provider(requested_model)
+    if requested_provider != session.client.provider:
+        session.console.print(
+            f"[red]/model: cannot switch providers mid-session "
+            f"(current: {session.client.provider}, requested: {requested_provider}). "
+            f"Restart dft chat with --model {requested_model}.[/red]"
+        )
+        return
+    new_client = create_client(model=requested_model)
+    session.client = new_client
+    # Clear messages: the new client has no previous_response_id chain,
+    # so prior tool-call history would cause LLMClientError on OpenAI.
+    session.messages.clear()
+    session.console.print(
+        f"[dim]Switched to model: {new_client.model} (history cleared)[/dim]"
+    )
+
+
+def _cmd_help(session: Session, args: list[str]) -> None:
+    lines = [
+        "[bold]/help[/bold]   — show this help",
+        "[bold]/clear[/bold]  — drop conversation history (keep session alive)",
+        "[bold]/model[/bold] [name] — show current model or switch to another (same provider only)",
+        "[bold]/exit[/bold]   — end the session",
+        "[bold]/quit[/bold]   — alias for /exit",
+        "",
+        "Ctrl+C  — cancel in-flight request; second Ctrl+C within 1s exits",
+        "Ctrl+D  — exit session",
+    ]
+    session.console.print("\n".join(lines))
+
+
+SLASH_COMMANDS: dict[str, Callable[[Session, list[str]], None]] = {
+    "exit": _cmd_exit,
+    "quit": _cmd_exit,
+    "clear": _cmd_clear,
+    "model": _cmd_model,
+    "help": _cmd_help,
+}
+
+
+def _dispatch_slash(session: Session, user_input: str) -> bool:
+    """Try to dispatch a slash command.
+
+    Returns True if the input was consumed as a slash command, False if the
+    input should be sent to the model.
+
+    Unknown slash commands pass through to the model verbatim (per spec: "no
+    magic").
+    """
+    stripped = user_input.strip()
+    if not stripped.startswith("/"):
+        return False
+    parts = stripped[1:].split()
+    if not parts:
+        return False
+    cmd_name = parts[0].lower()
+    cmd_args = parts[1:]
+    handler = SLASH_COMMANDS.get(cmd_name)
+    if handler is None:
+        return False
+    handler(session, cmd_args)
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Ctrl+C double-tap-exit helpers
+# ---------------------------------------------------------------------------
+
+
+def _check_double_tap_exit(last_sigint_ts: list[float]) -> bool:
+    """Record an interrupt timestamp and return True if it's a double-tap.
+
+    Args:
+        last_sigint_ts: A mutable list used as a single-slot timestamp store.
+            Pass an empty list on first call; the function appends/updates it.
+
+    Returns:
+        True if a second interrupt arrived within 1s of the previous one.
+    """
+    now = time.monotonic()
+    if last_sigint_ts and (now - last_sigint_ts[0]) < 1.0:
+        return True
+    if last_sigint_ts:
+        last_sigint_ts[0] = now
+    else:
+        last_sigint_ts.append(now)
+    return False
+
+
+# ---------------------------------------------------------------------------
+# Tool-result rendering
+# ---------------------------------------------------------------------------
+
+
+def _render_tool_result(name: str, result: Any, console: Console, tmpdir: Path) -> None:
+    """Render a tool result to the console.
+
+    Dispatch:
+    - Tabular execute_query result → rich.Table
+    - Error dict → red Panel
+    - Oversized JSON → truncated Panel with tmpfile reference
+    - Default → JSON Syntax Panel
+
+    tmpdir is a session-scoped directory for oversized-result files; the
+    caller cleans it up at session end.
+    """
+    # Tabular: execute_query returns {success, data: list[dict], columns, ...}
+    if (
+        isinstance(result, dict)
+        and isinstance(result.get("data"), list)
+        and isinstance(result.get("columns"), list)
+        and result.get("columns")
+        and not result.get("error")
+    ):
+        _render_tabular(name, result, console)
+        return
+
+    # Error dict
+    if isinstance(result, dict) and result.get("error"):
+        console.print(
+            Panel(
+                f"[red]{result['error']}[/red]",
+                title=f"Tool Error: {name}",
+                border_style="red",
+            )
+        )
+        return
+
+    # Default: JSON, possibly truncated
+    body = json.dumps(result, indent=2, default=str, sort_keys=True)
+    lines = body.splitlines()
+    if len(lines) > _JSON_TRUNCATE_LINES:
+        truncated_body = "\n".join(lines[:_JSON_TRUNCATE_LINES])
+        extra = len(lines) - _JSON_TRUNCATE_LINES
+        # Write the full result under the session tmpdir so it's cleaned up on exit.
+        # Sanitize name to avoid path separators from LLM-provided tool names.
+        safe_name = re.sub(r"[^\w-]", "_", name)
+        tmp_file = tmpdir / f"dft_{safe_name}_{int(time.monotonic() * 1000)}.json"
+        tmp_file.write_text(body)
+        summary = f"[dim]... ({extra} more lines, full result in {tmp_file})[/dim]"
+        console.print(
+            Panel(
+                Syntax(truncated_body, "json", word_wrap=True),
+                title=f"Tool Result: {name}",
+                border_style="blue",
+                subtitle=summary,
+            )
+        )
+    else:
+        console.print(
+            Panel(
+                Syntax(body, "json", word_wrap=True),
+                title=f"Tool Result: {name}",
+                border_style="blue",
+            )
+        )
+
+
+def _render_tabular(name: str, result: dict[str, Any], console: Console) -> None:
+    columns: list[str] = result["columns"]
+    data: list[dict[str, Any]] = result["data"]
+
+    table = Table(title=f"Tool Result: {name}", border_style="blue", show_lines=False)
+    for col in columns:
+        table.add_column(col, overflow="fold")
+    for row in data:
+        table.add_row(*[str(row.get(col, "")) for col in columns])
+
+    console.print(table)
+    if result.get("truncated"):
+        console.print(f"[dim](results truncated to {len(data)} rows)[/dim]")
+
+
+# ---------------------------------------------------------------------------
+# Agent stream loop — extracted so tests can call it directly
+# ---------------------------------------------------------------------------
+
+
+def _run_agent_session(
+    chat_session: _api.ChatSession,
+    prompt: str,
+    console: Console,
+    tmpdir: Path,
+    tools: list[dict[str, Any]] | None = None,
+    external_manager: ExternalMCPManager | None = None,
+) -> _SessionStatus:
+    """Stream one prompt's events through ``agent_api.chat.send_message``.
+
+    Installs a SIGINT handler for the duration that raises _AgentInterrupt.
+    Persistence is handled inside ``send_message`` (its ``try/finally`` writes
+    the turn to disk even when the generator is closed mid-stream).
+
+    Args:
+        tools: Merged tool list (built-in + external). Defaults to ALL_TOOLS.
+        external_manager: ExternalMCPManager for routing 'server__tool' calls.
+
+    Returns:
+        "done"        — completed normally
+        "interrupted" — cancelled by Ctrl+C
+        "errored"     — AgentError received (message already printed in red)
+    """
+    started_content = False
+    status: _SessionStatus = "done"
+
+    def _sigint_handler(_signum: int, _frame: Any) -> None:
+        raise _AgentInterrupt()
+
+    prev_handler = signal.signal(signal.SIGINT, _sigint_handler)
+
+    gen = _api.send_message(
+        chat_session,
+        prompt,
+        tools=tools,
+        external_manager=external_manager,
+    )
+    try:
+        for event in gen:
+            match event:
+                case ThinkingStatus(status=thk_status):
+                    if started_content:
+                        console.print()
+                        started_content = False
+                    console.print(f"[dim]{thk_status}[/dim]")
+                case ToolCallEvent(name=name, arguments=args):
+                    if started_content:
+                        console.print()
+                        started_content = False
+                    console.print(
+                        f"[bold cyan]Tool[/bold cyan] {name} "
+                        f"[dim]{json.dumps(args, default=str)}[/dim]"
+                    )
+                case ToolResultEvent(name=name, result=result):
+                    if started_content:
+                        console.print()
+                        started_content = False
+                    _render_tool_result(name, result, console, tmpdir)
+                case ContentDelta(delta=delta):
+                    if not started_content:
+                        console.print("[bold green]Assistant[/bold green]", end=" ")
+                        started_content = True
+                    console.out(delta, end="")
+                case AgentError(message=msg):
+                    if started_content:
+                        console.print()
+                        started_content = False
+                    console.print(f"[red]{msg}[/red]")
+                    status = "errored"
+                    return status
+                case AgentDone():
+                    pass
+    except _AgentInterrupt:
+        status = "interrupted"
+        gen.close()  # closes HTTP connection via generator protocol
+    finally:
+        signal.signal(signal.SIGINT, prev_handler)
+        if started_content:
+            console.print()
+
+    if status == "interrupted":
+        console.print("[dim]interrupted[/dim]")
+
+    return status
+
+
+# ---------------------------------------------------------------------------
+# Resume helpers
+# ---------------------------------------------------------------------------
+
+
+def _pick_session_interactive(console: Console) -> str | None:
+    """Show a rich.Table of recent sessions and prompt for a selection.
+
+    Returns the selected session_id, or None if the user aborts.
+    """
+    summaries = _api.list_sessions()
+    recent = summaries[:10]
+
+    if not recent:
+        console.print("[dim]No saved sessions found.[/dim]")
+        return None
+
+    table = Table(title="Recent sessions", border_style="dim", show_lines=False)
+    table.add_column("#", style="dim", width=3)
+    table.add_column("Started", style="cyan")
+    table.add_column("CWD", style="green")
+    table.add_column("ID", style="dim")
+    for i, entry in enumerate(recent, 1):
+        table.add_row(
+            str(i),
+            entry.started_at[:19],
+            str(entry.cwd),
+            entry.session_id[:8],
+        )
+    console.print(table)
+
+    raw = typer.prompt("Select session number (Enter to cancel)", default="")
+    if not raw.strip():
+        return None
+    try:
+        idx_choice = int(raw.strip()) - 1
+        if not (0 <= idx_choice < len(recent)):
+            raise ValueError
+        return recent[idx_choice].session_id
+    except ValueError:
+        console.print("[red]Invalid selection.[/red]")
+        return None
+
+
+def _resolve_resume_id(
+    continue_session: bool,
+    resume: str | None,
+    pick: bool,
+    cwd: Path,
+    console: Console,
+) -> str | None:
+    """Return the session_id to resume, or None for a fresh session."""
+    if resume is not None:
+        return resume.strip()
+    if pick:
+        return _pick_session_interactive(console)
+    if continue_session:
+        summaries = _api.list_sessions(owner=cwd)
+        if not summaries:
+            console.print(
+                "[dim]No prior session found for this directory. "
+                "Starting fresh.[/dim]"
+            )
+            return None
+        return summaries[0].session_id
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Public command
+# ---------------------------------------------------------------------------
+
+
+def chat_command(
+    prompt: Annotated[
+        str | None,
+        typer.Argument(help="Optional one-shot prompt"),
+    ] = None,
+    model: Annotated[
+        str | None,
+        typer.Option("--model", help="Model name, optionally provider-prefixed"),
+    ] = None,
+    continue_session: Annotated[
+        bool,
+        typer.Option(
+            "-c",
+            "--continue",
+            help="Resume the most recent session for the current directory",
+        ),
+    ] = False,
+    resume: Annotated[
+        str | None,
+        typer.Option(
+            "-r",
+            "--resume",
+            help="Resume a specific session by id (see --pick for the interactive list)",
+        ),
+    ] = None,
+    pick: Annotated[
+        bool,
+        typer.Option(
+            "-p",
+            "--pick",
+            help="Open an interactive picker to choose a session to resume",
+        ),
+    ] = False,
+) -> None:
+    """Chat with a terminal AI agent.
+
+    Sessions are auto-saved to ~/.dft/sessions/ and indexed by working directory.
+
+    Examples:
+      dft chat -c                   Resume last session for this directory
+      dft chat -r <session-id>      Resume a specific session by id
+      dft chat -p                   Pick from a list of recent sessions
+
+    External MCP servers are auto-loaded from the same config files that
+    `dft init mcp` writes (.cursor/mcp.json, .mcp.json, etc.). Dataface's own
+    MCP server entry is skipped — in-process tools are used instead.
+    External tools are namespaced as `server__tool` (double underscore, matching
+    Claude Code's MCP convention and compatible with LLM provider tool-name rules).
+    """
+    from dataface.cli._extras import require_extras
+
+    require_extras("chat")
+
+    import os
+
+    from dataface.ai.external_mcp import ExternalMCPManager, discover_servers
+    from dataface.ai.tool_schemas import ALL_TOOLS
+    from dataface.cli.commands._agent_server import (
+        start_http_server,
+        stop_http_server,
+    )
+
+    console = Console()
+
+    # Start the embedded preview HTTP server before building context, so
+    # render_dashboard URLs point at the agent's own server.
+    server_port_hint = int(os.getenv("DFT_CHAT_PORT", "8765"))
+    http_handle = start_http_server(port_hint=server_port_hint)
+    console.print(f"[dim]Preview server: http://localhost:{http_handle.port}/[/dim]")
+
+    cwd = Path.cwd().resolve()
+    session_id_to_resume = _resolve_resume_id(
+        continue_session, resume, pick, cwd, console
+    )
+
+    chat_session: _api.ChatSession | None = None
+    external_manager: ExternalMCPManager | None = None
+    try:
+        if session_id_to_resume:
+            try:
+                chat_session = _api.resume_session(
+                    session_id_to_resume,
+                    project_dir=Path.cwd(),
+                    server_port=http_handle.port,
+                    model=model,
+                )
+                console.print(
+                    f"[dim]Resumed session {session_id_to_resume[:8]} "
+                    f"({len(chat_session.messages)} messages)[/dim]"
+                )
+            except ValueError as exc:
+                console.print(f"[red]{exc}[/red]")
+                raise typer.Exit(1) from exc
+        else:
+            chat_session = _api.start_session(
+                model=model,
+                project_dir=Path.cwd(),
+                server_port=http_handle.port,
+            )
+
+        external_configs = discover_servers()
+        external_manager = ExternalMCPManager(external_configs)
+        external_manager.start(timeout=10.0)
+        merged_tools = ALL_TOOLS + external_manager.all_tools()
+
+        session = Session(
+            client=chat_session.client,
+            context=chat_session.context,
+            console=console,
+        )
+        session.messages = chat_session.messages
+
+        if prompt:
+            tmpdir = Path(tempfile.mkdtemp(prefix="dft_chat_"))
+            atexit.register(shutil.rmtree, str(tmpdir), True)
+            result_status = _run_agent_session(
+                chat_session,
+                prompt,
+                console=console,
+                tmpdir=tmpdir,
+                tools=merged_tools,
+                external_manager=external_manager,
+            )
+            if result_status == "interrupted":
+                raise typer.Exit(130)
+            if result_status == "errored":
+                raise typer.Exit(1)
+            return
+
+        _run_interactive(
+            session=session,
+            chat_session=chat_session,
+            console=console,
+            merged_tools=merged_tools,
+            external_manager=external_manager,
+        )
+    finally:
+        if external_manager is not None:
+            external_manager.close()
+        stop_http_server(http_handle)
+        if chat_session is not None:
+            chat_session.writer.close()
+
+
+def _run_interactive(
+    session: Session,
+    chat_session: _api.ChatSession,
+    console: Console,
+    merged_tools: list[dict[str, Any]],
+    external_manager: ExternalMCPManager,
+) -> None:
+    """Run the interactive REPL loop."""
+    tmpdir = Path(tempfile.mkdtemp(prefix="dft_chat_"))
+    atexit.register(shutil.rmtree, str(tmpdir), True)
+
+    hist_path = dft_home() / "chat_history"
+    input_layer = select_input_layer(hist_path, slash_commands=set(SLASH_COMMANDS))
+    stream_ctx: Callable[..., AbstractContextManager[None]]
+    if isinstance(input_layer, PromptToolkitInput):
+        from prompt_toolkit.patch_stdout import (
+            patch_stdout as _patch_stdout,
+        )  # noqa: PLC0415
+
+        stream_ctx = _patch_stdout
+    else:
+        stream_ctx = nullcontext
+
+    last_sigint_ts: list[float] = []
+
+    def _at_prompt_sigint(_signum: int, _frame: Any) -> None:
+        if _check_double_tap_exit(last_sigint_ts):
+            console.print("\n[dim]Session ended.[/dim]")
+            raise SystemExit(0)
+        console.print("\n[dim](Press Ctrl+C again within 1s to exit, or Ctrl+D)[/dim]")
+
+    signal.signal(signal.SIGINT, _at_prompt_sigint)
+
+    while True:
+        try:
+            user_input = input_layer.read("You: ")
+        except EOFError:
+            console.print("\n[dim]Session ended.[/dim]")
+            return
+
+        if not user_input.strip():
+            continue
+        if _dispatch_slash(session, user_input):
+            continue
+
+        user_input = expand_file_refs(user_input, project_dir=Path.cwd()).text
+
+        # Slash commands like /clear and /model rebuild session.client; mirror
+        # that onto chat_session so send_message picks up the new client.
+        chat_session.client = session.client
+
+        with stream_ctx():
+            _run_agent_session(
+                chat_session,
+                user_input,
+                console=session.console,
+                tmpdir=tmpdir,
+                tools=merged_tools,
+                external_manager=external_manager,
+            )